Unmanaged gradient accumulation: ZeRO offload support

deepspeed/runtime/base_optimizer.py

@@ -252,6 +252,23 @@ class ZeROOptimizer(DeepSpeedOptimizer):
 
     def __init__(self):
         self._backward_hook_state = BackwardHookStateManager()
+        # Mirrored copy of the engine GAS boundary for managed reduce/offload paths.
+        # Engine owns the source of truth (micro-step / step() / set_*); ZeRO reads this
+        # during backward. Prefer get/set methods over touching the private field.
+        self._is_gradient_accumulation_boundary = True
+
+    def is_gradient_accumulation_boundary(self) -> bool:
+        """Whether the current micro-batch is a gradient accumulation boundary.
+
+        Used by managed ZeRO reduce/partition/offload logic. Unmanaged mode still
+        mirrors True while ``engine.step()`` runs so late readers stay consistent;
+        deferred offload finalize does not branch on this flag.
+        """
+        return self._is_gradient_accumulation_boundary
+
+    def set_gradient_accumulation_boundary(self, is_boundary: bool) -> None:
+        """Mirror the engine's gradient accumulation boundary into this optimizer."""
+        self._is_gradient_accumulation_boundary = bool(is_boundary)
 
     # Delegate backward hook state management to the manager.
     # These properties provide backward compatibility with code that accesses

deepspeed/runtime/engine.py

@@ -1662,12 +1662,6 @@ def _do_sanity_check(self):
                 f'Client Optimizer (type = {type(self.client_optimizer)} is not instantiated but Client LR Scheduler is instantiated'
 
         if not self.managed_gradient_accumulation():
-            offload_optimizer = self.zero_offload_optimizer()
-            offload_param = self.zero_offload_param()
-            assert offload_optimizer is None or offload_optimizer.device == OffloadDeviceEnum.none, \
-                "managed_gradient_accumulation=False is not supported with ZeRO optimizer state offload"
-            assert offload_param is None or offload_param.device == OffloadDeviceEnum.none, \
-                "managed_gradient_accumulation=False is not supported with ZeRO parameter offload"
             assert self.zero_optimization_partition_gradients() or not self.zero_overlap_comm(), \
                 "managed_gradient_accumulation=False supports ZeRO overlap_comm only with ZeRO stage 2"
             assert not self.pipeline_parallelism, \
@@ -2856,7 +2850,8 @@ def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE):
             return
 
         # Pass (PP) gas boundary flag to optimizer (required for zero)
-        self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
+        if hasattr(self.optimizer, "set_gradient_accumulation_boundary"):
+            self.optimizer.set_gradient_accumulation_boundary(self.is_gradient_accumulation_boundary())
         if self.is_gradient_accumulation_boundary():
             self._reduce_autoep_folding_tp_replicated_gradients()
         # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well
@@ -2923,7 +2918,7 @@ def _backward_prologue(self):
             self.optimizer.zenflow_state ^= 1
 
         if self.zero_optimization():
-            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
+            self.optimizer.set_gradient_accumulation_boundary(self.is_gradient_accumulation_boundary())
 
         self._start_timers(self.engine_timers.backward_inner_timers)
 
@@ -3058,7 +3053,7 @@ def coalesce_grad_reduction(self):
             optimizer._coalesce_grad_reduction = False
             self.inside_no_sync_ctxt = False
             self._is_gradient_accumulation_boundary = True
-            optimizer.is_gradient_accumulation_boundary = True
+            optimizer.set_gradient_accumulation_boundary(True)
             try:
                 # Drive a single reduction pass over locally accumulated grads.
                 # Iterate explicitly (rather than calling reduce_gradients) so
@@ -3270,7 +3265,8 @@ def set_gradient_accumulation_boundary(self, is_boundary):
             "set_gradient_accumulation_boundary() is not supported with managed_gradient_accumulation=False; " \
             "the caller owns the boundary by calling step()"
         self._is_gradient_accumulation_boundary = is_boundary
-        self.optimizer.is_gradient_accumulation_boundary = is_boundary
+        if hasattr(self.optimizer, "set_gradient_accumulation_boundary"):
+            self.optimizer.set_gradient_accumulation_boundary(is_boundary)
 
     def zero_grad(self):
         """
@@ -3392,7 +3388,7 @@ def step(self, lr_kwargs=None):
         # Unmanaged mode: step() is the accumulation boundary.
         self._running_engine_step = True
 
-        # Unmanaged boundary: stage 2/3 already reduced/partitioned per backward so only finalize; stage 0/1/DDP reduce here.
+        # Unmanaged boundary: stage 2/3 already reduced/partitioned per backward so only finalize (incl. offload); stage 0/1/DDP reduce here.
         if not self.managed_gradient_accumulation():
             if self.zero_optimization_partition_gradients():
                 self.optimizer.finalize_gradient_accumulation_boundary()

deepspeed/runtime/superoffload/superoffload_stage3.py

@@ -151,7 +151,7 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
                 grad_buffer.copy_(cuda_grad_buffer, non_blocking=True)
                 grad_buffer = cuda_grad_buffer
 
-            if self.is_gradient_accumulation_boundary:
+            if self.is_gradient_accumulation_boundary():
                 self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer)
 
                 fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(
@@ -162,7 +162,7 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
             if self.sub_group_grad_partition_counts[i] == self.sub_group_to_param_num[i]:
                 completed_sub_groups.append(i)
 
-        if self.is_gradient_accumulation_boundary and completed_sub_groups:
+        if self.is_gradient_accumulation_boundary() and completed_sub_groups:
             get_accelerator().current_stream().synchronize()
             for i in completed_sub_groups:
                 if self.subgroup_to_device[i] == 'cpu' and not self.clip_grad:

deepspeed/runtime/zenflow/engine_stage3.py

@@ -325,7 +325,7 @@ def _process_selected_fp32_groups_grad(optimizer_z3, params_to_update, grad_part
 
             curr_buffer_idx += 1
 
-    if not optimizer_z3.is_gradient_accumulation_boundary:
+    if not optimizer_z3.is_gradient_accumulation_boundary():
         optimizer_z3.selective_optimizer.group_step(params_to_update)
     else:
         optimizer_z3.selective_optimizer.temp_copy_param(params_to_update)
@@ -494,7 +494,7 @@ def partition_grads(optimizer_z3, params_to_release: List[Parameter], grad_parti
             i, dest_offset, _ = optimizer_z3.grad_position[optimizer_z3.get_param_id(param)]
             now_state = optimizer_z3.get_overlap_step_state()
 
-            if optimizer_z3.is_gradient_accumulation_boundary:
+            if optimizer_z3.is_gradient_accumulation_boundary():
                 optimizer_z3.norm_for_param_grads[optimizer_z3.get_param_id(
                     param)] = optimizer_z3._constant_buffered_norm2(grad_buffer)
 

deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py

@@ -411,7 +411,7 @@ def _process_selected_fp32_groups_grad(self, tensor, total_size, communication_d
 
                 self.param_id_sum_buffer_offset[param_id] = []
 
-        if not self.is_gradient_accumulation_boundary:
+        if not self.is_gradient_accumulation_boundary():
             self.selective_optimizer.group_step(group_to_paramlist)
         else:
             self.selective_optimizer.temp_copy_param(group_to_paramlist)

deepspeed/runtime/zero/mics.py

@@ -408,7 +408,7 @@ def allreduce_mics_shard_grads(self, params, partitioned_grads_buffers: List[Ten
         """
         """
         # TODO: improve the condition check
-        if not self.is_gradient_accumulation_boundary or \
+        if not self.is_gradient_accumulation_boundary() or \
             len(partitioned_grads_buffers) == 0:
             return
 

deepspeed/runtime/zero/stage3.py

@@ -375,6 +375,8 @@ def _enforce_optimizer_offload():
         self.postscale_gradients = postscale_gradients
         self.gradient_accumulation_steps = gradient_accumulation_steps
         self.micro_step_id = 0
+        # ds_ids reduced since the last step(); used to finalize only active params in unmanaged offload mode.
+        self._offload_boundary_param_ids = set()
         self.reduce_bucket_size = int(reduce_bucket_size)
 
         if self.all2all_process_group is not None:
@@ -439,8 +441,6 @@ def _enforce_optimizer_offload():
         if self.swap_optimizer:
             self._configure_tensor_swapping(offload_optimizer_config, aio_config)
 
-        self.is_gradient_accumulation_boundary: bool = True
-
         # Toggled by DeepSpeedEngine.coalesce_grad_reduction().
         self._coalesce_grad_reduction = False
 
@@ -1353,10 +1353,52 @@ def independent_gradient_partition_epilogue(self):
         self._epilogue_ran_this_backward = True
 
     def finalize_gradient_accumulation_boundary(self):
-        # Unmanaged mode: grad partitions already accumulate across backwards via __param_id_to_grad_partition; nothing to finalize for non-offload.
-        assert not self.offload_optimizer and not self.offload_param, \
-            "unmanaged gradient accumulation does not support ZeRO offload"
-        self.is_gradient_accumulation_boundary = True
+        # Unmanaged mode: partitions already accumulate in __param_id_to_grad_partition; offload still needs deferred boundary copy.
+        # Mirror engine boundary for any managed-style readers during step(); finalize itself does not branch on it.
+        self.set_gradient_accumulation_boundary(True)
+        if self.offload_optimizer:
+            self._finalize_offload_gradient_accumulation()
+
+    def _offload_grad_partition_at_boundary(self, param, grad_buffer, offload_fp32_gradients, offload_fp32_offsets):
+        # Boundary-only: record grad norm and copy/swap into optimizer FP32 or NVMe buffers.
+        i, dest_offset, _ = self.grad_position[self.get_param_id(param)]
+        self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer)
+
+        if self._swappable_optimizer_subgroup(i):
+            if i not in offload_fp32_gradients.keys():
+                offload_fp32_gradients[i] = []
+                offload_fp32_offsets[i] = []
+
+            offload_fp32_gradients[i].append(grad_buffer.to(dtype=self.master_weights_and_grads_dtype))
+            offload_fp32_offsets[i].append(dest_offset)
+        else:
+            fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(0, dest_offset, grad_buffer.numel())
+            fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype), non_blocking=True)
+
+    def _swap_out_offload_fp32_gradients(self, offload_fp32_gradients, offload_fp32_offsets):
+        if not (self.offload_optimizer and self.swap_optimizer):
+            return
+        for i in offload_fp32_gradients.keys():
+            self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i],
+                                                      gradient_offsets=offload_fp32_offsets[i],
+                                                      gradient_tensors=offload_fp32_gradients[i])
+
+    def _finalize_offload_gradient_accumulation(self):
+        # Deferred boundary work for params reduced this window (matches managed offload; skips inactive params).
+        offload_fp32_gradients = {}
+        offload_fp32_offsets = {}
+        for param_group in self.fp16_groups:
+            for param in param_group:
+                if param.ds_id not in self._offload_boundary_param_ids:
+                    continue
+                if param.ds_id not in self.__param_id_to_grad_partition:
+                    continue
+                grad_buffer = self.__param_id_to_grad_partition[param.ds_id]
+                if not get_accelerator().on_accelerator(grad_buffer):
+                    grad_buffer = grad_buffer.to(get_accelerator().current_device_name(), non_blocking=True)
+                self._offload_grad_partition_at_boundary(param, grad_buffer, offload_fp32_gradients,
+                                                         offload_fp32_offsets)
+        self._swap_out_offload_fp32_gradients(offload_fp32_gradients, offload_fp32_offsets)
 
     def overlapping_partition_gradients_reduce_epilogue(self):
         self.independent_gradient_partition_epilogue()
@@ -1822,6 +1864,9 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
                 param.grad = None
                 continue
 
+            # Record active param so unmanaged offload finalize skips params unused this window.
+            self._offload_boundary_param_ids.add(param.ds_id)
+
             # move or accumulate gradient partition to target buffer
             grad_buffer = self.__param_id_to_grad_partition[param.ds_id].narrow(0, 0, grad_partition.numel())
             buffers.append(grad_buffer)
@@ -1844,35 +1889,17 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
 
             # offload the gradient partition if applicable
             if self.offload_optimizer:
-                i, dest_offset, _ = self.grad_position[self.get_param_id(param)]
-
-                if self.is_gradient_accumulation_boundary:
-                    self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer)
-
-                    if self._swappable_optimizer_subgroup(i):
-                        if i not in offload_fp32_gradients.keys():
-                            offload_fp32_gradients[i] = []
-                            offload_fp32_offsets[i] = []
-
-                        offload_fp32_gradients[i].append(grad_buffer.to(dtype=self.master_weights_and_grads_dtype))
-                        offload_fp32_offsets[i].append(dest_offset)
-                    else:
-                        fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(
-                            0, dest_offset, grad_buffer.numel())
-                        fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype),
-                                               non_blocking=True)
+                if self.is_gradient_accumulation_boundary():
+                    self._offload_grad_partition_at_boundary(param, grad_buffer, offload_fp32_gradients,
+                                                             offload_fp32_offsets)
 
             # free the gradient
             if not get_accelerator().is_synchronized_device():
                 if param.grad is not None:
                     param.grad.record_stream(get_accelerator().current_stream())
             param.grad = None
 
-        if self.offload_optimizer and self.swap_optimizer:
-            for i in offload_fp32_gradients.keys():
-                self.optimizer_swapper.swap_out_gradients(parameter=self.fp32_partitioned_groups_flat[i],
-                                                          gradient_offsets=offload_fp32_offsets[i],
-                                                          gradient_tensors=offload_fp32_gradients[i])
+        self._swap_out_offload_fp32_gradients(offload_fp32_gradients, offload_fp32_offsets)
         return buffers
 
     def _partitioned_buffers_all_gather(self, params: List[Parameter], buffers_to_allgather: List[Tensor],
@@ -2306,6 +2333,7 @@ def reset_cpu_buffers(self):
 
     def _pre_step(self):
         self.micro_step_id = 0
+        self._offload_boundary_param_ids = set()
         # Also reset the epilogue flag so the next iteration starts fresh.
         # Without this, the flag from the last backward before step() would cause
         # an increment in the next forward(), which is wrong.