GPU or MPS error when running run_clm.py
Author: oscar-defelice
Pull request: no

### System Info

## System Info

```bash
- `transformers` version: 4.37.0.dev0
- Platform: macOS-14.2-arm64-arm-64bit
- Python version: 3.11.7
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
```

---

Even if I am pasting this output, if I run on Ubuntu with 2 GPU I got the same issue.

### Who can help?

@ArthurZucker @muellerz 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I run 

```bash
python run_clm.py --model_name_or_path nferruz/ProtGPT2 --train_file data/fine_tune_data.txt --tokenizer_name nferruz/ProtGPT2 --do_train --output_dir models/ProtGPT/output --learning_rate 1e-06
```

And no matter what I try with batch_size and learning rate I always get

```bash
RuntimeError: MPS backend out of memory (MPS allocated: 78.40 GB, other allocations: 2.98 GB, max allowed: 81.60 GB). Tried to allocate 320.00 MB on private pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure).
```

### Expected behavior

It should work and finetune the model =) 

Comments:
Hi @oscar-defelice, thanks for raising an issue! 

When you say that you've tried varying the batch size, what values have you tested with. Are you able to run with a batch size of 1 e.g.:

```
python run_clm.py \
--model_name_or_path nferruz/ProtGPT2 \
--train_file data/fine_tune_data.txt \
--tokenizer_name nferruz/ProtGPT2 \
--do_train \
--output_dir models/ProtGPT/output \
--learning_rate 1e-06
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
```

How about running with a small model e.g.: 
```
python run_clm.py \
--model_name_or_path gpt2 \
--train_file data/fine_tune_data.txt \
--tokenizer_name nferruz/ProtGPT2 \
--do_train \
--output_dir models/ProtGPT/output \
--learning_rate 1e-06
--per_device_train_batch_size 1 \
--per_device_eval_batch_size 1 \
```
?

What is the size of the GPUs being used on the ubuntu machine? 

Hello @amyeroberts  thank you for your reply. Even with both train and eval batch size equal 1 I got the same error.

On Ubuntu the transformer-cli env command gives

```bash
- `transformers` version: 4.37.0.dev0
- Platform: Linux-6.2.0-39-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.0
- Accelerate version: 0.25.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
```

and for the GPUs 

```bash
display
    description: VGA compatible controller
    produit: NVIDIA Corporation
    fabricant: NVIDIA Corporation
    identifiant matériel: 0
    information bus: pci@0000:41:00.0
    version: a1
    bits: 64 bits
    horloge: 33MHz
    fonctionnalités: vga_controller bus_master cap_list rom
    configuration : driver=nvidia latency=0
    ressources : mémoireE/S:1100-10ff mémoireE/S:1180-117f irq:263 mémoire:f0000000-f0ffffff mémoire:11000000000-117ffffffff mémoire:11800000000-11801ffffff portE/S:4000(taille=128) mémoire:f1000000-f107ffff
 *-display
    description: VGA compatible controller
    produit: NVIDIA Corporation
    fabricant: NVIDIA Corporation
    identifiant matériel: 0
    information bus: pci@0000:61:00.0
    version: a1
    bits: 64 bits
    horloge: 33MHz
    fonctionnalités: vga_controller bus_master cap_list rom
    configuration : driver=nvidia latency=0
    ressources : mémoireE/S:1000-fff mémoireE/S:1080-107f irq:262 mémoire:f4000000-f4ffffff mémoire:10000000000-107ffffffff mémoire:10800000000-10801ffffff portE/S:f000(taille=128) mémoire:f5000000-f507ffff
 *-graphics
    produit: EFI VGA
    identifiant matériel: 2
    nom logique: /dev/fb0
    fonctionnalités: fb
    configuration : depth=32 resolution=1920,1080

@oscar-defelice Did you try running with a smaller model? Was that successful or did you hit the same issue? Do you see your GPUs being utilized when running the script? 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@amyeroberts I am facing the same issue. I have an A100 GPU machine with 40GB GPU memory.

With the ProtGPT2 model, I get out of memory error even with the batch size set to 1.
However, the smaller model (GPT2) works properly on this machine. It takes a maximum of 16GB GPU memory.

Also, I have tried running this on g5.12xlarge ec2 machine with 4 GPUs with 24GB memory each. On that I get this error
```
Traceback (most recent call last):
  File "/protgpt2/run_clm.py", line 657, in <module>
    main()
  File "/protgpt2/run_clm.py", line 605, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/transformers/trainer.py", line 1938, in train
    return inner_training_loop(
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/transformers/trainer.py", line 2279, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/transformers/trainer.py", line 3318, in training_step
    loss = self.compute_loss(model, inputs)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/transformers/trainer.py", line 3363, in compute_loss
    outputs = model(**inputs)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py", line 184, in forward
    replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py", line 189, in replicate
    return replicate(module, device_ids, not torch.is_grad_enabled())
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/parallel/replicate.py", line 110, in replicate
    param_copies = _broadcast_coalesced_reshape(params, devices, detach)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/parallel/replicate.py", line 83, in _broadcast_coalesced_reshape
    tensor_copies = Broadcast.apply(devices, *tensors)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/autograd/function.py", line 598, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/parallel/_functions.py", line 23, in forward
    outputs = comm.broadcast_coalesced(inputs, ctx.target_gpus)
  File "/root/miniconda3/envs/protgpt2/lib/python3.9/site-packages/torch/nn/parallel/comm.py", line 57, in broadcast_coalesced
    return torch._C._broadcast_coalesced(tensors, devices, buffer_size)
RuntimeError: NCCL Error 2: unhandled system error (run with NCCL_DEBUG=INFO for details)
``` 

I also have the logs with NCCL_DEBUG=INFO. Let me know if you need those too.

Hi @baksh97, thanks for sharing these details! 

For the NCCL error - this is probably a lower-level installation / configuration issue between torch, cuda and nccl in the environment. I'd suggest trying re-install torch to the latest version. 

Regarding the OOM, comparing at the configs of the models - [GPT2](https://huggingface.co/openai-community/gpt2/blob/main/config.json) and [nferruz/ProtGPT2](https://huggingface.co/nferruz/ProtGPT2/blob/main/config.json), you'll see that nferruz/ProtGPT2 is a far larger model e.g. `n_embed` is 1280 compared to `768`; `n_head` is 20 compared to 12; `n_layer` is 36 to 12 i.e. it is many times larger than GPT2. Testing on my machine of 4x A10G, I also hit OOM, but was able to get things running by adding `device_map="auto"` to the AutoModelForCausalLM call in the script and running with the `--torch_dtype bfloa16` flag