Hitting OOM in model loading in multi GPU setting

[rupeshgx]

My code:

def main():
    # 1. Load a model to finetune
    model = SentenceTransformer(
        model_name_or_path=Alibaba-NLP/gte-Qwen2-7B-instruct,
        tokenizer_kwargs={
            "model_max_length": 512,
            "truncation": True
        }
    )
    # set the max input seq length to 512
    model.max_seq_length = 512


if __name__ == "__main__":
    main()

When I run this code on one H100 GPU (which has 80 GB memory) I am able to load the model successfully. The model consumes 28GB of memory.

However, when I run the same code on 8 H100 GPUs (torchrun --nproc_per_node=8 finetuning.py) I get OOM with the following error message:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 260.00 MiB. GPU 0 has a total capacity of 79.10 GiB of which 57.06 MiB is free. Process 1552209 has 9.94 GiB memory in use. Process 1552212 has 9.94 GiB memory in use. Process 1552208 has 9.94 GiB memory in use. Process 1552211 has 9.68 GiB memory in use. Process 1552213 has 9.94 GiB memory in use. Process 1552210 has 9.94 GiB memory in use. Process 1552214 has 9.68 GiB memory in use. Process 1552215 has 9.94 GiB memory in use ...

What am I doing wrong?

[13918763630]

I also have this problem.