[BUG] The training process is stuck without any report

Describe the bug
When training llama2-7B, it gets stuck at the end of the first epoch without any bugs. ZeRO1, ZeRO3, fp16 and bf16 all ran into this problem.
The hardware environment is 8 A100 40G.
The frameworks used are transformers.

Screenshots
the loss print

{'loss': 0.0138, 'learning_rate': 6.989640463132237e-05, 'epoch': 0.97}
{'loss': 0.0135, 'learning_rate': 6.95917123705058e-05, 'epoch': 0.98}
{'loss': 0.015, 'learning_rate': 6.928702010968921e-05, 'epoch': 0.98}
{'loss': 0.0117, 'learning_rate': 6.898232784887264e-05, 'epoch': 0.99}
the tqdm print, the total num of epoch is 3
33%|███▎      | 1126/3384 [05:27<10:53,  3.46it/s]
33%|███▎      | 1127/3384 [05:27<10:53,  3.45it/s]
33%|███▎      | 1128/3384 [05:28<10:53,  3.45it/s]

It will be stuck at these outputs. It's an end of first epoch trainning.

### Tasks

No source file, test, configuration, or reproducible command is provided. Start by reviewing the DeepSpeed training configuration for the llama2-7B run with ZeRO and mixed precision, then reproduce the hang at the first epoch boundary on the stated 8-A100 setup. Done means identifying the cause and adding a verified fix or regression coverage.