Seq2Seq trainer cannot generate tokens with deepspeed and Zero3
Author: Lingy12
Pull request: no

### System Info

- `transformers` version: 4.31.0
- Platform: Linux-5.4.0-54-generic-x86_64-with-glibc2.27
- Python version: 3.10.11
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.0
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.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?: Yes 4 x A100 40GB
- Using distributed or parallel set-up in script?: Yes

### Who can help?

@ArthurZucker @pacman100 @gante 

### 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 sorry that I cannot give out the code directly here since the code have some confidential information, however, I think it's a general issue , rather than usecase dependent. Please let me know anything that is unclear so that I can give more detail.

I am using **torch.distributed.lanch** to lanch a finetuning on LLama2-7b model. I tried to use Seq2SeqTrainer and try to get the generated tokens. It does not work when I set **generation_max_length: Optional[int] = field(default=1024)**. The deepspeed is used here, provided the deepspeed config

```
{
    "bf16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 100,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1e-10
    },
    "zero_optimization": {
        "stage": 3,
        "allgather_partitions": true,
        "allgather_bucket_size": 1e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 1e8,
        "contiguous_gradients": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
```

The error I get is 

```
lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 633, in forward
    past_key_values_length = past_key_values[0][0].shape[2]
TypeError: 'NoneType' object is not subscriptable
```

I have tried some walkaround:

1. Setting **generation_max_length=20**, the code will give warning and it seems that it does not generate anything after the input. 
2.  Customize predict_step function for seq2seqtrainer, add **if self.state.is_world_process_zero**. The code will hang in the generate function. I think it's due to some synchronization issue between processes. 

Please let me know if more detail is required.


### Expected behavior

Seq2SeqTrainer should be able to generate tokens.

Comments:
I tried to upgrade the transformers and deepspeed version, I think the training loop stuck when evaluate is triggered. 

Hi @Lingy12 👋 

I see that the exception is `TypeError: 'NoneType' object is not subscriptable`. Can you confirm:
1. The type of `past_key_values` before the exception?
2. If `past_key_values` is `None`, whether the model's `config.use_cache` is `True`? 

Hi @gante.

Refering to this  
https://github.com/huggingface/transformers/blob/8eb9e29d8dc8b8bd98b4dd48317d1d596ec548f3/src/transformers/models/llama/modeling_llama.py#L845

The type of it should be List[torch.floatTensor]. 

I am using the default llama2 config. Hence config.use_cache should be true.


Hey @Lingy12 👋 

I'm aware of what the answers to the questions I asked above should be :) However, since you didn't share a script to reproduce the issue, the best I can do is to ask you to double-check a few things for us.

Hi @gante 

Sure. I try to log the past key type with 
`logger.info("past key Type " + str(type(past_key_values)))` 

And here is the output. 

<img width="341" alt="image" src="https://github.com/huggingface/transformers/assets/54443474/fb98bc07-4174-4e3a-80c2-80488995eb01">



If you store your model before the first evaluation round, can you generate with it?

If yes 👉 then it is a distributed-related issue, and we will need a short script to reproduce the issue
If no 👉 I may be able to provide further pointers depending on the error. A short script for reproducibility would help.

Yes. The model can generate if it's a saved checkpoint. It only have problem when using model parallel.


Let me try to put up a sample script here. 


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.

Hi,
I am facing the same issue, 

`[/usr/local/lib/python3.10/dist-packages/transformers/models/llama/modeling_llama.py](https://localhost:8080/#) in prepare_inputs_for_generation(self, input_ids, past_key_values, attention_mask, inputs_embeds, **kwargs)
   1082     ):
   1083         if past_key_values is not None:
-> 1084             past_length = past_key_values[0][0].shape[2]`

Is there a fix/what can be done to solve this problem? 
Thanks.


Hey @Nitish5499 if you want help we are going to need a reproducer!