Add configurable dtype for ZeRO checkpoint export (#8318)
## What this changes
Closes #4032.
This adds lower-precision ZeRO checkpoint export without changing the
existing `zero_to_fp32.py` behavior:
- Adds `zero_to_torch.py` with a required `--dtype` option for float32,
float16, or bfloat16 output.
- Adds `convert_zero_checkpoint_to_state_dict(...)` for the same
functionality from Python.
- Keeps `convert_zero_checkpoint_to_fp32_state_dict(...)` as a
backward-compatible fp32 wrapper.
- Copies both recovery scripts into newly saved DeepSpeed checkpoints.
- Uses the requested dtype while planning checkpoint shards as well as
while serializing them.
- Preserves shared-parameter aliases after conversion.
- Documents the CLI, Python API, and memory behavior.
Example:
```bash
./zero_to_torch.py . checkpoint-bf16 --dtype bfloat16
```
## Validation
Focused unit tests cover dtype validation, FP16/BF16 conversion, shared
tensors, saved checkpoint files, and CLI argument forwarding:
```text
3 passed in 3.67s
```
I also ran an end-to-end ZeRO-3 conversion with world size 2 on two RTX
4090 GPUs. The test saved a real distributed checkpoint, reconstructed
both formats, loaded both output files, and checked every tensor.
| Check | Result |
| --- | ---: |
| Trainable parameters | 100,600 |
| FP32 output | 405,019 bytes |
| BF16 output | 203,803 bytes |
| BF16 / FP32 size | 50.32% |
| Maximum absolute error | 0.0002441 |
| Shared-weight alias | Preserved |
| BF16 300KB shard plan | Single shard, as expected |
Repository checks:
- All pre-commit hooks passed, including YAPF, flake8, license,
torch-distributed, and CUDA checks.
- `git diff --check` passed.
- Commit includes the required DCO sign-off.
AI assisted with implementation suggestions and test preparation. I
reviewed and understood every change and accept responsibility for
maintenance and reviewer follow-up.
---------
Signed-off-by: gaoxiaomo <165135449+gaoxiaomo@users.noreply.github.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
