Deepspeed ZeRO2 + Trainer does not resume training after evaluation
### System Info

- `transformers` version: 4.30.2
- Platform: Linux-4.18.0-305.25.1.el8_4.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.0 (False)
- 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
- Using distributed or parallel set-up in script?: Parallel

### Who can help?

@pacman, @sgugger 

### 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

Steps to reproduce:

1. Use the [pytorch container from Nvidia](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch). 
2. Pip install missing dependencies (IIRC: flash-attn, deepspeed, einops, transformers, accelerate). Note that for deepspeed, I had to use this [specific PR](https://github.com/microsoft/DeepSpeed/issues/3678)
3. Download ShareGPT dataset from huggingface [here](https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered)
4. Run the script of my codebase [here](https://github.com/larrylawl/FastChat/blob/main/scripts/train_vicuna_13b_ds_debug.sh). You'll need to edit the filepaths. My codebase is a fork of FastChat.

### Expected behavior

As you can see from my [log file](https://github.com/huggingface/transformers/files/11766324/train.log), the training gets stuck after coming out of evaluation.

I expected the training script to continue.
When I killed the process, it gives the log 

```
^[[A^[[A^[[A^[[A^[[A^C[2023-06-16 08:32:00,972] [INFO] [launch.py:318:sigkill_handler] Killing subprocess 1947071
Traceback (most recent call last):
  File "/home/users/industry/dso/lannliat/.local/bin/deepspeed", line 6, in <module>
[2023-06-16 08:32:01,072] [INFO] [launch.py:318:sigkill_handler] Killing subprocess 1947071
    main()   
  File "/home/users/industry/dso/lannliat/.local/lib/python3.10/site-packages/deepspeed/launcher/runner.py", line 565, in main
    result.wait()
  File "/usr/lib/python3.10/subprocess.py", line 1207, in wait
    return self._wait(timeout=timeout)
  File "/usr/lib/python3.10/subprocess.py", line 1941, in _wait
    (pid, sts) = self._try_wait(0)
  File "/usr/lib/python3.10/subprocess.py", line 1899, in _try_wait
    (pid, sts) = os.waitpid(self.pid, wait_flags)
```

Seems like the process is waiting.<|||||>cc @pacman100 <|||||>Hello @larrylawl, can you provide a minimal reproducible example? The above example is very involved with a lot of dependencies like flash-attn ... 

When I run the following official example, everything is working fine:

```
cd transformers
export TASK_NAME=mrpc
export CUDA_VISIBLE_DEVICES="0,1"

torchrun --nnodes 1 --nproc-per-node 2 ./examples/pytorch/text-classification/run_glue.py --model_name_or_path bert-base-cased  --task_name $TASK_NAME  --do_train  --do_eval  --max_seq_length 128  --per_device_train_batch_size 16  --learning_rate 5e-5  --num_train_epochs 3  --output_dir /tmp/$TASK_NAME/ --overwrite_output_dir --deepspeed ~/transformers/tests/deepspeed/ds_config_zero2.json --save_steps 10 --evaluation_strategy "steps" --eval_steps 10
```

output:
```
[INFO|trainer.py:1682] 2023-06-16 13:14:12,536 >> ***** Running training *****
[INFO|trainer.py:1683] 2023-06-16 13:14:12,536 >>   Num examples = 3,668
[INFO|trainer.py:1684] 2023-06-16 13:14:12,536 >>   Num Epochs = 3
[INFO|trainer.py:1685] 2023-06-16 13:14:12,536 >>   Instantaneous batch size per device = 16
[INFO|trainer.py:1686] 2023-06-16 13:14:12,536 >>   Total train batch size (w. parallel, distributed & accumulation) = 32
[INFO|trainer.py:1687] 2023-06-16 13:14:12,536 >>   Gradient Accumulation steps = 1
[INFO|trainer.py:1688] 2023-06-16 13:14:12,536 >>   Total optimization steps = 345
[INFO|trainer.py:1689] 2023-06-16 13:14:12,537 >>   Number of trainable parameters = 108,311,810
[INFO|integrations.py:727] 2023-06-16 13:14:12,540 >> Automatic Weights & Biases logging enabled, to disable set os.environ["WANDB_DISABLED"] = "true"
wandb: Currently logged in as: smangrul. Use `wandb login --relogin` to force relogin
wandb: wandb version 0.15.4 is available!  To upgrade, please run:
wandb:  $ pip install wandb --upgrade
wandb: Tracking run with wandb version 0.13.3
wandb: Run data is saved locally in /home/sourab/transformers/wandb/run-20230616_131413-2fg1dtqg
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run distinctive-puddle-305
wandb: ⭐️ View project at https://wandb.ai/smangrul/huggingface
wandb: 🚀 View run at https://wandb.ai/smangrul/huggingface/runs/2fg1dtqg
  3%|█▊                                                            | 10/345 [00:01<01:02,  5.40it/s][INFO|trainer.py:773] 2023-06-16 13:14:21,035 >> The following columns in the evaluation set don't have a corresponding argument in `BertForSequenceClassification.forward` and have been ignored: idx, sentence1, sentence2. If idx, sentence1, sentence2 are not expected by `BertForSequenceClassification.forward`,  you can safely ignore this message.
[INFO|trainer.py:3079] 2023-06-16 13:14:21,037 >> ***** Running Evaluation *****
[INFO|trainer.py:3081] 2023-06-16 13:14:21,037 >>   Num examples = 408
[INFO|trainer.py:3084] 2023-06-16 13:14:21,037 >>   Batch size = 8
{'eval_loss': 0.6776408553123474, 'eval_accuracy': 0.6838235294117647, 'eval_f1': 0.8122270742358079, 'eval_combined_score': 0.7480253018237863, 'eval_runtime': 0.5202, 'eval_samples_per_second': 784.331, 'eval_steps_per_second': 49.982, 'epoch': 0.09}
  3%|█▊                                                            | 10/345 [00:02<01:02,  5.40it/s[INFO|trainer.py:2805] 2023-06-16 13:14:21,560 >> Saving model checkpoint to /tmp/mrpc/checkpoint-10 
[INFO|configuration_utils.py:458] 2023-06-16 13:14:21,561 >> Configuration saved in /tmp/mrpc/checkpoint-10/config.json
[INFO|modeling_utils.py:1844] 2023-06-16 13:14:22,280 >> Model weights saved in /tmp/mrpc/checkpoint-10/pytorch_model.bin
[INFO|tokenization_utils_base.py:2194] 2023-06-16 13:14:22,280 >> tokenizer config file saved in /tmp/mrpc/checkpoint-10/tokenizer_config.json
[INFO|tokenization_utils_base.py:2201] 2023-06-16 13:14:22,281 >> Special tokens file saved in /tmp/mrpc/checkpoint-10/special_tokens_map.json
[2023-06-16 13:14:22,308] [INFO] [logging.py:96:log_dist] [Rank 0] [Torch] Checkpoint global_step10 is about to be saved!
/home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/torch/nn/modules/module.py:1802: UserWarning: Positional args are being deprecated, use kwargs instead. Refer to https://pytorch.org/docs/master/generated/torch.nn.Module.html#torch.nn.Module.state_dict for details.
  warnings.warn(
/home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/torch/nn/modules/module.py:1802: UserWarning: Positional args are being deprecated, use kwargs instead. Refer to https://pytorch.org/docs/master/generated/torch.nn.Module.html#torch.nn.Module.state_dict for details.
  warnings.warn(
[2023-06-16 13:14:22,314] [INFO] [logging.py:96:log_dist] [Rank 0] Saving model checkpoint: /tmp/mrpc/checkpoint-10/global_step10/mp_rank_00_model_states.pt
[2023-06-16 13:14:22,314] [INFO] [torch_checkpoint_engine.py:21:save] [Torch] Saving /tmp/mrpc/checkpoint-10/global_step10/mp_rank_00_model_states.pt...
[2023-06-16 13:14:23,319] [INFO] [torch_checkpoint_engine.py:23:save] [Torch] Saved /tmp/mrpc/checkpoint-10/global_step10/mp_rank_00_model_states.pt.
[2023-06-16 13:14:23,320] [INFO] [torch_checkpoint_engine.py:21:save] [Torch] Saving /tmp/mrpc/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_optim_states.pt...
[2023-06-16 13:14:26,180] [INFO] [torch_checkpoint_engine.py:23:save] [Torch] Saved /tmp/mrpc/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_optim_states.pt.
[2023-06-16 13:14:26,181] [INFO] [engine.py:3228:_save_zero_checkpoint] zero checkpoint saved /tmp/mrpc/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_optim_states.pt
[2023-06-16 13:14:26,181] [INFO] [torch_checkpoint_engine.py:33:commit] [Torch] Checkpoint global_step10 is ready now!
  6%|███▌                                                          | 20/345 [00:08<01:20,  4.06it/s][INFO|trainer.py:773] 2023-06-16 13:14:28,075 >> The following columns in the evaluation set don't have a corresponding argument in `BertForSequenceClassification.forward` and have been ignored: idx, sentence1, sentence2. If idx, sentence1, sentence2 are not expected by `BertForSequenceClassification.forward`,  you can safely ignore this message.
[INFO|trainer.py:3079] 2023-06-16 13:14:28,077 >> ***** Running Evaluation *****
[INFO|trainer.py:3081] 2023-06-16 13:14:28,077 >>   Num examples = 408
[INFO|trainer.py:3084] 2023-06-16 13:14:28,077 >>   Batch size = 8
{'eval_loss': 0.6481188535690308, 'eval_accuracy': 0.6838235294117647, 'eval_f1': 0.8122270742358079, 'eval_combined_score': 0.7480253018237863, 'eval_runtime': 0.5179, 'eval_samples_per_second': 787.721, 'eval_steps_per_second': 50.198, 'epoch': 0.17}
  
...


52%|███████████████████████████████▊                             | 180/345 [02:11<00:44,  3.72it/s][INFO|trainer.py:773] 2023-06-16 13:16:30,237 >> The following columns in the evaluation set don't have a corresponding argument in `BertForSequenceClassification.forward` and have been ignored: idx, sentence1, sentence2. If idx, sentence1, sentence2 are not expected by `BertForSequenceClassification.forward`,  you can safely ignore this message.
[INFO|trainer.py:3079] 2023-06-16 13:16:30,239 >> ***** Running Evaluation *****
[INFO|trainer.py:3081] 2023-06-16 13:16:30,239 >>   Num examples = 408
[INFO|trainer.py:3084] 2023-06-16 13:16:30,239 >>   Batch size = 8
{'eval_loss': 0.36363303661346436, 'eval_accuracy': 0.8186274509803921, 'eval_f1': 0.8673835125448028, 'eval_combined_score': 0.8430054817625975, 'eval_runtime': 0.5171, 'eval_samples_per_second': 789.05, 'eval_steps_per_second': 50.283, 'epoch': 1.57}

```


<|||||>I encouter the same problem. Even with deepspeed and FSDP.
It feels like when model save its weights and stuck.<|||||>Hello @Ricardokevins, please provide a minimal reproducible example. I clearly show above that things are working fine
<|||||>@Ricardokevins I hypothesise that it's a flash attention issue. It works fine with deepspeed only (for my case) and fsdp only (for @pacman100 )<|||||>> @Ricardokevins I hypothesise that it's a flash attention issue. It works fine with deepspeed only (for my case) and fsdp only (for @pacman100 )

This issue may be quite complex and unusual. Initially, when I was training the code in Alpaca-Lora using DeepSpeed, I encountered this problem (training got stuck, and the GPU utilization of some GPUs remained at 0%). This was before using Flash-attn.

Subsequently, I started training the code in FastChat using FSDP (which includes Flash-attn), and encountered similar issues.

Yesterday, I reinstalled all the environments, replaced the code with flash-attn from the FastChat issue, and started training using DeepSpeed. So far, I haven't encountered any problems.
Currently, I'm still unsure about the root cause of the issue, as I haven't faced it recently. If I encounter it again in the future, I will continue the discussion and seek your assistance. Thank you.<|||||>@Ricardokevins Oh nice that you fixed it! Can I ask for some advice since I'm still facing the issue:
- What do you mean by "replaced the code with flash-attn from the FastChat issue"? Did you mean [this patch from FastChat](https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama_flash_attn_monkey_patch.py)?
- What's the cuda version of your system and environment? If you used a docker image, can you share which one worked for you?
- Do you mind sharing your training loss curves? I'm facing a strange issue where my deepspeed + flash attention setting yielded very volatile curves...

![image](https://github.com/huggingface/transformers/assets/40198156/e86acd43-4a4b-4587-9a5d-b56a3487ad04)

But FSDP + flash attention yielded smoother curves

![image](https://github.com/huggingface/transformers/assets/40198156/772f8857-bdd3-4821-b5c4-6e1e73826858)
<|||||>> @Ricardokevins Oh nice that you fixed it! Can I ask for some advice since I'm still facing the issue:
> 
> * What do you mean by "replaced the code with flash-attn from the FastChat issue"? Did you mean [this patch from FastChat](https://github.com/lm-sys/FastChat/blob/main/fastchat/train/llama_flash_attn_monkey_patch.py)?
> * What's the cuda version of your system and environment? If you used a docker image, can you share which one worked for you?
> * Do you mind sharing your training loss curves? I'm facing a strange issue where my deepspeed + flash attention setting yielded very volatile curves...
> 
> ![image](https://user-images.githubusercontent.com/40198156/249803109-e86acd43-4a4b-4587-9a5d-b56a3487ad04.png)
> 
> But FSDP + flash attention yielded smoother curves
> 
> ![image](https://user-images.githubusercontent.com/40198156/249803545-772f8857-bdd3-4821-b5c4-6e1e73826858.png)

1. i use the code from here: https://github.com/lm-sys/FastChat/commit/3adc92d405038d316a3cb908886261231b058590?diff=split
2. cuda version 11.7
<img width="862" alt="image" src="https://github.com/huggingface/transformers/assets/43642508/7b9bf3a7-98ea-4972-9142-ee25325575c2">
<|||||>Thanks @Ricardokevins ! Btw I've fixed the issue by setting number of threads used for intraop parallelism to 1
```
torch.set_num_threads(1) 
```

This [thread](https://discuss.pytorch.org/t/cpu-usage-far-too-high-and-training-inefficient/57228) explains why the above works.

Also, the Vicuna repo now supports [xformers](https://github.com/lm-sys/FastChat/pull/1255). FYI<|||||>> Thanks @Ricardokevins ! Btw I've fixed the issue by setting number of threads used for intraop parallelism to 1
> 
> ```
> torch.set_num_threads(1) 
> ```
> 
> This [thread](https://discuss.pytorch.org/t/cpu-usage-far-too-high-and-training-inefficient/57228) explains why the above works.
> 
> Also, the Vicuna repo now supports [xformers](https://github.com/lm-sys/FastChat/pull/1255). FYI

wow, i will try this if I encounter the problem again!<|||||>> Thanks @Ricardokevins ! Btw I've fixed the issue by setting number of threads used for intraop parallelism to 1
> 
> ```
> torch.set_num_threads(1) 
> ```
> 
> This [thread](https://discuss.pytorch.org/t/cpu-usage-far-too-high-and-training-inefficient/57228) explains why the above works.
> 
> Also, the Vicuna repo now supports [xformers](https://github.com/lm-sys/FastChat/pull/1255). FYI

Hey,i try to use it,but it doesn't work