Support the parallel conversion from ZeRO checkpoints to FP32/FP16/BF16 param weight

[xylian86]

This PR adds parallel conversion support for the zero_to_fp32 script using universal checkpoints, addressing the request at ISSUE#6526. "The previous version was written 3 years ago models were small and converted fast. Now with 70B+ models the conversion can take hours."

• Switch to universal checkpoint API
• Support Frozen Parameters
• Support Shared Parameters
• Add the support for output to SateTensor
• Add the support for output to FP16/BF16

[loadams]

Hi @xylian86 thanks for the contribution, could you run the pre-commit formatter locally to fix the Formatting error?

[xylian86]

@loadams Thank you for the reminder. I run the pre-commit formatter locally before the PR but It seems that formatting outputs can vary across different environments. After switching to a new machine with ubuntu-22.04 and re-running pre-commit formatter, I've resolved the formatting issues.

[loadams]

@loadams Thank you for the reminder. I run the pre-commit formatter locally before the PR but It seems that formatting outputs can vary across different environments. After switching to a new machine with ubuntu-22.04 and re-running pre-commit formatter, I've resolved the formatting issues.

Hi @xylian86 - looks like there are a few merge conflicts now - could you resolve those, then we can get this merged?