DeepSpeed

8af59b2e - Restore fp32 masters for ZeRO-3 and the ZeRO-1 BF16 optimizer (#8892)

Restore fp32 masters for ZeRO-3 and the ZeRO-1 BF16 optimizer (#8892) ## Summary - Follow-up to #8891 (merged as `2dd22dc`). The delta against `master` is ZeRO-3 and the dedicated ZeRO-1 BF16 optimizer, plus the review fixes listed below. - `load_checkpoint` still always restores saved fp32 masters. `load_optimizer_states=False` keeps those masters and starts Adam with fresh moments. `load_module_only=True` rebuilds the dedicated BF16 optimizer from the loaded bf16 weights. ZeRO-3 still reads its fp32 partition on a module-only load, because the model checkpoint does not hold the partitioned weights. - `ZeRO-3` and `BF16_Optimizer.load_state_dict(..., load_from_fp32_weights=False)` now rebuild the masters from the bit16 weights, matching ZeRO stage 1 and 2. `DeepSpeedEngine.load_checkpoint` passes true, so a normal resume keeps the saved masters. `zero_optimization.load_from_fp32_weights` stays deprecated and ignored. - Review fixes: a pipeline-serialized load now checks that ZeRO shards exist before waiting on the previous rank, so a missing shard raises instead of leaving ranks blocked in `dist.recv`. A universal load that skips optimizer state keeps `optim_fragment`, so `safe_get_full_optimizer_state` and `safe_set_full_optimizer_state` still work. A leftover debug `print` in `BF16_Optimizer` is removed. AI assistance was used to implement and test this change. This does not duplicate an open DeepSpeed PR. ## Test plan - [x] CPU, torch 2.9.1+cu130, Python 3.12.12, Intel Xeon Platinum 8488C, gloo, world size 1, `tests/unit/checkpoint/test_fp32_master_only.py` and `tests/unit/checkpoint/test_fp32_master_stage3_bf16.py`: 10 passed. Both files ran on the tree that became `a22f0cc4`. - The dedicated BF16 optimizer keeps the saved masters when optimizer state is skipped and rebuilds them for `load_module_only=True`. ZeRO-3 keeps the saved masters for both of those loads. Adam moments stay fresh. - A direct `load_state_dict(..., load_from_fp32_weights=False)` after a module-only load rebuilds the masters from the bit16 weights, for both BF16 and ZeRO-3. - Optimizer globals (`clip_grad`, `overflow`, `dynamic_loss_scale`) are not restored on a masters-only load. - A universal masters-only load into a stepped engine keeps `safe_get_full_optimizer_state` and `safe_set_full_optimizer_state` working. - A pipeline-serialized load with missing shards raises before `dist.recv`. The NVMe masters-only copy restores only the fp32 swap file. - [x] 4×H200, torch 2.9.1+cu130, Python 3.12.12, same two files on tree `a22f0cc4`: 10 passed in 52s. This includes the three data-parallel change tests (universal, legacy, elastic), which spawn 2 ranks on `cuda:0`/`cuda:1`. The #8891 suite there does not cover these two optimizers. Validated cells: dedicated ZeRO-1 BF16 optimizer and ZeRO-3, same data-parallel size, legacy checkpoint, bf16, no offload. ZeRO-1/2 universal masters-only load covered by the fragment test. Open cells: fp16, MoE, real NVMe offload (the NVMe test uses a stub swapper), and a data-parallel size change. --------- Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>