Preserve non-floating buffers in ZeRO checkpoint conversion (#8515)
ZeRO checkpoint reconstruction currently converts every persistent
buffer to fp32. This silently changes integer values before
`load_state_dict()` gets a chance to restore their dtype. For example,
an int64 buffer containing `16777217` becomes `16777216`; reloading a
small model that uses those offsets changes its output from `3.98585796`
to `1.99292898`. The newer dtype-selectable exporter also casts all
buffers while planning and materializing shards, so fixing
reconstruction alone is insufficient.
Restrict precision conversion to floating-point tensors. Integer,
boolean, and other non-floating buffers retain their dtype and values
during reconstruction and export, while floating buffers still recover
to fp32 and follow the requested export dtype. The same dtype selection
is used for empty shard-planning tensors and materialized output,
keeping shard sizes accurate. Shared tensor aliases remain preserved.
The regression tests save actual ZeRO-2 and ZeRO-3 checkpoints after a
small SGD training step, reconstruct and export them, reload the model,
and compare its output. They cover int64 offsets above the fp32
exact-integer range, boolean masks, fp16 source buffers, fp32/fp16/bf16
exports, and sharded-file metadata. The expected precision-converted
model comes from PyTorch's `Module.to(dtype).float()` behavior.
The existing conversion integration test also now uses `tmp_path`: its
previous `tmpdir` object returned a legacy `py.path` stat wrapper that
raises `AttributeError: st_st_size` at the existing `.stat().st_size`
assertions. I reproduced that failure with the unmodified converter
before updating the fixture.
Validation on an Apple M3 Max, 48 GiB RAM, macOS 27, Python 3.12.9,
PyTorch 2.8.0:
- Entire `tests/unit/checkpoint/test_convert_checkpoint.py`: **17
passed**. The new integration cases each execute with **one and two real
CPU/gloo ranks**, for both ZeRO stages and all three output dtypes.
- The same final tests with the original converter: **12 failed, 5
passed**; failures are incorrect non-floating dtypes and changed model
outputs.
- All applicable `pre-commit run --files deepspeed/utils/zero_to_fp32.py
tests/unit/checkpoint/test_convert_checkpoint.py` hooks passed; `git
diff --check` passed.
No CUDA/GPU training, model-convergence run, or safetensors
serialization test was performed. The tests use real DeepSpeed
initialization, training, checkpoint saving, reconstruction, and PyTorch
serialization; no model weights, external services, or compiled
DeepSpeed kernels are needed for these CPU cases.
Developed with OpenAI Codex assistance and independently reviewed by
another coding agent and myself manually. All reported checks were run
locally.
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>