[Dist] Refactor ZeRO-1 (#5145)

Commit

2 years ago

[Dist] Refactor ZeRO-1 (#5145) * refactor * fix * fix * add padding * more robust save/load