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

deepspeed/runtime/zero/stage3.py

@@ -1095,6 +1095,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,11 +1163,12 @@ 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:
-                    # 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.
+                        self.fp16_partitioned_groups_flat[i].device == torch.device(self.device):
+                    # 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,
@@ -3852,6 +3856,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,
@@ -3874,14 +3899,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)
 
@@ -3902,6 +3927,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
@@ -3945,10 +3971,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)
 
@@ -3958,6 +3984,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):

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()

tests/unit/v1/zero/test_zero.py

@@ -1919,6 +1919,102 @@ 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()
+
+
+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