batch generation with Llama: IndexError: index out of range in self
Author: arian-askari
Pull request: no

### System Info

I am using Cuda 11.6, and latest version of transformer which is 4.29.1 to the best of my knowledge.

### Who can help?

@gante 

### Information

- [ ] The official example scripts
- [X] 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 am facing the below error with batch generation:

```
from transformers import AutoTokenizer, LlamaForCausalLM
model_name = "huggyllama/llama-7b"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = LlamaForCausalLM.from_pretrained(model_name)

tokenizer.add_special_tokens({'pad_token': '[PAD]'}) 

prompt = "Hey, are you consciours? Can you talk to me?"
inputs = tokenizer([prompt, prompt + " blah blah"], return_tensors="pt", padding=True, truncation=True)

# Generate
generate_ids = model.generate(inputs.input_ids, max_length=30)
tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
```

### Expected behavior

The above code works if the length of sentences in the batch be equal. E.g., if you initiate inputs variable with the below command then everything works perfectly:

`inputs = tokenizer([prompt, prompt], return_tensors="pt", padding=True, truncation=True)`


Here is the error:

`

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ in <module>:1                                                                                    │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/torch/utils/_contextlib.py:115 in decorate_context                             │
│                                                                                                  │
│   112 │   @functools.wraps(func)                                                                 │
│   113 │   def decorate_context(*args, **kwargs):                                                 │
│   114 │   │   with ctx_factory():                                                                │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                                                   │
│   116 │                                                                                          │
│   117 │   return decorate_context                                                                │
│   118                                                                                            │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/transformers/generation/utils.py:1515 in generate                              │
│                                                                                                  │
│   1512 │   │   │   │   )                                                                         │
│   1513 │   │   │                                                                                 │
│   1514 │   │   │   # 11. run greedy search                                                       │
│ ❱ 1515 │   │   │   return self.greedy_search(                                                    │
│   1516 │   │   │   │   input_ids,                                                                │
│   1517 │   │   │   │   logits_processor=logits_processor,                                        │
│   1518 │   │   │   │   stopping_criteria=stopping_criteria,                                      │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/transformers/generation/utils.py:2332 in greedy_search                         │
│                                                                                                  │
│   2329 │   │   │   model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)  │
│   2330 │   │   │                                                                                 │
│   2331 │   │   │   # forward pass to get next token                                              │
│ ❱ 2332 │   │   │   outputs = self(                                                               │
│   2333 │   │   │   │   **model_inputs,                                                           │
│   2334 │   │   │   │   return_dict=True,                                                         │
│   2335 │   │   │   │   output_attentions=output_attentions,                                      │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/torch/nn/modules/module.py:1501 in _call_impl                                  │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/transformers/models/llama/modeling_llama.py:688 in forward                     │
│                                                                                                  │
│   685 │   │   return_dict = return_dict if return_dict is not None else self.config.use_return   │
│   686 │   │                                                                                      │
│   687 │   │   # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)    │
│ ❱ 688 │   │   outputs = self.model(                                                              │
│   689 │   │   │   input_ids=input_ids,                                                           │
│   690 │   │   │   attention_mask=attention_mask,                                                 │
│   691 │   │   │   position_ids=position_ids,                                                     │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/torch/nn/modules/module.py:1501 in _call_impl                                  │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/transformers/models/llama/modeling_llama.py:531 in forward                     │
│                                                                                                  │
│   528 │   │   │   position_ids = position_ids.view(-1, seq_length).long()                        │
│   529 │   │                                                                                      │
│   530 │   │   if inputs_embeds is None:                                                          │
│ ❱ 531 │   │   │   inputs_embeds = self.embed_tokens(input_ids)                                   │
│   532 │   │   # embed positions                                                                  │
│   533 │   │   if attention_mask is None:                                                         │
│   534 │   │   │   attention_mask = torch.ones(                                                   │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/torch/nn/modules/module.py:1501 in _call_impl                                  │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/torch/nn/modules/sparse.py:162 in forward                                      │
│                                                                                                  │
│   159 │   │   │   │   self.weight[self.padding_idx].fill_(0)                                     │
│   160 │                                                                                          │
│   161 │   def forward(self, input: Tensor) -> Tensor:                                            │
│ ❱ 162 │   │   return F.embedding(                                                                │
│   163 │   │   │   input, self.weight, self.padding_idx, self.max_norm,                           │
│   164 │   │   │   self.norm_type, self.scale_grad_by_freq, self.sparse)                          │
│   165                                                                                            │
│                                                                                                  │
│ /data/arianaskariaa/conda_envs/envs/transformer_4_29_1_efficientLlaMa_int8_python_env/lib/python │
│ 3.8/site-packages/torch/nn/functional.py:2210 in embedding                                       │
│                                                                                                  │
│   2207 │   │   #   torch.embedding_renorm_                                                       │
│   2208 │   │   # remove once script supports set_grad_enabled                                    │
│   2209 │   │   _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)                    │
│ ❱ 2210 │   return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)        │
│   2211                                                                                           │
│   2212                                                                                           │
│   2213 def embedding_bag(                                                                        │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
IndexError: index out of range in self
`

Thanks in advance!

Comments:
Hey @arian-askari 👋 

The exception pops up because you are defining a new token (`[PAD]`), which causes the exception in the embedding layer (it doesn't know the embeddings for the new token until you define them). 

Most decoder-only models have the same "issue" where the padding token is not defined. The standard workaround is as follows:
```py
from transformers import AutoTokenizer, LlamaForCausalLM
model_name = "huggyllama/llama-7b"

tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
model = LlamaForCausalLM.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
model.generation_config.pad_token_id = model.generation_config.eos_token_id

prompt = "Hey, are you consciours? Can you talk to me?"
inputs = tokenizer([prompt, prompt + " blah blah"], return_tensors="pt", padding=True, truncation=True)

# Generate
generate_ids = model.generate(inputs.input_ids, max_length=30)
tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
```

Please note that I added `padding_side="left"` on the tokenizer -- it is critical for decoder-only models like Llama!

Hey @gante,

Thanks a lot! It got fixed with the suggested modification. 

@gante, why need set padding_side="left" for decoder-only models?

@akk-123 these models predict the next token at any given point of the sequence, using the embedding of the latest token as a critical input. If your latest token is a pad token and/or if it is masked by the attention mask, your next token will be unrelated to the sequence -- mostly because the models are not trained to handle this case.

Left-padding ensures the phenomenon above doesn't occur.

@gante thanks a lot! but it seems origin llama model padding side is 'right' ?

Yes -- at train time, you want right-padded sequences. 

I am confuse about it, you mean at train time, we need set right-padding, at inference time,  we should set left-padding? what's more, we will set attention mask when padding, maybe attention mask will avoid the problem you mentioned?
```
these models predict the next token at any given point of the sequence, using the embedding of the latest token as a critical input. If your latest token is a pad token and/or if it is masked by the attention mask, your next token will be unrelated to the sequence -- mostly because the models are not trained to handle this case.
```


The attention mask will not solve it, you need left-padding at generation time.

There's nothing like playing with the model to understand what would happen :)

> Hey @arian-askari 👋
> 
> The exception pops up because you are defining a new token (`[PAD]`), which causes the exception in the embedding layer (it doesn't know the embeddings for the new token until you define them).
> 
> Most decoder-only models have the same "issue" where the padding token is not defined. The standard workaround is as follows:
> 
> ```python
> from transformers import AutoTokenizer, LlamaForCausalLM
> model_name = "huggyllama/llama-7b"
> 
> tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
> model = LlamaForCausalLM.from_pretrained(model_name)
> tokenizer.pad_token = tokenizer.eos_token
> model.generation_config.pad_token_id = model.generation_config.eos_token_id
> 
> prompt = "Hey, are you consciours? Can you talk to me?"
> inputs = tokenizer([prompt, prompt + " blah blah"], return_tensors="pt", padding=True, truncation=True)
> 
> # Generate
> generate_ids = model.generate(inputs.input_ids, max_length=30)
> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
> ```
> 
> Please note that I added `padding_side="left"` on the tokenizer -- it is critical for decoder-only models like Llama!

When using left padding, do we need to set the mask matrix for the left padding, or is there an automatic mask mechanism inside ?


@renmengjie7 the masking mechanism is the same, the only difference is the mask that comes out of the tokenizer (`inputs.attention_mask` in the snippet above) :) 

@gante got it ! Thank you very much.