DeepSpeed
a20acd1d - Preserve non-floating buffers in ZeRO checkpoint conversion (#8515)

Commit
6 days ago
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>
Author
Parents
Loading