llama2 training has nan
Author: LZY-the-boys
Pull request: no

### System Info

- `transformers` version: 4.31.0
- Platform: Linux-5.19.0-42-generic-x86_64-with-glibc2.35
- Python version: 3.9.16
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: NO
        - mixed_precision: fp16
        - use_cpu: False
        - num_processes: 1
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: all
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.0.0+cu117 (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?

_No response_

### Information

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

### Tasks

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

### Reproduction

I am training llama2-7b using huggingface trainer. I find that it will occur nan in **forward** when training with batchsize > 1 for llama2, however when batchsize=1 there is no error.
I dive into it and find that the **nan** occurs in layer.31.input_layer_norm, which is caused by **inf** in layers.30.mlp forward after the post_layer_norm, and this **inf** may comes from huge value in hidden_size. However, this doesn't explain why llama1 and llama2 with batchsize=1 can work, which also has huge outliners in hidden_size. 
The code I use is like this: 
The dataset format is: `write a xxx. ###sentences: xxxx`
The checkpoint of `meta-llama/Llama-2-7b-hf` and original meta checkpoint converted by transformers code are both tried.
```
trainer = transformers.Trainer(
    model=model,
    train_dataset=train_data,
    eval_dataset=val_data,
    args=transformers.TrainingArguments(
        per_device_train_batch_size=args.micro_batch,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        warmup_ratio=args.warmup_ratio,
        num_train_epochs=args.num_epoch,
        learning_rate=3e-4,
        fp16=True,
        logging_steps=args.log_steps,
        logging_first_step=True, # convenient
        evaluation_strategy="no",
        save_strategy=args.save_strategy,
        eval_steps=None,
        save_steps=args.save_steps,
        output_dir=args.output_path,
        load_best_model_at_end= False,
        ddp_find_unused_parameters=False if ddp else None,
        report_to="wandb" if args.wandb else [],
        ignore_data_skip=args.ignore_data_skip,
    ),
    data_collator=PROMPT.data_collator()
)
model.config.use_cache = False
if list(pathlib.Path(args.output_path).glob("checkpoint-*")):
    trainer.train(resume_from_checkpoint=True)
else:
    trainer.train()
trainer.save_state()
model.save_pretrained(args.output_path)
```

### Expected behavior

The training should has no nan in forward, thus loss will be normal in backward

Comments:
What are you setting as your `pad_token`? I was getting `nan` while running batched inference (see post [here](https://discuss.huggingface.co/t/llama2-pad-token-for-batched-inference/48020)) and am wondering if this might be related to your issue?

I set tokenizer.pad_token_id = 0, which is the same as llama1 (:

I have the same question!

Me too！

Hey! Thanks all for reporting this. I would suggest using the following: 
```python 
        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights + attention_mask
            dtype_min = torch.tensor(
                torch.finfo(attn_weights.dtype).min, device=attn_weights.device, dtype=attn_weights.dtype
            )
            attn_weights = torch.max(attn_weights, dtype_min)
``` 
This was removed from:

https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L343-L348 

because it is not used in the original model. This should help. Otherwise adding 
```python 
        # clamp inf values to enable fp16 training
        if hidden_states.dtype == torch.float16:
            max_dtype = torch.finfo(hidden_states.dtype).max
            clamp_value = torch.where(torch.isinf(hidden_states).any(), max_dtype - 1000, max_dtype)
            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
```
will also help mitigate this 


Another fix is to train with bfloat16

> Another fix is to train with bfloat16

+1, this actually worked for me with accelerate FSDP training

Can you share what config you used for training @arazd  I still see nan with bf16


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.

> I set tokenizer.pad_token_id = 0, which is the same as llama1 (:

@LZY-the-boys, by the way, could you please tell how do you exclude padding tokens from loss computation then? As model expects padding tokens in the `labels` to be `-100`, what `CrossEntropyLoss` expects:

https://github.com/huggingface/transformers/blob/000e52aec8850d3fe2f360adc6fd256e5b47fe4c/src/transformers/models/llama/modeling_llama.py#L792-L796

Something like this 
`labels = torch.where(inputs.input_ids == tokenizer.pad_token_id, labels, -100)` 
would do the trick. 

> Hey! Thanks all for reporting this. I would suggest using the following:
> 
> ```python
>         if attention_mask is not None:
>             if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
>                 raise ValueError(
>                     f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
>                 )
>             attn_weights = attn_weights + attention_mask
>             dtype_min = torch.tensor(
>                 torch.finfo(attn_weights.dtype).min, device=attn_weights.device, dtype=attn_weights.dtype
>             )
>             attn_weights = torch.max(attn_weights, dtype_min)
> ```
> 
> This was removed from:
> 
> https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L343-L348
> 
> because it is not used in the original model. This should help. Otherwise adding
> 
> ```python
>         # clamp inf values to enable fp16 training
>         if hidden_states.dtype == torch.float16:
>             max_dtype = torch.finfo(hidden_states.dtype).max
>             clamp_value = torch.where(torch.isinf(hidden_states).any(), max_dtype - 1000, max_dtype)
>             hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
> ```
> 
> will also help mitigate this

Hi @ArthurZucker,

I was wondering if it would be possible to integrate the provided code snippet(the 2nd approach) into the transformer's repository.

I have concerns about using bf16, as it may result in some loss of precision during inference. Additionally, older GPUs like V100 do not support bf16. Currently, in order to bypass the issue of `nan` values in padded input data, I have had to manually copy a significant amount of code from the modeling_llama.py file and make the aforementioned modifications.

Furthermore, to ensure the safety of these modifications, would it be possible to apply the clamp only to the padded token section? According to theory, `nan` values can occur at any token's position, not just limited to padded tokens. However, based on the investigation conducted in pull request #25284, it appears that `nan` values are most likely to occur at padded tokens. Therefore, I believe it would be sufficient to modify the logits solely for padded tokens. If, by any chance, the same error occurs at unpadded tokens, it might be preferable to consider changing the precision to fp32/bf16, as altering the logits of unpadded tokens could degrade the model's performance.

Thank you for your assistance.

Hey! I'll deep dive a bit on this issue this week, making sure we can fix this without any issues but remember that the goal of `transformers` is to make it easy to develop on top of it! 
To me the changes required in a codebase look pretty minimal, clamping just before softmax. Remember that clamping has a cost in terms of speed.  

Hi Arthur, thank you for your suggestion! I appreciate your willingness to address the issue. Yeah, I agree clamping should do the work. And more precisely, we should clamp before layernorm. Given the chain of error:

30.layernorm -> 30.attention -> 30.mlp[-inf appears] -> 31.layernorm[nan appears] -> 31.attention[more nans appear]

to avoid nan, we'd better clamp somewhere before layernorm. Hence I prefer the second approach, which rescue the logits after layer 30's mlp.

However, I am uncertain if this solution addresses the root cause of the problem. As highlighted in pr #25284, there seems to be an issue with the attention mask for left-padded tokens. Some of the padded tokens have masks like [-65534, -Inf, -Inf, -65534, -65534], resulting in attention weights of [0.33, 0, 0, 0.33, 0.33]. Consequently, the padded token at position 1 is attending to tokens after it (position 4 and 5). This behavior is unexpected in a causal language model and is not something the model should have learned during pretraining. I wonder if correcting the attention behavior would eliminate the generation of -inf values or if there might be another underlying issue.

For now, I will adhere to the solution of clamping for the sake of simplicity and ease of use.

Thank you again for your attention and assistance.

Sure! Correcting the padded causal mask is also important! I remember these issues appeared when @younesbelkada was porting the LlamaFlashAttention, I'll definitely have a look!  

I am observing the exact same behavior as @VeryLazyBoy on the 30th and 31st decoder layer of llama-2. I'm using llama-2-chat-7b. Is there any update on this issue or should I still use the workaround?

Update: FYI, did some investigation and found that the nan values come from

-  The 30.MLP gives inf values in ``hidden_states``, same as what @VeryLazyBoy said
-  In 31.input_layernorm, there is ``variance = hidden_states.pow(2).mean(-1, keepdim=True)``. Variance now has a few inf values from ``hidden_states``.
- Then, the next line in 31.input_layernorm is ``hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)``. ``torch.rsqrt(variance + self.variance_epsilon)`` turns those inf values to 0 values. Then, you multiply the inf values in ``hidden_states`` with those 0 values and produce the few nan values.
- Lastly, the nan values spread out in the attention calculations and result in a lot of nan values in the last hidden states.

My workaround is to add the clamp before self.input_layernorm is called in LlamaDecoderLayer's forward. Let me know if there is a better solution.

Mmm the issue mentioned by @VeryLazyBoy regarding the attention mask was fixed by #27114 and should not longer affect the forward. 
If the training is done in fp16 instead of bf16 I think it's pretty much expected that there will be some overflow. 
Clamping can significantly slow down the training, but might be the best if you have to use float16 

[This comment](https://github.com/huggingface/transformers/pull/27114#issuecomment-1848235762) might be helpful

> I am observing the exact same behavior as @VeryLazyBoy on the 30th and 31st decoder layer of llama-2. I'm using llama-2-chat-7b. Is there any update on this issue or should I still use the workaround?
> 
> Update: FYI, did some investigation and found that the nan values come from
> 
> * The 30.MLP gives inf values in `hidden_states`, same as what @VeryLazyBoy said
> * In 31.input_layernorm, there is `variance = hidden_states.pow(2).mean(-1, keepdim=True)`. Variance now has a few inf values from `hidden_states`.
> * Then, the next line in 31.input_layernorm is `hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)`. `torch.rsqrt(variance + self.variance_epsilon)` turns those inf values to 0 values. Then, you multiply the inf values in `hidden_states` with those 0 values and produce the few nan values.
> * Lastly, the nan values spread out in the attention calculations and result in a lot of nan values in the last hidden states.
> 
> My workaround is to add the clamp before self.input_layernorm is called in LlamaDecoderLayer's forward. Let me know if there is a better solution.

@BaleChen Hi, I'd like to ask how did you add the clamp, I added the clamp and still get nan