DeepSpeed

ed827399 - ZeRO-3: reuse the weight partition as the master copy when it already has the master dtype (#8842)

ZeRO-3: reuse the weight partition as the master copy when it already has the master dtype (#8842) # ZeRO-3: reuse the weight partition as the master copy when it already has the master dtype ## Summary `DeepSpeedZeroOptimizer_Stage3._create_fp32_partitions` has a branch that uses a sub-group's weight partition as its master copy when the two would be identical (same dtype, on the device). Its check compared `tensor.device == self.device`, where `self.device` is a string such as `"cuda:0"` (`get_accelerator().current_device_name()`). A `torch.device` never equals a string, so the branch never ran, and ZeRO-3 kept a second copy of every partition. This PR compares with `torch.device(self.device)`. The branch now runs in these ZeRO-3 runs without optimizer offload: | Run | Config | Device memory saved per rank | |---|---|---| | Plain fp32 | no `bf16`, `fp16` or `torch_autocast` section | 4 bytes per parameter in the rank's partition | | fp32 model under torch autocast | `torch_autocast` enabled, no `bf16` / `fp16` section | 4 bytes per parameter in the rank's partition | | bf16 master weights | `"bf16": {"enabled": true, "bf16_master_weights_and_grads": true}` | 2 bytes per parameter in the rank's partition | Runs with `"bf16": {"enabled": true}` and fp32 master weights are not affected: their partitions are bf16 and the master copy is fp32, so the branch does not apply with or without this PR. Runs with optimizer offload or NVMe are not affected either: those branches come first. In these runs the master and the weight partition are the same tensor, so `offload_states` / `reload_states` must not move them apart. This PR makes those sub-groups move with `lp_params` (details below). ## Duplicate check Searched open pull requests on 2026-11-01 ("master copy stage3", "fp32_partitioned_groups_flat", "torch.device self.device", "offload_states alias", "master weights autocast zero3"). None changes this comparison or `offload_states`. #8892 also edits `stage3.py`, but only the restore of fp32 masters from a checkpoint. ## The change `deepspeed/runtime/zero/stage3.py`: | Where | What | |---|---| | `_create_fp32_partitions` | Compare with `torch.device(self.device)`. Record the sub-groups whose master is the weight partition in `aliased_fp32_sub_groups`. | | `offload_states` / `reload_states`, `hp_params` | Move only the masters that have their own storage (`_fp32_partitions_with_own_storage`). An aliased master moves with `lp_params`; moving it alone would split it from the weights. | | `offload_states` / `reload_states`, `lp_params` | After `lp_param_buffer` moves, point each aliased master at its slice of the moved buffer again (`_point_aliased_fp32_partitions_at_lp_buffer`). The master stays the tensor the optimizer holds. | | `offload_states`, `hp_params` without `lp_params` | Log a warning once: the aliased sub-groups stay on the device, so the call frees no memory for them. The DeepSpeed logger prints only warnings by default, so an info log would not show. | `docs/code-docs/source/zero3.rst`: the `hp_params` entry says that, in these runs, the fp32 copies move only with `lp_params`. ## Behaviour changes 1. **Memory.** As in the Summary table. 2. **`offload_states` with `hp_params` alone** frees no device memory for aliased sub-groups (the warning and the docs say so). With `lp_params` in the same call, or with `include=None`, the weights and their masters move together. 3. **`GatheredParameters` edits now persist in these runs.** A weight edit under `deepspeed.zero.GatheredParameters(..., modifier_rank=...)` writes into the weight partition, which is now the master. So the next `step()` updates the edited value. Before this PR, `step()` copied the separate master back over the edit, unless `refresh_fp32_params()` was called. This is intended: there is no second copy to restore from. Runs with bf16 / fp16 weights and fp32 masters keep the old behaviour. `refresh_fp32_params()` still works; for an aliased sub-group it copies the partition onto itself. ## Measured MiniMax-M3 built from its config (random weights), ZeRO-3, bf16 weights with bf16 master weights and Adam states, no CPU offload. The fix was tested in a DeepSpeed study tree before this branch existed. | Setup | Without the fix | With the fix | |---|---|---| | 1 node, 2 H200, first 4 layers, AutoEP 2, 4,096 tokens | 51.62 GiB allocated between steps, 62.04 GiB peak | 41.50 GiB allocated, 51.91 GiB peak; the same losses at every printed step (13.4375, 13.34375, 13.229167, 13.140625) | | 4 nodes x 8 H200, all 60 layers (426B), AutoTP 4 x AutoEP 8, 8,192 tokens per GPU | out of memory in step 2 at 125.0 GiB allocated | trains: 117.9 GiB peak, 15.14 s per step (mean of steps 2-4) | | the same at 16,384 tokens per GPU | | trains: 124.1 GiB peak, 25.17 s per step | ## Tests New tests, all with 2 ranks and ZeRO-3 without optimizer offload: | Test | What it checks | |---|---| | `TestZero3AutocastSharesMasterWeights` (`test_zero.py`) | torch autocast (bf16), fp32 weights: every master partition is the weight partition itself; 2 training steps run. | | `TestZero3Fp32SharesMasterWeights` (`test_zero.py`) | Plain fp32: every master partition is the weight partition; 3 steps match `torch.optim.Adam` on an unsharded copy (losses and final weights, rtol 1e-5, atol 1e-6, gradient clipping off). | | `TestOffloadStatesZero3AliasedMaster` (`test_offload_states.py`; fp32 and autocast, `pin_memory` False / True) | Offload and reload `[lp_params]`, `[hp_params]`, `[lp_params]`, `[hp_params, lp_params]`, then all states, with a step before each. After each round trip the weights are bitwise unchanged and each master is still the weight partition inside `lp_param_buffer`; offloading `lp_params` frees at least the weights' bytes. At the end the weights equal those of a run without offloading, bitwise. The autocast case skips on accelerators without bf16. | On master the first two fail at `assert master is model_partition`, and the third fails at its first check that the master is the weight partition. Hardware: 1 node of 8 H200, torch 2.8.0+cu126, Python 3.10.12, pytest 9.1.1. The driver version was not recorded in this run; the earlier runs on these nodes used 580.126.09. ``` # master c6291b7e1 + only this PR's two test files: pytest unit/v1/zero/test_zero.py unit/v1/zero/test_offload_states.py \ -k "SharesMasterWeights or AliasedMaster" # 6 failed (2 + 4, as above) # this branch, merged with master c6291b7e1: pytest unit/v1/zero/test_zero.py unit/v1/zero/test_offload_states.py \ -k "SharesMasterWeights or OffloadPinnedCopyBack or AliasedMaster" # 8 passed pytest unit/v1/zero/test_offload_states.py # 132 passed pytest unit/v1/zero/test_zero.py # 110 passed, 1 skipped (TestZero3ParamPartitioningBaseBF16, marked "not working") pytest unit/v1/zero/test_zero_autocast.py # 24 passed ``` The `hp_params`-only round trip in `TestOffloadStatesZero3AliasedMaster` prints the new warning once per process. `pre-commit run --files` passes on the changed files. ## Compatibility | Cell | Validated | |---|---| | ZeRO-3, no offload, plain fp32 | yes (tests, including parity with `torch.optim.Adam` and the offload round trips) | | ZeRO-3, no offload, fp32 model under torch autocast | yes (tests) | | ZeRO-3, bf16 master weights (AutoTP + AutoEP) | yes (1 node and 4 nodes, see Measured; before this branch existed) | | ZeRO-3 with optimizer offload, NVMe, or parameter offload | branch not reached; `aliased_fp32_sub_groups` stays empty and `offload_states` behaves as before | | ZeRO-1 / ZeRO-2 | not touched | | Accelerators other than CUDA | `torch.device(name)` accepts every accelerator's device name; not run | | Multi-node | not run for the new tests | ## AI assistance This change, its tests and this description were prepared with AI assistance (Claude). --------- Signed-off-by: pengdurice <pengduhit@gmail.com>