zero3: defer param release during retain_graph backward #7352

deepspeed/runtime/base_optimizer.py

@@ -237,6 +237,7 @@ class ZeROOptimizer(DeepSpeedOptimizer):
 
     def __init__(self):
         self._backward_hook_state = BackwardHookStateManager()
+        self.retain_graph_on_current_backward = False
 
     # Delegate backward hook state management to the manager.
     # These properties provide backward compatibility with code that accesses
@@ -399,10 +400,14 @@ def backward(self, loss, **kwargs):
 
         scaled_loss = self.backward_prologue(loss)
         retain_graph = kwargs.pop('retain_graph', False)
+        self.retain_graph_on_current_backward = retain_graph
         self.enter_backward()
-        scaled_loss.backward(retain_graph=retain_graph)
-        self.backward_epilogue()
-        self.exit_backward()
+        try:
+            scaled_loss.backward(retain_graph=retain_graph)
+            self.backward_epilogue()
+        finally:
+            self.exit_backward()
+            self.retain_graph_on_current_backward = False
 
     def register_grad_acc_post_hook(self, hook):
         """Register a callback to run when all gradient hooks have fired."""

deepspeed/runtime/engine.py

@@ -2504,6 +2504,11 @@ def _backward_epilogue(self):
             if not bf16_optimizer:
                 self.optimizer.backward_epilogue()
             self.optimizer.exit_backward()
+            # Clear the retained-backward flag here so it is reset for both the
+            # engine.backward() path and the torch-style manual path
+            # (engine.scale(loss).backward(retain_graph=True)), which both reach
+            # this epilogue after the gradient hooks have run.
+            self.optimizer.retain_graph_on_current_backward = False
 
         see_memory_usage("Engine after backward", force=self.memory_breakdown())
         self._stop_timers(self.engine_timers.backward_reduce_timers)
@@ -2662,7 +2667,7 @@ def _flush_coalesced_reduction_zero3(self, optimizer):
                     optimizer.reduce_ready_partitions_and_remove_grads(param)
         optimizer.independent_gradient_partition_epilogue()
 
-    def scale(self, loss):
+    def scale(self, loss, retain_graph=False):
         r"""Apply loss scaler for manual backward pass.
 
         Use this method when calling loss.backward() directly instead of engine.backward().
@@ -2680,6 +2685,12 @@ def scale(self, loss):
 
         Arguments:
             loss: Scalar loss tensor to be scaled
+            retain_graph: bool, default: false
+                Set to true when the upcoming manual backward keeps the graph alive
+                (``scaled_loss.backward(retain_graph=True)``) so a second backward can
+                run over the same forward. For ZeRO-3 this defers parameter release so
+                the retained graph's saved tensors stay valid, matching the behavior of
+                ``engine.backward(loss, retain_graph=True)``.
 
         Returns:
             Scaled loss tensor ready for .backward() call
@@ -2706,6 +2717,9 @@ def scale(self, loss):
         scaled_loss = loss
         if isinstance(self.optimizer, ZeROOptimizer):
             scaled_loss = self.optimizer.scale_if_loss(loss)
+            # The manual path bypasses engine.backward(), so propagate retain_graph here.
+            # Cleared in _backward_epilogue() once the gradient hooks have run.
+            self.optimizer.retain_graph_on_current_backward = retain_graph
         elif self.torch_autocast_z0_gradscaler:
             scaled_loss = self.torch_autocast_z0_gradscaler.scale(loss)
 
@@ -2745,23 +2759,27 @@ def backward(self, loss, retain_graph=False, scale_wrt_gas=True):
         # TODO: handle these scaling with direct calls to loss.backward()
         if isinstance(self.optimizer, ZeROOptimizer):
             loss = self.optimizer.scale_if_loss(loss)
+            self.optimizer.retain_graph_on_current_backward = retain_graph
         elif self.torch_autocast_z0_gradscaler:
             loss = self.torch_autocast_z0_gradscaler.scale(loss)
 
-        with compiled_autograd(self._is_compiled_autograd_enabled, self._compile_kwargs):
-            if self.zero_optimization() or not self.amp_enabled():
-                loss.backward(**backward_kwargs)
-            elif self.amp_enabled():
-                # AMP requires delaying unscale when inside gradient accumulation boundaries
-                # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
-                delay_unscale = not self.is_gradient_accumulation_boundary()
-                with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:
-                    scaled_loss.backward(**backward_kwargs)
-
-            # backward_epilogue is not called in a hook when self._support_torch_style_backward is False
-            self._backward_epilogue()
-
-        self._running_engine_backward = False
+        try:
+            with compiled_autograd(self._is_compiled_autograd_enabled, self._compile_kwargs):
+                if self.zero_optimization() or not self.amp_enabled():
+                    loss.backward(**backward_kwargs)
+                elif self.amp_enabled():
+                    # AMP requires delaying unscale when inside gradient accumulation boundaries
+                    # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
+                    delay_unscale = not self.is_gradient_accumulation_boundary()
+                    with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:
+                        scaled_loss.backward(**backward_kwargs)
+
+                # backward_epilogue is not called in a hook when self._support_torch_style_backward is False
+                self._backward_epilogue()
+        finally:
+            self._running_engine_backward = False
+            if isinstance(self.optimizer, ZeROOptimizer):
+                self.optimizer.retain_graph_on_current_backward = False
 
         return gas_scaled_loss
 

deepspeed/runtime/zero/parameter_offload.py

@@ -552,7 +552,13 @@ def post_sub_module_backward_function(self, sub_module):
             for param in params_to_fetch:
                 param.data = param.data.t() if len(param.ds_shape) != 1 else param.data
 
-        self.get_param_coordinator().release_sub_module(sub_module, forward=False)
+        # Keep gathered params alive when the current backward retains the graph,
+        # so a second backward over the same forward can reuse valid saved tensors.
+        zero_optimizer = getattr(self, "zero_optimizer", None)
+        retain_graph_backward = bool(zero_optimizer is not None
+                                     and getattr(zero_optimizer, "retain_graph_on_current_backward", False))
+        if not retain_graph_backward:
+            self.get_param_coordinator().release_sub_module(sub_module, forward=False)
 
         see_memory_usage(
             f"After sub module backward function {sub_module.__class__.__name__} {sub_module.ds_id} after release",

deepspeed/runtime/zero/stage3.py

@@ -274,6 +274,7 @@ def __init__(
             zero_module_granularity_threshold=zero_module_granularity_threshold,
             log_trace_cache_warnings=log_trace_cache_warnings,
         )
+        self.parameter_offload.zero_optimizer = self
 
         self.persistent_parameters = self.parameter_offload.persistent_parameters
         self._configure_offloading(offload_optimizer_config, offload_param_config)

tests/unit/v1/zero/test_zero_user_backward.py

@@ -539,6 +539,126 @@ def test_separate_loss_function(self, zero_stage):
 
         model_engine.destroy()
 
+    def test_two_losses_separate_backward_gas1(self, zero_stage):
+        """Regression test for https://github.com/deepspeedai/DeepSpeed/issues/7352
+
+        A single forward followed by two separate backward passes, with
+        zero_grad() in between, must produce independent gradients for each
+        loss when gradient_accumulation_steps == 1. Previously the first
+        backward was treated as the accumulation boundary, which froze loss1's
+        gradients so that zero_grad() had no effect and loss2's gradients were
+        doubled (grad2 == grad1 + grad2). Each DeepSpeed gradient set is
+        compared against an equivalent PyTorch DDP baseline.
+        """
+        hidden_dim = 4
+        batch_size = 2
+
+        # Default gradient_accumulation_steps=1, which is the failing case.
+        model_ddp, optimizer_ddp, model_engine, device, dtype = setup_models_and_engines(
+            model_class=SimpleOutputModel, zero_stage=zero_stage, hidden_dim=hidden_dim)
+
+        loss_fn = torch.nn.CrossEntropyLoss()
+
+        # Two different targets so loss1 and loss2 yield distinct gradients;
+        # this makes accidental accumulation (grad1 + grad2) detectable.
+        torch.manual_seed(456)
+        x = torch.randn(batch_size, hidden_dim, device=device, dtype=dtype)
+        y1 = torch.randint(0, hidden_dim, (batch_size, ), device=device)
+        y2 = torch.randint(0, hidden_dim, (batch_size, ), device=device)
+
+        # DDP baseline: separate backward for each loss with zero_grad in between.
+        output_ddp = model_ddp(x)
+        loss1_ddp = loss_fn(output_ddp, y1)
+        loss2_ddp = loss_fn(output_ddp, y2)
+
+        optimizer_ddp.zero_grad()
+        loss1_ddp.backward(retain_graph=True)
+        grads1_ddp = collect_ddp_gradients(model_ddp)
+
+        optimizer_ddp.zero_grad()
+        loss2_ddp.backward()
+        grads2_ddp = collect_ddp_gradients(model_ddp)
+
+        # DeepSpeed: identical sequence.
+        output_ds = model_engine(x)
+        loss1_ds = loss_fn(output_ds, y1)
+        loss2_ds = loss_fn(output_ds, y2)
+
+        model_engine.zero_grad()
+        model_engine.backward(loss1_ds, retain_graph=True)
+        grads1_ds = collect_gradients_safe(model_engine)
+
+        model_engine.zero_grad()
+        model_engine.backward(loss2_ds)
+        grads2_ds = collect_gradients_safe(model_engine)
+
+        # The second backward must NOT accumulate loss1's gradients on top of
+        # loss2; both gradient sets must match their DDP counterparts.
+        compare_gradients(grads1_ddp, grads1_ds, "loss1")
+        compare_gradients(grads2_ddp, grads2_ds, "loss2")
+
+        model_engine.destroy()
+
+    def test_two_losses_separate_manual_backward_gas1(self, zero_stage):
+        """Regression test for https://github.com/deepspeedai/DeepSpeed/issues/7352
+
+        Same scenario as test_two_losses_separate_backward_gas1, but using the
+        torch-style manual path engine.scale(loss, retain_graph=True).backward(
+        retain_graph=True) instead of engine.backward(loss, retain_graph=True).
+        The manual path bypasses engine.backward(), so retain_graph must be
+        propagated through scale(). For ZeRO-3 this defers parameter release so
+        the retained graph's saved tensors stay valid for the second backward
+        over the same forward.
+        """
+        hidden_dim = 4
+        batch_size = 2
+
+        # Default gradient_accumulation_steps=1, which is the failing case.
+        model_ddp, optimizer_ddp, model_engine, device, dtype = setup_models_and_engines(
+            model_class=SimpleOutputModel, zero_stage=zero_stage, hidden_dim=hidden_dim)
+
+        loss_fn = torch.nn.CrossEntropyLoss()
+
+        # Two different targets so loss1 and loss2 yield distinct gradients;
+        # this makes accidental accumulation (grad1 + grad2) detectable.
+        torch.manual_seed(456)
+        x = torch.randn(batch_size, hidden_dim, device=device, dtype=dtype)
+        y1 = torch.randint(0, hidden_dim, (batch_size, ), device=device)
+        y2 = torch.randint(0, hidden_dim, (batch_size, ), device=device)
+
+        # DDP baseline: separate backward for each loss with zero_grad in between.
+        output_ddp = model_ddp(x)
+        loss1_ddp = loss_fn(output_ddp, y1)
+        loss2_ddp = loss_fn(output_ddp, y2)
+
+        optimizer_ddp.zero_grad()
+        loss1_ddp.backward(retain_graph=True)
+        grads1_ddp = collect_ddp_gradients(model_ddp)
+
+        optimizer_ddp.zero_grad()
+        loss2_ddp.backward()
+        grads2_ddp = collect_ddp_gradients(model_ddp)
+
+        # DeepSpeed: identical sequence via the manual scale().backward() path.
+        output_ds = model_engine(x)
+        loss1_ds = loss_fn(output_ds, y1)
+        loss2_ds = loss_fn(output_ds, y2)
+
+        model_engine.zero_grad()
+        model_engine.scale(loss1_ds, retain_graph=True).backward(retain_graph=True)
+        grads1_ds = collect_gradients_safe(model_engine)
+
+        model_engine.zero_grad()
+        model_engine.scale(loss2_ds).backward()
+        grads2_ds = collect_gradients_safe(model_engine)
+
+        # The second backward must NOT accumulate loss1's gradients on top of
+        # loss2; both gradient sets must match their DDP counterparts.
+        compare_gradients(grads1_ddp, grads1_ds, "loss1")
+        compare_gradients(grads2_ddp, grads2_ds, "loss2")
+
+        model_engine.destroy()
+
 
 class LeafModuleModel(torch.nn.Module):
     """Model with ModuleList that uses all parameters - for testing leaf module compatibility"""