From 1409708dc4688bea5f035c1de545f365132225b7 Mon Sep 17 00:00:00 2001 From: pengdurice Date: Thu, 15 Oct 2026 19:31:16 +0000 Subject: [PATCH 1/2] ZeRO-3: reuse the weight partition as the master copy when it already has the master dtype Signed-off-by: pengdurice --- deepspeed/runtime/zero/stage3.py | 2 +- tests/unit/v1/zero/test_zero.py | 46 ++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 215eee472980..b5f9866dbf5a 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1159,7 +1159,7 @@ def _create_fp32_partitions(self): self.fp32_partitioned_groups_flat.append(converted.clone().detach()) self._create_momentum_buffer(num_elements, i, ds_id) elif self.fp16_partitioned_groups_flat[i].dtype == self.master_weights_and_grads_dtype and \ - self.fp16_partitioned_groups_flat[i].device == self.device: + self.fp16_partitioned_groups_flat[i].device == torch.device(self.device): # When torch autocast is enabled, weights in the provided model (and thus groups in the so-called # "fp16" partitioned groups) are already in and updated using fp32. In such cases we don't need # another copy of the weights. diff --git a/tests/unit/v1/zero/test_zero.py b/tests/unit/v1/zero/test_zero.py index 0475dc0b2250..33525c093198 100644 --- a/tests/unit/v1/zero/test_zero.py +++ b/tests/unit/v1/zero/test_zero.py @@ -1919,6 +1919,52 @@ def test(self, force_ds_optim): model, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict) +@pytest.mark.skipif(not get_accelerator().is_bf16_supported(), reason="bf16 is not supported on this accelerator") +class TestZero3AutocastSharesMasterWeights(DistributedTest): + """With torch autocast and no bf16 / fp16 section, the model keeps its weights in float32, the master dtype, so + ZeRO-3 without optimizer offload must use each partition as its own master copy (_create_fp32_partitions) instead + of keeping a second float32 copy. The check compares the partition's torch.device with the optimizer's device + name. (With "bf16": {"enabled": true} as well, the partitions are bf16 and this branch is not taken.)""" + world_size = 2 + + def test(self): + hidden_dim = 64 + config_dict = { + "train_micro_batch_size_per_gpu": 2, + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-3 + } + }, + "torch_autocast": { + "enabled": True, + "dtype": "bfloat16" + }, + "zero_optimization": { + "stage": 3 + }, + } + torch.manual_seed(0) + model = SimpleModel(hidden_dim, nlayers=2) + model, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict) + optimizer = model.optimizer + partitions = [p for p in optimizer.fp16_partitioned_groups_flat if p is not None] + assert partitions and all(p.dtype == torch.float32 for p in partitions) + for model_partition, master in zip(optimizer.fp16_partitioned_groups_flat, + optimizer.fp32_partitioned_groups_flat): + assert master is model_partition + data_loader = random_dataloader(model=model, + total_samples=4, + hidden_dim=hidden_dim, + device=model.device, + dtype=torch.float32) + for batch in data_loader: + loss = model(batch[0], batch[1]) + model.backward(loss) + model.step() + + @pytest.mark.parametrize("training", [True, False]) class TestZeroPartitionCache(DistributedTest): world_size = 1 From 8045140f22e4543538a985798164af2c996bd609 Mon Sep 17 00:00:00 2001 From: pengdurice Date: Wed, 21 Oct 2026 13:46:35 +0000 Subject: [PATCH 2/2] ZeRO-3: keep aliased master partitions with the weights in offload_states With the master partition reused as the weight partition (same dtype, on the device), offload_states and reload_states moved the two apart: an hp_params round trip rebound the shared tensor to new storage while lp_param_buffer kept the old one, so after the next step a weight round trip restored stale weights, and an lp_params offload freed no device memory because the master still held it. Sub-groups whose master partition is the weight partition are now recorded. hp_params offload and reload leave them alone; they move with lp_params, and after lp_param_buffer moves each of them is pointed at its slice again, keeping the tensor the optimizer holds. Tests: offload round trips for fp32 and torch autocast models (weights unchanged, still one copy, the weights' memory freed, training identical to a run without offloading), and fp32 ZeRO-3 training against torch.optim.Adam. Signed-off-by: pengdurice --- deepspeed/runtime/zero/stage3.py | 43 ++++++++--- tests/unit/v1/zero/test_offload_states.py | 87 +++++++++++++++++++++++ tests/unit/v1/zero/test_zero.py | 50 +++++++++++++ 3 files changed, 172 insertions(+), 8 deletions(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index b5f9866dbf5a..dc10882bc5f0 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1094,6 +1094,9 @@ def _create_fp32_partitions(self): nvme_fp16_num_elems = [] nvme_fp32_dest_tensors = [] fp32_element_size = torch.tensor([], dtype=self.master_weights_and_grads_dtype).element_size() + # Sub-groups whose master partition is the weight partition itself (no second copy); offload_states and + # reload_states move those with lp_params. + self.aliased_fp32_sub_groups = set() # Assign portion of subgroup to cpu, the other to gpu. if self.offload_optimizer: @@ -1160,10 +1163,11 @@ def _create_fp32_partitions(self): self._create_momentum_buffer(num_elements, i, ds_id) elif self.fp16_partitioned_groups_flat[i].dtype == self.master_weights_and_grads_dtype and \ self.fp16_partitioned_groups_flat[i].device == torch.device(self.device): - # When torch autocast is enabled, weights in the provided model (and thus groups in the so-called - # "fp16" partitioned groups) are already in and updated using fp32. In such cases we don't need - # another copy of the weights. + # When the weights are already in the master dtype (a float32 model, with or without torch + # autocast), the groups in the so-called "fp16" partitioned groups are updated in place and we + # don't need another copy of the weights. self.fp32_partitioned_groups_flat.append(self.fp16_partitioned_groups_flat[i]) + self.aliased_fp32_sub_groups.add(i) self._create_momentum_buffer(num_elements, i, ds_id) else: converted = self.fp16_partitioned_groups_flat[i].to(self.device, @@ -3851,6 +3855,27 @@ def checkpoint_event_epilogue(self): def empty_partition_cache(self): self.parameter_offload.empty_partition_cache() + def _fp32_partitions_with_own_storage(self): + """The master partitions that are not also a weight partition (see aliased_fp32_sub_groups). An aliased one has + no storage of its own, so moving it apart from lp_params would split the two copies.""" + return [ + partition for i, partition in enumerate(self.fp32_partitioned_groups_flat) + if i not in self.aliased_fp32_sub_groups + ] + + def _point_aliased_fp32_partitions_at_lp_buffer(self): + """After lp_param_buffer moves, point each aliased master partition at its slice of it again. The master + tensor stays the one the optimizer holds, and on the device it is again the sub-group's weight partition.""" + offset = 0 + for i, sub_group in enumerate(self.fp16_groups): + sub_group_numel = sum(param.partition_numel() for param in sub_group) + if i in self.aliased_fp32_sub_groups: + master = self.fp32_partitioned_groups_flat[i] + master.data = self.lp_param_buffer.narrow(0, offset, sub_group_numel) + if i < len(self.fp16_partitioned_groups_flat): + self.fp16_partitioned_groups_flat[i] = master + offset += sub_group_numel + def offload_states(self, include: Container[OffloadStateTypeEnum] = None, device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, @@ -3873,14 +3898,14 @@ def needs_offload(target): if not hasattr(self, "hp_params_pin_buffers"): self.hp_params_pin_buffers = [ get_accelerator().pin_memory(torch.empty_like(t, device=device)) - for t in self.fp32_partitioned_groups_flat + for t in self._fp32_partitions_with_own_storage() ] - for src_tensor, dest_buf in zip(self.fp32_partitioned_groups_flat, self.hp_params_pin_buffers): + for src_tensor, dest_buf in zip(self._fp32_partitions_with_own_storage(), self.hp_params_pin_buffers): dest_buf.copy_(src_tensor, non_blocking=non_blocking) src_tensor.data = dest_buf else: - for buf in self.fp32_partitioned_groups_flat: + for buf in self._fp32_partitions_with_own_storage(): buf.data = buf.data.to(device, non_blocking=non_blocking) self.offloaded_states.add(OffloadStateTypeEnum.hp_params) @@ -3901,6 +3926,7 @@ def needs_offload(target): tensor.data = cpu_buffer.narrow(0, offset, tensor_numel) self.fp16_partitioned_groups_flat.clear() + self._point_aliased_fp32_partitions_at_lp_buffer() self.offloaded_states.add(OffloadStateTypeEnum.lp_params) # LP grad @@ -3944,10 +3970,10 @@ def reload_states(self, non_blocking: bool = False): # HP param if OffloadStateTypeEnum.hp_params in self.offloaded_states: if hasattr(self, "hp_params_pin_buffers"): - for src, dest in zip(self.hp_params_pin_buffers, self.fp32_partitioned_groups_flat): + for src, dest in zip(self.hp_params_pin_buffers, self._fp32_partitions_with_own_storage()): dest.data = src.to(device, non_blocking=non_blocking) else: - for buf in self.fp32_partitioned_groups_flat: + for buf in self._fp32_partitions_with_own_storage(): buf.data = buf.data.to(device, non_blocking=non_blocking) self.offloaded_states.remove(OffloadStateTypeEnum.hp_params) @@ -3957,6 +3983,7 @@ def reload_states(self, non_blocking: bool = False): self, "lp_param_contiguous_pin_buffer") else self.lp_param_buffer self.lp_param_buffer.data = cpu_buffer.data.to(device, non_blocking=non_blocking) self._set_fp16_partitioned_groups_flat() + self._point_aliased_fp32_partitions_at_lp_buffer() parameter_partitions = self._get_parameter_partitions() for tensor, offset, tensor_numel in get_mapping_to_flat_buffer(parameter_partitions): diff --git a/tests/unit/v1/zero/test_offload_states.py b/tests/unit/v1/zero/test_offload_states.py index 88781595ab23..2491e10e1d6d 100644 --- a/tests/unit/v1/zero/test_offload_states.py +++ b/tests/unit/v1/zero/test_offload_states.py @@ -418,3 +418,90 @@ def test_dynamic_offload_states_zero3(self, included_state, pin_memory, non_bloc offloaded_states = None if included_state is None else [included_state] run_model_zero3(model, param_groups, config_dict, hidden_dim, torch.bfloat16, offloaded_states, pin_memory, non_blocking) + + +@pytest.mark.parametrize("precision", ["fp32", "autocast"]) +@pytest.mark.parametrize("pin_memory", [False, True]) +class TestOffloadStatesZero3AliasedMaster(DistributedTest): + """A float32 model (with or without torch autocast) under ZeRO-3 without optimizer offload uses each weight + partition as its own master copy. Offloading and reloading states must keep it one copy: after every round trip + the weights are unchanged and training continues exactly as without offloading, and offloading lp_params frees the + weights' device memory.""" + world_size = 2 + + def test_round_trips_keep_one_copy(self, precision, pin_memory): + hidden_dim = 256 + config_dict = { + "train_micro_batch_size_per_gpu": 1, + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-2 + } + }, + "zero_optimization": { + "stage": 3 + }, + } + if precision == "autocast": + config_dict["torch_autocast"] = {"enabled": True, "dtype": "bfloat16"} + + def build(): + torch.manual_seed(0) + model = SimpleModel(hidden_dim, nlayers=2) + engine, _, _, _ = deepspeed.initialize(model=model, + model_parameters=model.parameters(), + config=config_dict) + return engine + + def train(engine, batch): + loss = engine(batch[0], batch[1]) + engine.backward(loss) + engine.step() + + def one_copy(engine): + optimizer = engine.optimizer + masters = optimizer.fp32_partitioned_groups_flat + weights = optimizer.fp16_partitioned_groups_flat + buffer = optimizer.lp_param_buffer.untyped_storage().data_ptr() + return (all(master is weight for master, weight in zip(masters, weights)) + and all(p.ds_tensor.untyped_storage().data_ptr() == buffer for p in engine.module.parameters())) + + engine = build() + reference = build() + assert one_copy(engine) + batches = list( + random_dataloader(model=engine, + total_samples=6, + hidden_dim=hidden_dim, + device=engine.device, + dtype=torch.float32)) + lp_bytes = engine.optimizer.lp_param_buffer.numel() * engine.optimizer.lp_param_buffer.element_size() + # A weight round trip first (the master holds the same storage, which must not keep it on the device), then + # a master round trip and a step before the next weight round trip (two copies would part there). + rounds = [[OffloadStateTypeEnum.lp_params], [OffloadStateTypeEnum.hp_params], [OffloadStateTypeEnum.lp_params], + [OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.lp_params], None] + for batch, include in zip(batches, rounds): + train(engine, batch) + train(reference, batch) + weights_before = [p.ds_tensor.clone() for p in engine.module.parameters()] + # offload_states first releases the gathered parameters ZeRO-3 keeps after a step; release them here so + # that only the offload itself is measured. + engine.empty_partition_cache() + alloc_before = get_accelerator().memory_allocated() + engine.offload_states(include=include, device=OffloadDeviceEnum.cpu, pin_memory=pin_memory) + freed = alloc_before - get_accelerator().memory_allocated() + if allocator_backed_memory_stats and offload_leaves_accelerator and ( + include is None or OffloadStateTypeEnum.lp_params in include): + assert freed >= lp_bytes, f"offloading lp_params freed {freed} bytes, the weights take {lp_bytes}" + engine.reload_states() + for before, param in zip(weights_before, engine.module.parameters()): + assert torch.equal(before, param.ds_tensor), f"weights changed across offloading {include}" + assert one_copy(engine), f"master and weights are two copies after offloading {include}" + + train(engine, batches[-1]) + train(reference, batches[-1]) + for param, reference_param in zip(engine.module.parameters(), reference.module.parameters()): + assert torch.equal(param.ds_tensor, reference_param.ds_tensor) + engine.destroy() + reference.destroy() diff --git a/tests/unit/v1/zero/test_zero.py b/tests/unit/v1/zero/test_zero.py index 33525c093198..b464fe2822cf 100644 --- a/tests/unit/v1/zero/test_zero.py +++ b/tests/unit/v1/zero/test_zero.py @@ -1965,6 +1965,56 @@ def test(self): model.step() +class TestZero3Fp32SharesMasterWeights(DistributedTest): + """A float32 model with no bf16 / fp16 section is in the master dtype too, so ZeRO-3 without optimizer offload + uses each partition as its own master copy, and training must match torch.optim.Adam on an unsharded copy.""" + world_size = 2 + + def test(self): + hidden_dim = 64 + config_dict = { + "train_micro_batch_size_per_gpu": 2, + "optimizer": { + "type": "Adam", + "params": { + "lr": 1e-2 + } + }, + "zero_optimization": { + "stage": 3 + }, + # The reference does not clip; DeepSpeed clips the global gradient norm to 1.0 by default. + "gradient_clipping": 0.0, + } + torch.manual_seed(0) + model = SimpleModel(hidden_dim, nlayers=2) + device = get_accelerator().current_device_name() + reference = deepcopy(model).to(device) + reference_optimizer = torch.optim.Adam(reference.parameters(), lr=1e-2) + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config_dict) + optimizer = engine.optimizer + for model_partition, master in zip(optimizer.fp16_partitioned_groups_flat, + optimizer.fp32_partitioned_groups_flat): + assert master is model_partition + + for step in range(3): + # The same batch on every rank, so the averaged gradient is the reference's. + generator = torch.Generator().manual_seed(100 + step) + x = torch.randn(2, hidden_dim, generator=generator).to(device) + y = torch.randint(0, hidden_dim, (2, ), generator=generator).to(device) + loss = engine(x, y) + engine.backward(loss) + engine.step() + reference_loss = reference(x, y) + reference_loss.backward() + reference_optimizer.step() + reference_optimizer.zero_grad() + torch.testing.assert_close(loss, reference_loss, rtol=1e-5, atol=1e-6) + + for param, reference_param in zip(engine.module.parameters(), reference.parameters()): + torch.testing.assert_close(safe_get_full_fp32_param(param), reference_param, rtol=1e-5, atol=1e-6) + + @pytest.mark.parametrize("training", [True, False]) class TestZeroPartitionCache(DistributedTest): world_size = 1