Padding token results in Nan for logits and loss for Llama2
Author: NamburiSrinath
Pull request: no
Labels: name: Generation

### System Info

- `transformers` version: 4.34.0
- Platform: Linux-5.15.0-100-generic-x86_64-with-glibc2.31
- Python version: 3.12.1
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.0 (True)

### Who can help?
I think folks working on generation and/or documentation can be helpful
@LysandreJik @patrickvonplaten @sgugger 

### Reproduction

I am trying to implement the generate logic by myself for Llama-2-7B and the below code works fine 😃 
```
model = AutoModelForCausalLM.from_pretrained(
    'Llama-2-7B-hf'
)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Update few config parameters to satisfy padding constraints!
tokenizer.padding_side = 'left'
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
model.config.pad_token_id = tokenizer.pad_token_id
model.config.pad_token = tokenizer.pad_token

## Modified generate logic using model.forward()
max_length = 30
def updated_generate(prompts, model):
    # tokenize batch using truncation and padding
    tokenized_input = tokenizer(prompts, return_tensors="pt")

    # Get input ids and attention mask from tokenizer 
    input_ids = tokenized_input.input_ids
    attention_mask = tokenized_input.attention_mask

    # Process position ids from attention mask
    position_ids = attention_mask.long().cumsum(-1) - 1
    position_ids.masked_fill_(attention_mask == 0, 1)

    # Process labels from input ids i.e replace all the padded tokens with -100
    input_ids_masked = torch.zeros(tokenized_input.input_ids.shape, 
                                dtype=torch.int64)
    torch.where(tokenized_input.input_ids == tokenizer.pad_token_id,
                torch.tensor(-100, dtype=torch.int64),
                tokenized_input.input_ids,
                out=input_ids_masked)

    for i in range(max_length):
        outputs = model.forward(input_ids=input_ids, 
                        attention_mask = attention_mask, 
                        labels = input_ids_masked,
                        position_ids=position_ids
                        )

        # Extract logits and do greedy sampling
        logits = outputs.logits[:,-1,:]
        predicted_id = logits.argmax(-1)
        
        # Update input ids, mask, labels and position ids
        # input ids -> add the generated seq
        # attention mask -> add 1
        # labels (input_ids_masked) -> add the generated seq (same as input ids)
        # position ids -> get the max position and add 1
        
        input_ids = torch.cat([input_ids, predicted_id.unsqueeze(1)], dim=1)
        attention_mask = torch.cat([attention_mask, torch.ones(1).unsqueeze(1)], dim=1)
        input_ids_masked = torch.cat([input_ids_masked, predicted_id.unsqueeze(1)], dim=1)
        position_ids = torch.cat([position_ids, (position_ids.max().unsqueeze(0) + 1).unsqueeze(1)], dim=1)
    
    print(tokenizer.batch_decode(input_ids))
```

But the moment I add padding i.e
```
tokenized_input = tokenizer(prompts, return_tensors="pt", 
                        padding='max_length', max_length=128
                        )
```
the output loss and logits are nan 😢 . 
And because of that, only 0 gets added to the input ids! I've printed and the variables are properly getting updated.


### Expected behavior

This might not be a bug from Huggingface and I might be missing some update here most likely, but I've referred to the [LlamaForCausalLM class](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1129 ) and don't see anything additional that I might be missing apparently.

Why? 
1. Padding is needed for batch generation. When I don't do padding, my loop is working fine 😄 
2. Instead of using model.generate(), I would like to know in-depth on model.forward() -> Access to logits, loss and other use-cases which needs a fine-grained analysis.

I'll also redirect it to Huggingface forums. 

Comments:
Hi @NamburiSrinath, thanks for opening this issue! 

Yes, this is a question best placed in our forums

cc @gante for reference in case there's something you immediately spot 

Hi @NamburiSrinath 👋 

At first glance, everything looks correct. You're not committing common mistakes like not using left-padding or not updating the attention mask/position ids.

What I would do, ordered:
1. try removing the `labels` input, see if there is any change;
2. compare your unpadded outputs vs generate, confirm that they match at a logit level (using `output_scores=True`)
3. place a `breakpoint()` inside `generate` before `forward`. Confirm that the inputs match.

If you find the mismatch cause, please let us know -- it would be useful information for me :)

Hi @gante and @amyeroberts, 

Thanks for the response. I've observed something peculiar

The moment I initialized `tokenizer.pad_token` to `tokenizer.bos_token` in the above code (i.e `tokenizer.pad_token = tokenizer.bos_token`), I am able to generate proper generations.

While in theory, the pad token shouldn't matter i.e because the labels and attention mask will take care to not attend to it anyways. So, somehow this one line of change has an unexplained behavior or I might be missing something basic.

I will try to debug further and work on the ordered list as suggested, but let me know if you can think of a reason for this behavior!



@NamburiSrinath would you be able to share the exact (stand-alone) example you ran that resulted in `Nan`? We've seen other Llama issues where `Nan` were obtained in seemingly normal situations, but we were never able to pin an example. 

For instance, if I run the script below (based on your script), no `Nan` are returned. 

```py
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = 'meta-llama/Llama-2-7b-hf'
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Update few config parameters to satisfy padding constraints!
tokenizer.padding_side = 'left'
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
model.config.pad_token_id = tokenizer.pad_token_id
model.config.pad_token = tokenizer.pad_token

## Modified generate logic using model.forward()
max_length = 30
def updated_generate(prompts, model):
    # tokenize batch using truncation and padding
    # tokenized_input = tokenizer(prompts, return_tensors="pt").to(model.device)
    tokenized_input = tokenizer(prompts, return_tensors="pt",  padding='max_length', max_length=128).to(model.device)

    # Get input ids and attention mask from tokenizer
    input_ids = tokenized_input.input_ids
    attention_mask = tokenized_input.attention_mask

    # Process position ids from attention mask
    position_ids = attention_mask.long().cumsum(-1) - 1
    position_ids.masked_fill_(attention_mask == 0, 1)

    # Process labels from input ids i.e replace all the padded tokens with -100
    input_ids_masked = torch.zeros(tokenized_input.input_ids.shape, dtype=torch.int64, device=model.device)
    torch.where(tokenized_input.input_ids == tokenizer.pad_token_id,
                torch.tensor(-100, dtype=torch.int64),
                tokenized_input.input_ids,
                out=input_ids_masked)

    for i in range(max_length):
        outputs = model.forward(input_ids=input_ids,
                        attention_mask = attention_mask,
                        labels = input_ids_masked,
                        position_ids=position_ids
                        )

        # Extract logits and do greedy sampling
        logits = outputs.logits[:,-1,:]
        if logits.isnan().any():
            breakpoint()
        predicted_id = logits.argmax(-1)

        # Update input ids, mask, labels and position ids
        # input ids -> add the generated seq
        # attention mask -> add 1
        # labels (input_ids_masked) -> add the generated seq (same as input ids)
        # position ids -> get the max position and add 1

        input_ids = torch.cat([input_ids, predicted_id.unsqueeze(1)], dim=1)
        attention_mask = torch.cat([attention_mask, torch.ones(1, device=model.device).unsqueeze(1)], dim=1)
        input_ids_masked = torch.cat([input_ids_masked, predicted_id.unsqueeze(1)], dim=1)
        position_ids = torch.cat([position_ids, (position_ids.max().unsqueeze(0) + 1).unsqueeze(1)], dim=1)

    print(tokenizer.batch_decode(input_ids))

prompts = ["Once upon a time"]
updated_generate(prompts, model)
```

Also would recommend you to upgrade to latest releases of transformers, might be a rope precision related issue

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.