[BUG]: ChatGPT can't load the optimizer weights saved before on Multi-GPU Zero2 mode

[hijkzzz]

🐛 Describe the bug

The root cause is that the state_dict does not match.
We should gather the optimizer weights before save funciton.

Environment

No response

[binmakeswell]

@hijkzzz Thanks for your proposal, we will update save/load next week.