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

[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).