[pengdurice]

ZeRO-3: reuse the weight partition as the master copy when it already has the master dtype

Summary

DeepSpeedZeroOptimizer_Stage3._create_fp32_partitions reuses a sub-group's weight partition as its master copy when
the two would be identical (same dtype, on the device): under torch autocast with an fp32 model, and with bf16 master
weights. The check compares 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 runs and
ZeRO-3 keeps a second copy of every partition: +4 bytes per parameter for an fp32 model under autocast, +2 with bf16
master weights. The fix compares with torch.device(self.device).

Duplicate check

Searched open and recent pull requests on 2026-10-15 ("master weights", "fp32 partition", "autocast", "stage3
device"): none touches this comparison.

The change

deepspeed/runtime/zero/stage3.py, one line:

-                        self.fp16_partitioned_groups_flat[i].device == self.device:
+                        self.fp16_partitioned_groups_flat[i].device == torch.device(self.device):

The branch it enables already exists (the code comment above it explains it); offload, NVMe and the other
branches are unchanged.

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

Runs with "bf16": {"enabled": true} and fp32 master weights are not affected: there the partitions are bf16 and
the master copy fp32, so the branch does not apply with or without the fix.

Tests

TestZero3AutocastSharesMasterWeights (2 ranks, tests/unit/v1/zero/test_zero.py): torch autocast (bf16) without a
bf16 / fp16 section, so the model weights stay fp32; ZeRO-3 without offload. Every master partition must be the
weight partition itself (master is model_partition), then 2 training steps run.

Hardware: 1 node of 8 H200 (driver 580.126.09), torch 2.8.0+cu126, pytest 9.1.1.

# master 5f266e26d + only the new test: 1 failed, `assert master is model_partition`
# this branch:
pytest unit/v1/zero/test_zero.py::TestZero3AutocastSharesMasterWeights    # 1 passed
# the rest of the file on this branch:
pytest unit/v1/zero/test_zero.py                                          # 106 passed, 1 skipped (+ 1 failed: an
                                                                          # earlier version of the new test that
                                                                          # assumed fp32 partitions under bf16)

pre-commit run --files passes on both changed files.

Compatibility

Cell	Validated
ZeRO-3, no offload, fp32 model under torch autocast	yes (test)
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	branch not reached (offload and NVMe branches come first)
ZeRO-1 / ZeRO-2	not touched
Accelerators other than CUDA	torch.device(name) accepts every accelerator's device name; not run

AI assistance

This change, its test and this description were prepared with AI assistance (Claude).

[0z5a]

The fix itself is right: torch.device == str was always False, so the alias branch was dead. But that dead branch was also what kept hp and lp storage separate everywhere, and offload_states / reload_states rely on that. Turning it on breaks the offload round trip, and it now also fires for plain fp32 ZeRO-3, not only autocast / bf16-master runs.

1. [blocking] offload_states / reload_states desync the weights once the partitions alias — stage3.py:1162, stage3.py:3881, stage3.py:3948

Once fp32_partitioned_groups_flat[i] is fp16_partitioned_groups_flat[i], offloading hp_params rebinds .data of the shared tensor while the ds_tensors still point into lp_param_buffer. After the next step() they get repointed to the new storage, lp_param_buffer goes stale, and a later lp_params offload/reload restores stale weights.

Repro on the CPU accelerator (1 rank, plain fp32 nn.Sequential, Adam, ZeRO-3, no optimizer offload): step → offload/reload [hp_params] → step → offload/reload [lp_params]. On this branch the module weights change by up to 1e-2 across the lp_params round trip and no longer share storage with lp_param_buffer; on the merge base the same sequence is exact.

Also, while aliased, offloading lp_params can't free device memory (the master still references the old storage), and after a reload the two lists are separate tensors again, so the saving silently disappears. TestDynamicOffloadStatesZero3 only runs bf16, where nothing aliases, so it doesn't catch this. Suggest moving hp and lp together for aliased groups (and re-aliasing after reload), plus an fp32 / autocast parametrization of that test.

2. [should-fix] This now applies to every plain fp32 ZeRO-3 run — stage3.py:1162

The branch also fires without any bf16 / fp16 / autocast section, but the summary lists only autocast with an fp32 model and bf16 master weights, so reviewers may read plain fp32 runs as untouched. Please call that out, and add a numerical parity test (aliased vs. a forced copy over a few steps) in addition to the identity check in TestZero3AutocastSharesMasterWeights.

3. [question] GatheredParameters edits now persist in fp32 runs

With the alias, weight edits under GatheredParameters(..., modifier_rank=0) now go straight into the master copy in fp32 runs. On master the next step() overwrote them unless refresh_fp32_params() was called (in my check the weight mean is -0.011 on master vs. 0.500 kept here); bf16 runs keep the old behaviour. Is that intended? If so it deserves a note. Related: with #8843 the fp32 sub-groups from keep_in_fp32_modules will alias too, and #8843 still allows hp_params offload for mixed dtypes, which is the half that breaks the alias.

[pengdurice]

@0z5a , thank you for the review. fixed accordingly.