[BUG]: pretraing llama2 using "gemini" plugin, can not resume from saved checkpoints

[jiejie1993]

🐛 Describe the bug

pretrain llama2-7b can resume when using "zero2" plugin, but can not resume when using "gemini" plugin, when using "gemini" plugin, the resume process will stuck, the cuda memory do not change in "nvtop" monitor.

Environment

16 * 8 * H100
torch 2.0.0

[Edenzzzz]

Thanks for your issue. Could you share which script did you run, or a minimal reproducible example?