Lllama2 layer outputs NaN when using dual GPU
Author: ekim322
Pull request: no

### System Info

- `transformers` version: 4.33.1
- Platform: Linux-5.15.0-83-generic-x86_64-with-glibc2.31
- Python version: 3.9.17
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.3
- Accelerate version: 0.22.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker @SunMarc @pacman100 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

I'm facing errors when attempting to QLoRA fine tuning LLAMA2 using a dual GPU set up.

The code runs properly when if I use only one of the two GPUs I have (RTX4090 Founders Edition and RTX4090 Liquid Suprim). The exact code also runs fine on a GCP vm instance with a two Nvidia T4s. 
It's just not working when I run it on my local dual GPU set up.

I use trl.SFTTrainer to start training, but the errors seem to be occuring from the hugging face library so I am opening the issue here.

```
from trl import SFTTrainer

trainer = SFTTrainer(
    model = self.model,
    tokenizer=self.tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    args=train_args,
    max_seq_length=self.cfg.hparams.max_token_len,
    formatting_func=data_format_fn,
    data_collator=data_collator,
    callbacks=callbacks
)

trainer.train()
```

Here is the error I am facing 
```
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [0,0,0] Assertion `t >= 0 && t < n_classes` failed.
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [1,0,0] Assertion `t >= 0 && t < n_classes` failed.
...
...
...
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [30,0,0] Assertion `t >= 0 && t < n_classes` failed.
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [31,0,0] Assertion `t >= 0 && t < n_classes` failed.
Traceback (most recent call last):
  File "/home/ekim/test/trainer.py", line 178, in train
    trainer.train()
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 1553, in train
    return inner_training_loop(
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 1835, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 2679, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 2704, in compute_loss
    outputs = model(**inputs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/accelerate/utils/operations.py", line 632, in forward
    return model_forward(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/accelerate/utils/operations.py", line 620, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/amp/autocast_mode.py", line 14, in decorate_autocast
    return func(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/peft/peft_model.py", line 918, in forward
    return self.base_model(
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/peft/tuners/tuners_utils.py", line 94, in forward
    return self.model.forward(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py", line 861, in forward
    loss = loss_fct(shift_logits, shift_labels)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/functional.py", line 3029, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```


The error seems to be occuring because there are NaN values passed into torch.F.cross_entropy. 

### Expected behavior

When checking the forward pass in LlamaModel ([link](https://github.com/huggingface/transformers/blob/ce2e7ef3d96afaf592faf3337b7dd997c7ad4928/src/transformers/models/llama/modeling_llama.py#L615C13-L615C13)), there seems to be 40 layers. The first 17 layers are allocated to GPU:0 and the rest of the layers (18-40) are on GPU:1. 

Observations:
- The outputs from first 17 layers (on GPU:0) are as expected
- Layers 18 to 20 (on GPU:1) also produce expected outputs

However, starting from layer 21 (on GPU:1), all the layer_outputs ([link](https://github.com/huggingface/transformers/blob/ce2e7ef3d96afaf592faf3337b7dd997c7ad4928/src/transformers/models/llama/modeling_llama.py#L701)) yield NaNs. This is not expected behavior and is the source of the cuda error. 

The same exact code runs without problems on a single GPU or another GCP instance with dual GPU. I'm thinking it might be some communication issue between the two GPUs I have but I am not sure how to fix this. 

Comments:
Hey, is the `self.model` in `fp16` or `bfloat16`?

I was loading in `fp16`. Here are my configs

```
bnb_cfg = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16 # torch.bfloat16
)

train_args = TrainingArguments(
    output_dir=os.path.join("runs", cfg.wandb.project, cfg.wandb.group, cfg.exp),
    num_train_epochs=cfg.hparams.epochs,
    per_device_train_batch_size=cfg.hparams.batch_size,
    gradient_accumulation_steps=cfg.hparams.grad_accum,
    gradient_checkpointing=True,
    # logging_steps=self.cfg.hparams.grad_accum, 
    logging_steps=1, 
    save_strategy="steps",
    save_steps=self.cfg.save_steps,
    report_to="wandb",
    optim=cfg.hparams.optimizer,
    learning_rate=cfg.hparams.learning_rate,
    lr_scheduler_type=cfg.hparams.lr_scheduler_type,
    max_grad_norm=cfg.hparams.max_grad_norm,
    warmup_ratio=cfg.hparams.warmup_ratio,
    evaluation_strategy='steps',
    eval_steps=self.cfg.eval_steps, 
    eval_accumulation_steps=1, 
    fp16=True, #  bf16=True,
    ddp_find_unused_parameters=False,
)
```

When I switch to `bf16`, none of the 40 layers output NaN values. But I still get errors in the cross entropy call. 
```
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [0,0,0] Assertion `t >= 0 && t < n_classes` failed.
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [1,0,0] Assertion `t >= 0 && t < n_classes` failed.
...
...
...
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [30,0,0] Assertion `t >= 0 && t < n_classes` failed.
../aten/src/ATen/native/cuda/Loss.cu:240: nll_loss_forward_reduce_cuda_kernel_2d: block: [0,0,0], thread: [31,0,0] Assertion `t >= 0 && t < n_classes` failed.
Traceback (most recent call last):
  File "/home/ekim/test/trainer.py", line 178, in train
    trainer.train()
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 1553, in train
    return inner_training_loop(
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 1835, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 2679, in training_step
    loss = self.compute_loss(model, inputs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/trainer.py", line 2704, in compute_loss
    outputs = model(**inputs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/accelerate/utils/operations.py", line 632, in forward
    return model_forward(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/accelerate/utils/operations.py", line 620, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/amp/autocast_mode.py", line 14, in decorate_autocast
    return func(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/peft/peft_model.py", line 918, in forward
    return self.base_model(
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/peft/tuners/tuners_utils.py", line 94, in forward
    return self.model.forward(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/transformers/models/llama/modeling_llama.py", line 861, in forward
    loss = loss_fct(shift_logits, shift_labels)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/home/ekim/anaconda3/envs/torch/lib/python3.9/site-packages/torch/nn/functional.py", line 3029, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

[This](https://github.com/pytorch/pytorch/blob/e7bd9c53158fd6a08b0f93462ee883646c7451f9/torch/nn/functional.py#L3053) is the specific line that causes `CUDA error: device-side assert triggered`. 

`torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)`

```
input
tensor([[ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        ...,
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906]],
       device='cuda:1', grad_fn=<ViewBackward0>)

target
tensor([3577293384251152737, 3577293384251152737, 3577293384251152737,
        3577293384251152737, 3577293384251152737, 3577293384251152737,
        3577293384251152737, 3577293384251152737, 3577293384251152737,
        3577293384251152737, 3577293384251152737, 3577293384251152737,
        3577293384251152737 ...... 3577293384251152736], device='cuda:1')

weight
None

reduction
'mean'

ignore_index
-100

label_smoothing
0.0
```

When I run my code with a single GPU `target` is different than when I use two GPUs (I was expecting it to be the same). 
```
target # with single GPU
tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100, ..., ..., ..., ... -100,  -100,  -100,  -100,
           13,   450, 17097,  8818, 29769,  1353,   310,   278,  8010,   402,
         6935,  2552,  5449,   653,  9000, 10606,   338, 29871, 29953, 29946,
        29899, 29953, 29906, 29906, 29896, 29945, 29946, 29896, 29889,  -100],
       device='cuda:0')
```

The unexpected `target` values for dual GPU seem to originate from [here](https://github.com/huggingface/transformers/blob/ce2e7ef3d96afaf592faf3337b7dd997c7ad4928/src/transformers/models/llama/modeling_llama.py#L851). `shift_labels = shift_labels.to(shift_logits.device)` 


```
if labels is not None:
    # Shift so that tokens < n predict n
    shift_logits = logits[..., :-1, :].contiguous()   # slogitA
    shift_labels = labels[..., 1:].contiguous()      # slabelA
    # Flatten the tokens
    loss_fct = CrossEntropyLoss()
    shift_logits = shift_logits.view(-1, self.config.vocab_size) # slogitB
    shift_labels = shift_labels.view(-1)                                    # slabelB
    # Enable model parallelism
    shift_labels = shift_labels.to(shift_logits.device)               # slabelC
    loss = loss_fct(shift_logits, shift_labels)
```

```
logits
tensor([[[ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         ...,
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906]]],
       device='cuda:1', grad_fn=<ToCopyBackward0>)

labels
tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100, 
          -100,  -100,  -100,  -100,  -100,  -100,  ...., ...., ....,   -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,    13,   450, 17097,  8818, 29769,  1353,   310,   278,  8010,
           402,  6935,  2552,  5449,   653,  9000, 10606,   338, 29871, 29953,
         29946, 29899, 29953, 29906, 29906, 29896, 29945, 29946, 29896, 29889,
          -100]], device='cuda:0')

shift_logits # slogitA
tensor([[[ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         ...,
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
         [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906]]],
       device='cuda:1', grad_fn=<SliceBackward0>)

shift_labels # slabelA
tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100, 
          -100,  -100,  -100,  -100,  -100,  -100,  ..., ..., ...,  -100,  -100,
            13,   450, 17097,  8818, 29769,  1353,   310,   278,  8010,   402,
          6935,  2552,  5449,   653,  9000, 10606,   338, 29871, 29953, 29946,
         29899, 29953, 29906, 29906, 29896, 29945, 29946, 29896, 29889,  -100]],
       device='cuda:0')

shift_logits # slogitB
tensor([[ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        ...,
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906],
        [ 2.1562,  9.9375, -0.0903,  ...,  1.7656, -0.4590,  1.3906]],
       device='cuda:1', grad_fn=<ViewBackward0>)


shift_labels # slabelB
tensor([ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
         -100,  -100,  -100,  -100,  -100,  -100,  ..., ..., ..., ... ,  -100,  -100,
           13,   450, 17097,  8818, 29769,  1353,   310,   278,  8010,   402,
         6935,  2552,  5449,   653,  9000, 10606,   338, 29871, 29953, 29946,
        29899, 29953, 29906, 29906, 29896, 29945, 29946, 29896, 29889,  -100],
       device='cuda:0')


shift_labels # slabelC
tensor([3577293384251152737, 3577293384251152737, 3577293384251152737,
        3577293384251152737, ..., ..., ..., 3974362539628152236], device='cuda:1')
```



Also 

Sometimes `CUDA error: device-side assert triggered` happens [here](https://github.com/huggingface/transformers/blob/ce2e7ef3d96afaf592faf3337b7dd997c7ad4928/src/transformers/models/llama/modeling_llama.py#L184) `cos = cos[position_ids].unsqueeze(1)`. I am using vscode debugger and the error is not consistent and seem to occur randomly. 




You should train the model in bfloat16. It's expected to not work if you train in float16 you need to monkey patch the code !

bfloat16 is not working either.
When I transfer data from GPU to GPU the values are getting modified randomly. 

GPU to CPU and CPU to GPU is working fine.
For example if I send tensor from CPU -> GPU0 -> CPU -> GPU1, all the values are preserved.

```
labels = torch.tensor([[ 0.3316, -2.1440, -0.6160, -0.0321,  0.0256]])
labels = labels.to('cuda:0') # tensor([[ 0.3316, -2.1440, -0.6160, -0.0321,  0.0256]], device='cuda:0')
labels = labels.to('cpu') # tensor([[ 0.3316, -2.1440, -0.6160, -0.0321,  0.0256]])
labels = labels.to('cuda:1') # tensor([[ 0.3316, -2.1440, -0.6160, -0.0321,  0.0256]], device='cuda:1')
```

However when I transfer from GPU0 -> GPU1, the tensor values change. 
```
labels = torch.tensor([[ 0.3316, -2.1440, -0.6160, -0.0321,  0.0256]])
labels = labels.to('cuda:0') # tensor([[ 0.3316, -2.1440, -0.6160, -0.0321,  0.0256]], device='cuda:0')
labels = labels.to('cuda:1') # tensor([[0., 0., 0., 0., 0.]], device='cuda:1')
```

I've tried [this](https://discuss.pytorch.org/t/tensor-totally-changes-when-allocating-moving-from-gpu-to-gpu/73930/16) `export NCCL_P2P_DISABLE=1` but this is not working either. 




My temporary fix is to move the data to CPU first than to GPU rather than direct GPU to GPU
```
shift_labels = shift_labels.to('cpu')
shift_labels = shift_labels.to(shift_logits.device)
```

However, there are numerous other variables that have this issue as well and I don't think editing the source code one by one is the best approach since the library gets updated quite often. 

Any thoughts on how to fix this?

Okay, I'll let @pacman100 handle this! But we need you to share a full reproducer, as minimal as possible ( which model you are using from which repo etc)

Hello, this seems to Hardware issue, referring to https://discuss.pytorch.org/t/problem-transfering-tensor-between-gpus/60263 and https://discuss.pytorch.org/t/tensor-totally-changes-when-allocating-moving-from-gpu-to-gpu/73930/16

This was a communication issue between the two GPUs (2x4090) I had. 


@pacman100 The solutions in the links you copied did not work for me. 

[Here](https://forums.developer.nvidia.com/t/standard-nvidia-cuda-tests-fail-with-dual-rtx-4090-linux-box/233202/22?page=2) is how I found out my solution.

I faced errors when using Nvidia Driver 520 and 530. 
Using Nvidia Driver 525.105.17 with Cuda 11.8 solved the issue

> This was a communication issue between the two GPUs (2x4090) I had.
> 
> @pacman100 The solutions in the links you copied did not work for me.
> 
> [Here](https://forums.developer.nvidia.com/t/standard-nvidia-cuda-tests-fail-with-dual-rtx-4090-linux-box/233202/22?page=2) is how I found out my solution.
> 
> I faced errors when using Nvidia Driver 520 and 530. Using Nvidia Driver 525.105.17 with Cuda 11.8 solved the issue

Was the version you were using before 11.7? 

Just to note here that I had the same problem with nvidia drivers 545 and cuda 12.1. Switch to 525 and 11.8 and it's all working again with a dual 4090 setup.

I finally solved this by disabling ACS in bios, ref https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#pci-access-control-services-acs. Changing nvidia driver and cuda version doesn't help.

This test is very helpful. https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#gpu-to-gpu-communication