ZeRO-1/2: preserve gradients across request boundaries until step

deepspeed/runtime/zero/stage_1_and_2.py

@@ -650,6 +650,8 @@ def _enforce_cpu_offload():
             self.grad_position = {}
             # Param ids reduced since the last step(); used to finalize only active params in unmanaged mode.
             self._offload_accumulated_param_ids = set()
+            # FP32-master gradients first seen at a boundary stay in the partition until another backward.
+            self._offload_partition_only_param_ids = set()
             self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel,
                                                                 device=self.device,
                                                                 dtype=self.dtype)
@@ -1052,11 +1054,11 @@ def independent_gradient_partition_epilogue(self):
                         device=get_accelerator().current_device_name(),
                         param_group_idx=i,
                         return_tensor_list=True)
-                    # Clear all_grad_tensors after use. With reentrant checkpointing,
-                    # the epilogue may run multiple times per backward pass. Each time,
-                    # we read the cumulative grad_accum (which PyTorch naturally accumulates)
-                    # and the final phase will have all gradients.
-                    self.all_grad_tensors[i] = None
+                    # grad_accum is cumulative across reentrant checkpoint phases, so
+                    # reread it rather than adding it twice. Otherwise keep the running
+                    # sum until step(): more request boundaries may precede the update.
+                    if self.use_grad_accum_attribute:
+                        self.all_grad_tensors[i] = None
 
         self._release_ipg_buffers()
 
@@ -1099,6 +1101,8 @@ def _finalize_cpu_offload_gradient_accumulation(self):
                     continue
                 if self.get_param_id(param) not in self._offload_accumulated_param_ids:
                     continue
+                if self.get_param_id(param) in self._offload_partition_only_param_ids:
+                    continue
                 if self._use_compact_grad_partition_buffer(group_idx):
                     self.copy_grads_in_partition(param, group_idx)
                     continue
@@ -1829,7 +1833,13 @@ def buffer_to_accumulate_to_in_cpu():
         #accumulate gradients into param.grad_accum or parts of it that belongs to this partition
         def accumulate_gradients():
             grad_accum = self.get_param_gradient_attribute(param)
-            if not self.low_precision_master_weights_and_grads:
+            if param_id in self._offload_partition_only_param_ids:
+                # Recover only this rank's fragment, converting back to the accumulation dtype.
+                previous = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements)
+                local_buffer = dest_buffer.narrow(0, source_offset, num_elements)
+                local_buffer.copy_(previous, non_blocking=True)
+                grad_accum.data.view(-1).narrow(0, source_offset, num_elements).add_(local_buffer)
+            elif not self.low_precision_master_weights_and_grads:
                 dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True)
                 grad_accum.data.view(-1).add_(dest_buffer)
             else:
@@ -1852,9 +1862,12 @@ def copy_gradients_to_cpu():
         if param_id not in self.accumulated_grads_in_cpu:
             self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu()
 
-        if self.micro_step_id > 0:
+        # Other parameters may have run backwards before this parameter's first use.
+        # Only its participation in this window makes the persistent host buffer current.
+        if param_id in self._offload_accumulated_param_ids:
             accumulate_gradients()
         copy_gradients_to_cpu()
+        self._offload_partition_only_param_ids.discard(param_id)
         self._record_offload_copy()
 
     def set_norm_for_param_grad(self, param):
@@ -2159,22 +2172,30 @@ def copy_grads_in_partition(self, param, group_idx=None):
                 self.clear_grad_attribute(param)
             return
         if self.cpu_offload:
-            # Accumulate when there were prior backwards in this step (restore from
-            # CPU buffer) or more will follow (save to CPU buffer). Skipping only
-            # the lone backward of a step preserves the existing fast path for
-            # ga_steps=1 + single backward.
-            if self.micro_step_id > 0 or not self.is_gradient_accumulation_boundary():
+            param_id = self.get_param_id(param)
+            first_use = param_id not in self._offload_accumulated_param_ids
+            boundary = self.is_gradient_accumulation_boundary()
+            if first_use and boundary:
+                if self.low_precision_master_weights_and_grads:
+                    # Low-precision masters accumulate in the optimizer's partition grad.
+                    [i, _, dest_offset, num_elements] = self.grad_position[param_id]
+                    self.accumulated_grads_in_cpu[param_id] = self.single_partition_of_fp32_groups[i].grad.view(
+                        -1).narrow(0, dest_offset, num_elements)
+                else:
+                    # The boundary copy is sufficient unless another backward follows before step().
+                    self._offload_partition_only_param_ids.add(param_id)
+            else:
                 self.async_accumulate_grad_in_cpu_via_gpu(param)
-                # Record active param so unmanaged finalize skips params unused this window.
-                self._offload_accumulated_param_ids.add(self.get_param_id(param))
 
-            if self.is_gradient_accumulation_boundary():
+            if boundary:
                 self.set_norm_for_param_grad_in_gpu(param)
 
                 self.update_offload_overflow_tracker_for_param_grad(param)
 
                 self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param)
 
+            # Include the shared-buffer fast path in subsequent accumulation and finalization.
+            self._offload_accumulated_param_ids.add(param_id)
             return
         if self._use_compact_grad_partition_buffer(group_idx):
             if self.use_grad_accum_attribute:
@@ -3018,6 +3039,7 @@ def step(self, closure=None):
         self._wait_for_offload_copies()
         if self.cpu_offload:
             self._offload_accumulated_param_ids = set()
+            self._offload_partition_only_param_ids = set()
 
         if not self.compute_grad_norm:
             self._global_grad_norm = None

tests/unit/v1/zero/test_zero2_gradient_safety.py

@@ -50,6 +50,7 @@ def make_optimizer(dtype=torch.bfloat16, device=None, low_precision=False, cpu_o
     opt.micro_step_id = 0
     opt._pending_offload_events = {}
     opt._offload_accumulated_param_ids = set()
+    opt._offload_partition_only_param_ids = set()
     opt.accumulated_grads_in_cpu = {}
     opt.norm_for_param_grads = {}
     opt.local_overflow = False
@@ -138,12 +139,73 @@ def test_oversized_copy_is_independent_and_preserves_routing(dtype, cpu_offload,
     assert opt.ipg_buckets[torch.float32].params == [(0, 0, 0)]
 
 
-def test_default_offload_accumulates_across_nonboundary_backwards():
-    opt, param, master = make_optimizer()
-    backward_values(opt, param, [1, 2, 4], [False, False, True])
+@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"])
+@pytest.mark.parametrize("first_boundary", [False, True], ids=["nonboundary-first", "boundary-first"])
+def test_default_offload_accumulates_across_nonboundary_backwards(low_precision, first_boundary):
+    opt, param, master = make_optimizer(low_precision=low_precision)
+    backward_values(opt, param, [1, 2, 4], [first_boundary, False, True])
     assert torch.equal(master.grad, torch.full_like(master, 7))
 
 
+@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"])
+def test_boundary_offload_can_be_finalized_without_another_backward(low_precision):
+    opt, param, master = make_optimizer(low_precision=low_precision)
+    backward_values(opt, param, [3], [True])
+    opt.finalize_gradient_accumulation_boundary()
+    opt.finalize_gradient_accumulation_boundary()
+    opt._wait_for_offload_copies()
+    assert torch.equal(master.grad, torch.full_like(master, 3))
+
+
+@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"])
+def test_single_boundary_offload_avoids_auxiliary_storage_and_copy(monkeypatch, low_precision):
+    opt, param, master = make_optimizer(low_precision=low_precision)
+    prepare_step(opt, param, master, monkeypatch)
+    ref = torch.nn.Parameter(master.detach().clone())
+    ref_optimizer = torch.optim.Adam([ref], lr=0.01)
+    original_copy = torch.Tensor.copy_
+    copies = []
+
+    def count_copy(destination, source, *args, **kwargs):
+        copies.append(destination.data_ptr())
+        return original_copy(destination, source, *args, **kwargs)
+
+    for value in [2, 3, 4]:
+        copies.clear()
+        with monkeypatch.context() as patch:
+            patch.setattr(torch.Tensor, "copy_", count_copy)
+            backward_values(opt, param, [value], [True])
+        # Pin the resource regression: one partition copy, no additional host storage.
+        assert copies == [master.grad.data_ptr()]
+        for buffer in opt.accumulated_grads_in_cpu.values():
+            assert buffer.untyped_storage().data_ptr() == master.grad.untyped_storage().data_ptr()
+        assert torch.equal(master.grad, torch.full_like(master, value))
+        opt.finalize_gradient_accumulation_boundary()
+        opt.step()
+        ref.grad = torch.full_like(ref, value)
+        ref_optimizer.step()
+        assert torch.equal(master, ref)
+
+
+@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"])
+@pytest.mark.parametrize("source_offset,dest_offset,num_elements", [(0, 0, 8), (2, 1, 3), (5, 4, 3)])
+def test_boundary_first_offload_preserves_owned_fragment_and_accumulation_dtype(low_precision, source_offset,
+                                                                                dest_offset, num_elements):
+    opt, param, master = make_optimizer(low_precision=low_precision)
+    opt.grad_position[0] = [0, source_offset, dest_offset, num_elements]
+    master.grad.fill_(123)
+    reference = torch.nn.Parameter(torch.zeros_like(param))
+    # BF16 rounds 256 + 1 back to 256; accumulating in the FP32 master would differ.
+    for value, boundary in [(256, True), (1, False), (2, True)]:
+        (reference * value).sum().backward()
+        backward_values(opt, param, [value], [boundary])
+        opt.finalize_gradient_accumulation_boundary()
+        opt._wait_for_offload_copies()
+        expected = torch.full_like(master, 123)
+        expected.narrow(0, dest_offset, num_elements).copy_(reference.grad.narrow(0, source_offset, num_elements))
+        assert torch.equal(master.grad, expected)
+
+
 def test_reduction_and_offload_consume_owned_oversized_buffer():
     opt, param, master = make_optimizer()
     opt.is_param_in_current_partition = {0: True}
@@ -306,11 +368,19 @@ def test_record_stream_is_separate_from_readiness(monkeypatch):
     assert not stream.waited
 
 
-def test_cpu_accumulator_waits_before_pageable_source_is_read():
+@pytest.mark.parametrize("first_boundary", [False, True], ids=["host-accumulator", "optimizer-partition"])
+def test_cpu_accumulator_waits_before_pageable_source_is_read(first_boundary):
     opt, param, master = make_optimizer()
-    backward_values(opt, param, [1], [False])
-    # The earlier copy into the CPU accumulator only lands once its event completes.
-    event = FakeEvent(lambda: opt.accumulated_grads_in_cpu[0].fill_(3))
+    backward_values(opt, param, [1], [first_boundary])
+
+    # The earlier CPU copy only lands once its event completes.
+    def complete_copy():
+        if first_boundary:
+            master.grad.fill_(3)
+        if 0 in opt.accumulated_grads_in_cpu:
+            opt.accumulated_grads_in_cpu[0].fill_(3)
+
+    event = FakeEvent(complete_copy)
     opt._pending_offload_events["producer"] = event
     opt.micro_step_id = 1
     opt.set_gradient_accumulation_boundary(True)
@@ -369,6 +439,54 @@ def test_step_consumes_gradients_after_offload_copies_complete(monkeypatch):
     assert torch.equal(master, ref)
 
 
+@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"])
+@pytest.mark.parametrize("first_boundary", [False, True], ids=["nonboundary-first", "boundary-first"])
+@pytest.mark.parametrize("overflow", [False, True], ids=["successful-step", "overflow-step"])
+def test_offload_first_use_after_another_backward_starts_fresh(monkeypatch, low_precision, first_boundary, overflow):
+    opt, param, master = make_optimizer(low_precision=low_precision)
+    prepare_step(opt, param, master, monkeypatch)
+    ref = torch.nn.Parameter(master.detach().clone())
+    ref_optimizer = torch.optim.Adam([ref], lr=0.01)
+    # Warm the persistent auxiliary buffer before testing first use in a later window.
+    backward_values(opt, param, [8], [False])
+    opt.finalize_gradient_accumulation_boundary()
+    opt.step()
+    ref.grad = torch.full_like(ref, 8)
+    ref_optimizer.step()
+    assert torch.equal(master, ref)
+    if overflow:
+        # A boundary-only overflow must clear partition-only state without reusing the warmed buffer.
+        backward_values(opt, param, [7], [True])
+        opt.finalize_gradient_accumulation_boundary()
+        opt.local_overflow = True
+        opt.step()
+        assert opt.overflow
+        assert torch.equal(master, ref)
+    # Other parameters have already contributed to the new window; this one has not.
+    opt.micro_step_id = 1
+    opt.set_gradient_accumulation_boundary(first_boundary)
+    param.grad = torch.full_like(param, 2)
+    opt.copy_grads_in_partition(param)
+    param.grad = None
+    opt.micro_step_id = 2
+    opt.set_gradient_accumulation_boundary(True)
+    param.grad = torch.full_like(param, 4)
+    opt.copy_grads_in_partition(param)
+    opt._wait_for_offload_copies()
+    assert torch.equal(master.grad, torch.full_like(master, 6))
+    opt.step()
+    ref.grad = torch.full_like(ref, 6)
+    ref_optimizer.step()
+    assert torch.equal(master, ref)
+    # A later boundary-only window must not reuse the warmed auxiliary sum.
+    backward_values(opt, param, [3], [True])
+    opt.finalize_gradient_accumulation_boundary()
+    opt.step()
+    ref.grad = torch.full_like(ref, 3)
+    ref_optimizer.step()
+    assert torch.equal(master, ref)
+
+
 @pytest.mark.parametrize("operation", ["reset", "legacy_load", "universal_load"])
 def test_cpu_buffer_reuse_waits_for_offload_copies(operation):
     opt, _, _ = make_optimizer()