Finetuning m2m100 with run_translation_no_trainer.py using ZERO stage 3 hangs when evaluation after first epoch
Author: cokuehuang
Pull request: no

### System Info

- `transformers` version: 4.22.0.dev0
- Platform: Linux-5.15.0-41-generic-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.0+cu113 (True)
- Tensorflow version (GPU?): 2.10.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.4.1 (gpu)
- Jax version: 0.3.5
- JaxLib version: 0.3.5
- Using GPU in script?: <yes>
- Using distributed or parallel set-up in script?: <yes>


### Who can help?

_No response_

### Information

- [x] 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

1. accelerate config
Accelerate configs as follows:
```
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 4
use_cpu: false
```
2. Run finetuning script with command: 
`accelerate launch run_translation_no_trainer.py --model_name_or_path facebook/m2m100_418M    --source_lang ro     --target_lang zh    --train_file teddata/train.json --validation_file teddata/val.json   --output_dir ./m2m100_418M --max_source_length 128     --max_target_length 128    --per_device_train_batch_size=8     --per_device_eval_batch_size=4 --forced_bos_token zh`

Traing output infos:
11/09/2022 11:02:34 - INFO - __main__ - ***** Running training *****
11/09/2022 11:02:34 - INFO - __main__ -   Num examples = 1000
11/09/2022 11:02:34 - INFO - __main__ -   Num Epochs = 3
11/09/2022 11:02:34 - INFO - __main__ -   Instantaneous batch size per device = 8
11/09/2022 11:02:34 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 32
11/09/2022 11:02:34 - INFO - __main__ -   Gradient Accumulation steps = 1
11/09/2022 11:02:34 - INFO - __main__ -   Total optimization steps = 94
 33%|███████████████████████████                                                             32/94[18:31<39:25,  9.20s/it]

Finetuning hangs here, all  GPU-Util is almost 100%.
While accelerate config set zero stage 2, finetuning is success .


### Expected behavior

Success finish finetuning m2m100 with  run_translation_no_trainer.py using ZERO stage 3.

Comments:
cc @pacman100 

Hello @cokuehuang, please provide a minimal script along with the minimal dataset in order to reproduce this issue. I am unable to reproduce this using below steps:
1. run `accelerate env` command to see the config being used:
```
- `Accelerate` version: 0.15.0.dev0
- Platform: Linux-5.4.0-125-generic-x86_64-with-glibc2.31
- Python version: 3.10.4
- Numpy version: 1.23.1
- PyTorch version (GPU?): 1.12.1 (True)
- `Accelerate` default config:
	- compute_environment: LOCAL_MACHINE
	- distributed_type: DEEPSPEED
	- mixed_precision: no
	- use_cpu: False
	- dynamo_backend: NO
	- num_processes: 4
	- machine_rank: 0
	- num_machines: 1
	- gpu_ids: None
	- main_process_ip: None
	- main_process_port: None
	- rdzv_backend: static
	- same_network: True
	- main_training_function: main
	- deepspeed_config: {'gradient_accumulation_steps': 1, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'cpu', 'offload_param_device': 'cpu', 'zero3_init_flag': True, 'zero3_save_16bit_model': True, 'zero_stage': 3}
	- fsdp_config: {}
	- megatron_lm_config: {}
	- downcast_bf16: no
	- tpu_name: None
	- tpu_zone: None
	- command_file: None
	- commands: None
```  

2. Run below command:
```
accelerate launch run_translation_no_trainer.py --model_name_or_path facebook/m2m100_418M --source_lang en --target_lang ro --dataset_name wmt16 --output_dir ./m2m100_418M --max_source_length 128 --max_target_length 128 --per_device_train_batch_size 8 --per_device_eval_batch_size 4 --dataset_config_name ro-en
```
For this to work, change the following line in `run_translation_no_trainer.py` :
```diff
- if isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):
+ if isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast, M2M100Tokenizer)):
```

3. Output logs:
```
11/10/2022 07:21:52 - INFO - __main__ - ***** Running training *****
11/10/2022 07:21:52 - INFO - __main__ -   Num examples = 610320
11/10/2022 07:21:52 - INFO - __main__ -   Num Epochs = 3
11/10/2022 07:21:52 - INFO - __main__ -   Instantaneous batch size per device = 8
11/10/2022 07:21:52 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 32
11/10/2022 07:21:52 - INFO - __main__ -   Gradient Accumulation steps = 1
11/10/2022 07:21:52 - INFO - __main__ -   Total optimization steps = 57219
  0%|▎                                                                                  | 209/57219 [05:48<26:29:02,  1.67s/it]
```

What I think is happening is that at step 32 epoch 1 is over and now the eval loop starts which is using `accelerator.unwrap_model(model).generate()`. Now, this might be taking long time when being offloaded to CPU and as a result one might feel like code has hanged. Can you try ZeRO Stage-3 without offloading anything to `CPU` and let us know if that resolves the issue?


  

@pacman100 
1. accelerate env:
`- `Accelerate` version: 0.12.0
- Platform: Linux-5.15.0-41-generic-x86_64-with-glibc2.17
- Python version: 3.8.13
- Numpy version: 1.23.3
- PyTorch version (GPU?): 1.12.0+cu113 (True)
- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - mixed_precision: no
        - use_cpu: False
        - num_processes: 4
        - machine_rank: 0
        - num_machines: 1
        - main_process_ip: None
        - main_process_port: None
        - main_training_function: main
        - deepspeed_config: {'gradient_accumulation_steps': 1, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'cpu', 'offload_param_device': 'cpu', 'zero3_init_flag': True, 'zero3_save_16bit_model': True, 'zero_stage': 3}
        - fsdp_config: {}
        - downcast_bf16: no
`

2.  My training script and datas:
[scriptanddatas.zip](https://github.com/huggingface/transformers/files/9979662/scriptanddatas.zip)

3. Yes you're right, by adding log , at step 32, it's start eval loop and 'hangs' at generate(). I'll try  ZeRO Stage-3 without offloading cpu later.

1. Without cpu offloading :
`- `Accelerate` version: 0.12.0
- Platform: Linux-5.15.0-41-generic-x86_64-with-glibc2.17
- Python version: 3.8.13
- Numpy version: 1.23.3
- PyTorch version (GPU?): 1.12.0+cu113 (True)
- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - mixed_precision: no
        - use_cpu: False
        - num_processes: 4
        - machine_rank: 0
        - num_machines: 1
        - main_process_ip: None
        - main_process_port: None
        - main_training_function: main
        - deepspeed_config: {'gradient_accumulation_steps': 1, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'none', 'offload_param_device': 'none', 'zero3_init_flag': True, 'zero3_save_16bit_model': True, 'zero_stage': 3}
        - fsdp_config: {}
        - downcast_bf16: no
`
trainging output infos:
11/10/2022 18:21:48 - INFO - __main__ - ***** Running training *****
11/10/2022 18:21:48 - INFO - __main__ -   Num examples = 1000
11/10/2022 18:21:48 - INFO - __main__ -   Num Epochs = 3
11/10/2022 18:21:48 - INFO - __main__ -   Instantaneous batch size per device = 8
11/10/2022 18:21:48 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 32
11/10/2022 18:21:48 - INFO - __main__ -   Gradient Accumulation steps = 1
11/10/2022 18:21:48 - INFO - __main__ -   Total optimization steps = 96
 33%|████████████████████████████████████████████████▎                                                                                                | 32/96 [03:06<06:42,  6.29s/it]

Hangs already 1hour at 33% and eval data size is only 200.

Hello @cokuehuang, Thank you for giving the minimal script and data for reproducing the issue on our end. When using ZeRO stage-3 following needs to passed to `generate` function call:
```
if accelerator.state.deepspeed_plugin.zero_stage == 3:
            gen_kwargs["synced_gpus"] = True #required for ZeRO Stage 3
``` 
after adding it, everything should work just fine when using DS ZeRO-3 with/without cpu offloading
```
11/10/2022 14:09:03 - INFO - __main__ - ***** Running training *****
11/10/2022 14:09:03 - INFO - __main__ -   Num examples = 1000
11/10/2022 14:09:03 - INFO - __main__ -   Num Epochs = 3
11/10/2022 14:09:03 - INFO - __main__ -   Instantaneous batch size per device = 16
11/10/2022 14:09:03 - INFO - __main__ -   Total train batch size (w. parallel, distributed & accumulation) = 32
11/10/2022 14:09:03 - INFO - __main__ -   Gradient Accumulation steps = 1
11/10/2022 14:09:03 - INFO - __main__ -   Total optimization steps = 96
 33%|█████████████████████                                          | 32/96 [01:14<02:28,  2.32s/it]{'max_length': 128, 'num_beams': None, 'synced_gpus': True}
 33%|█████████████████████                                          | 32/96 [01:14<02:28,  2.32s/it]{'max_length': 128, 'num_beams': None, 'synced_gpus': True}
11/10/2022 14:13:04 - INFO - __main__ - {'bleu': 6.697252711851462}
 67%|██████████████████████████████████████████                     | 64/96 [05:13<01:14,  2.32s/it]{'max_length': 128, 'num_beams': None, 'synced_gpus': True}
 67%|██████████████████████████████████████████                     | 64/96 [05:13<01:14,  2.32s/it]{'max_length': 128, 'num_beams': None, 'synced_gpus': True}
11/10/2022 14:16:52 - INFO - __main__ - {'bleu': 6.944214970589274}
100%|███████████████████████████████████████████████████████████████| 96/96 [09:02<00:00,  2.33s/it]{'max_length': 128, 'num_beams': None, 'synced_gpus': True}
100%|███████████████████████████████████████████████████████████████| 96/96 [09:02<00:00,  2.33s/it]{'max_length': 128, 'num_beams': None, 'synced_gpus': True}
11/10/2022 14:20:52 - INFO - __main__ - {'bleu': 6.8998500689065}
Configuration saved in ./m2m100_418M/config.json
100%|███████████████████████████████████████████████████████████████| 96/96 [11:48<00:00,  7.38s/it]
Model weights saved in ./m2m100_418M/pytorch_model.bin
tokenizer config file saved in ./m2m100_418M/tokenizer_config.json
Special tokens file saved in ./m2m100_418M/special_tokens_map.json
100%|███████████████████████████████████████████████████████████████| 96/96 [11:48<00:00,  7.38s/it]
```

@pacman100 Yes! It works! Thanks very much!