Support safetensors export

[xu-song]

Feature

This commit implements the following features:

• support saving checkpoint as safetensors (more commonly used format)
• support sharding checkpoints (which is important for very large models)

Most of the codes are borrowed from https://github.com/huggingface/transformers/blob/v4.45.1/src/transformers/modeling_utils.py#L2490

Usage

For pytorch_model.bin export

python zero_to_fp32.py . output_dir/

For model.safetensors export

python zero_to_fp32.py . output_dir/ --safe_serialization

[tohtana]

I think DeepSpeed is designed to work in the environment where HF is not installed.

requirements/requirements.txt doesn't have transformer (requirements/requirements-dev.txt does) and there are try-except to imports for transformers (example). Can we have local import to attempt to load transformers and huggingface_hub only when the safetensor option is specified? We can also show an error message when the option is enabled and the modules are not available.

[xu-song]

@tohtana Following your suggestions, I made some optimization for import issue.

[tohtana]

Thank you @xu-song! The formatting test fails, can you check it?

[xu-song]

formatting has been fixed