Unmanaged gradient accumulation: ZeRO offload support

deepspeed/runtime/engine.py

@@ -1648,12 +1648,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, \
@@ -3378,7 +3372,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/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:
@@ -1353,10 +1355,51 @@ 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"
+        # Unmanaged mode: partitions already accumulate in __param_id_to_grad_partition; offload still needs deferred boundary copy.
         self.is_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 +1865,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 +1890,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)
+                    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 +2334,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.

deepspeed/runtime/zero/stage_1_and_2.py

@@ -579,6 +579,8 @@ def _enforce_cpu_offload():
             self.norm_for_param_grads = {}
             self.local_overflow = False
             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()
             self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel,
                                                                 device=self.device,
                                                                 dtype=self.dtype)
@@ -948,9 +950,11 @@ def independent_gradient_partition_epilogue(self):
         see_memory_usage("End ipg_epilogue")
 
     def finalize_gradient_accumulation_boundary(self):
-        # Unmanaged mode: grads were reduced/accumulated into all_grad_tensors each backward; finalize averaged_gradients for step().
-        assert not self.cpu_offload, "unmanaged gradient accumulation does not support ZeRO optimizer state offload"
+        # Unmanaged mode: grads accumulated each backward; finalize for step() (averaged_gradients or offload fp32 copy).
         self.is_gradient_accumulation_boundary = True
+        if self.cpu_offload:
+            self._finalize_cpu_offload_gradient_accumulation()
+            return
         for i, _ in enumerate(self.bit16_groups):
             self.averaged_gradients[i] = self.get_flat_partition(self.params_in_partition[i],
                                                                  self.first_offset[i],
@@ -961,6 +965,36 @@ def finalize_gradient_accumulation_boundary(self):
                                                                  return_tensor_list=True)
             self.all_grad_tensors[i] = None
 
+    def _finalize_cpu_offload_gradient_accumulation(self):
+        # Deferred boundary work for params reduced this window (matches managed offload; skips inactive params).
+        for group in self.params_in_partition:
+            for param in group:
+                if not param.requires_grad:
+                    continue
+                if self.get_param_id(param) not in self._offload_accumulated_param_ids:
+                    continue
+                self._restore_cpu_offload_grad_to_gpu(param)
+                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)
+
+    def _restore_cpu_offload_grad_to_gpu(self, param):
+        # Last non-boundary epilogue cleared param.grad; reload accumulated CPU grads for boundary helpers.
+        param_id = self.get_param_id(param)
+        [_, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
+        dest_buffer = self.temp_grad_buffer_for_gpu_offload.view(-1).narrow(0, 0, param.numel())
+        if not self.low_precision_master_weights_and_grads:
+            dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True)
+        else:
+            dest_buffer.narrow(0, source_offset, num_elements).copy_(self.accumulated_grads_in_cpu[param_id].view(-1),
+                                                                     non_blocking=True)
+        # Clone so the shared temp buffer can be reused for the next parameter.
+        restored = dest_buffer.view_as(param).clone()
+        if self.use_grad_accum_attribute:
+            param.grad_accum = restored
+        else:
+            param.grad = restored
+
     def clear_backward_seen_flag(self):
         """Clear the backward seen flag and do deferred cleanup.
 
@@ -1595,6 +1629,8 @@ def copy_grads_in_partition(self, param):
             # ga_steps=1 + single backward.
             if self.micro_step_id > 0 or not self.is_gradient_accumulation_boundary:
                 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:
                 self.set_norm_for_param_grad_in_gpu(param)
@@ -2220,6 +2256,8 @@ def step(self, closure=None):
         Not supporting closure.
         """
         self.micro_step_id = INITIAL_MICRO_STEP_ID
+        if self.cpu_offload:
+            self._offload_accumulated_param_ids = set()
 
         see_memory_usage("In step before checking overflow")
 

docs/_pages/config-json.md

@@ -33,7 +33,7 @@ toc_label: "Contents"
 
 | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              | Default |
 | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
-| Controls how gradient accumulation boundaries are managed. When `true`, DeepSpeed tracks micro-steps and applies the optimizer step only at the accumulation boundary, so `forward`/`backward`/`step` can be called symmetrically on every micro-batch. When `false`, micro-step tracking is disabled and the client is responsible for calling `step()` at the accumulation boundary; each `step()` finalizes the locally-accumulated gradients and applies an optimizer update. The `false` setting currently supports ZeRO stage 0/1/2/3 (and DDP); ZeRO offload (optimizer state and parameter) is planned but not yet available (rejected at initialization). It is also incompatible with pipeline parallelism, DeepCompile, and Apex AMP. ZeRO `overlap_comm` is supported only with ZeRO stage 2 (rejected for stage 0/1, where reduction is deferred to `step()`). | `true`  |
+| Controls how gradient accumulation boundaries are managed. When `true`, DeepSpeed tracks micro-steps and applies the optimizer step only at the accumulation boundary, so `forward`/`backward`/`step` can be called symmetrically on every micro-batch. When `false`, micro-step tracking is disabled and the client is responsible for calling `step()` at the accumulation boundary; each `step()` finalizes the locally-accumulated gradients and applies an optimizer update. The `false` setting supports ZeRO stage 0/1/2/3 (and DDP), including ZeRO optimizer-state and parameter offload (CPU/NVMe). It is incompatible with pipeline parallelism, DeepCompile, and Apex AMP. ZeRO `overlap_comm` is supported only with ZeRO stage 2 (rejected for stage 0/1, where reduction is deferred to `step()`). | `true`  |
 
 
 

docs/code-docs/source/training.rst

@@ -104,11 +104,13 @@ This is useful when ``backward`` and ``step`` must be decoupled and the number o
 client- or RPC-driven RL backends where a single optimizer step arrives as ``N`` ``backward()``
 calls followed by one ``step()``, with ``N`` unknown at configuration time.
 
-Unmanaged mode currently supports **ZeRO stage 0/1/2/3 and DDP**. For ZeRO stage 0/1 (and DDP),
-``backward()`` accumulates gradients locally and ``step()`` performs the gradient all-reduce followed
-by the optimizer update. For ZeRO stage 2/3, gradients are reduced/partitioned on every ``backward()``
-(as in managed mode) and ``step()`` finalizes the accumulated partition gradients before the optimizer
-update; ZeRO ``overlap_comm`` is supported in the stage-2 path.
+Unmanaged mode currently supports **ZeRO stage 0/1/2/3 and DDP**, including ZeRO optimizer-state
+and parameter offload (CPU/NVMe). For ZeRO stage 0/1 (and DDP), ``backward()`` accumulates gradients
+locally and ``step()`` performs the gradient all-reduce followed by the optimizer update. For ZeRO
+stage 2/3, gradients are reduced/partitioned on every ``backward()`` (as in managed mode) and
+``step()`` finalizes the accumulated partition gradients (including deferred offload norm/FP32 or
+NVMe copy when offload is enabled) before the optimizer update; ZeRO ``overlap_comm`` is supported
+in the stage-2 path.
 
 .. note::
    By default ``backward()`` scales the loss and gradients by the configured
@@ -120,12 +122,10 @@ update; ZeRO ``overlap_comm`` is supported in the stage-2 path.
    applies and no manual averaging is needed.)
 
 .. note::
-   Unmanaged mode is being added incrementally. ZeRO stage 0/1/2/3 (and DDP) is supported today;
-   ZeRO offload (optimizer state and parameter) is planned but **not yet available** -- enabling it
-   with ``managed_gradient_accumulation=false`` raises an ``AssertionError`` at initialization.
-   Unmanaged mode is likewise incompatible with pipeline parallelism, DeepCompile, and Apex AMP,
-   which are also rejected at initialization. ZeRO ``overlap_comm`` is supported only with ZeRO
-   stage 2 (rejected for stage 0/1, where reduction is deferred to ``step()``).
+   Unmanaged mode supports ZeRO stage 0/1/2/3 (and DDP), including ZeRO optimizer-state and parameter
+   offload (CPU/NVMe). It is incompatible with pipeline parallelism, DeepCompile, and Apex AMP, which
+   are rejected at initialization. ZeRO ``overlap_comm`` is supported only with ZeRO stage 2
+   (rejected for stage 0/1, where reduction is deferred to ``step()``).
 
 .. autofunction:: deepspeed.DeepSpeedEngine.set_gradient_accumulation_boundary