Merge pull request #26 from nathon-lee/feature/zero-multi-loss-separate-backward-test

Fix ZeRO-3 so two separate backward passes on the same forward graph work correctly when `retain_graph=True` is used on the first backward.

Author: nathon-lee

deepspeed/runtime/base_optimizer.py

@@ -252,6 +252,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
@@ -419,10 +420,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

@@ -222,11 +222,6 @@ def active_timers(self):
         return self.micro_timers + self.global_timers
 
 
-def _eigenvalue_summary_events(block_eigenvalue, global_samples):
-    return [(f"Train/Eigenvalues/ModelBlockParam_{i}", ev_value[0], global_samples)
-            for i, ev_value in enumerate(block_eigenvalue.values())]
-
-
 class DeepSpeedEngine(Module):
     r"""DeepSpeed engine for training."""
 
@@ -2759,6 +2754,7 @@ def _backward_epilogue(self):
             if not bf16_optimizer:
                 self.optimizer.backward_epilogue()
             self.optimizer.exit_backward()
+            self.optimizer.retain_graph_on_current_backward = False
 
         if self.is_deepcompile_active():
             deepcompile_backward_epilogue()
@@ -2920,7 +2916,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().
@@ -2938,6 +2934,8 @@ def scale(self, loss):
 
         Arguments:
             loss: Scalar loss tensor to be scaled
+            retain_graph: bool, default: false
+                forward on user defined choice of retain_graph
 
         Returns:
             Scaled loss tensor ready for .backward() call
@@ -2963,6 +2961,7 @@ def scale(self, loss):
         # Apply loss scaler based on optimizer type
         scaled_loss = loss
         if isinstance(self.optimizer, ZeROOptimizer):
+            self.optimizer.retain_graph_on_current_backward = retain_graph
             scaled_loss = self.optimizer.scale_if_loss(scaled_loss)
         elif self.torch_autocast_z0_gradscaler:
             scaled_loss = self.torch_autocast_z0_gradscaler.scale(scaled_loss)
@@ -3003,25 +3002,29 @@ 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
-
-        return gas_scaled_loss
+        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()
+
+            return gas_scaled_loss
+        finally:
+            self._running_engine_backward = False
+            if isinstance(self.optimizer, ZeROOptimizer):
+                self.optimizer.retain_graph_on_current_backward = False
 
     def is_gradient_accumulation_boundary(self):
         """
@@ -3223,8 +3226,13 @@ def step(self, lr_kwargs=None):
 
                     if (self.eigenvalue_enabled()
                             and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution()):
-                        self.summary_events.extend(
-                            _eigenvalue_summary_events(self.block_eigenvalue, self.global_samples))
+                        ev_values = self.block_eigenvalue.values()
+                        for i in range(len(ev_values)):
+                            self.summary_events.append((
+                                f"Train/Eigenvalues/ModelBlockParam_{i}",
+                                self.ev_values[i][0],
+                                self.global_samples,
+                            ))
                     self.monitor.write_events(self.summary_events)
 
         # Check flops profiling
@@ -4327,9 +4335,6 @@ def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True,
         process with rank 0.
 
         """
-        if not save_dir:
-            raise ValueError(f"save_dir must be a non-empty string, got {save_dir!r}")
-
         if self._optimizer_has_ckpt_event_prologue():
             # Custom preparation for checkpoint save, if applicable
             self.optimizer.checkpoint_event_prologue()

deepspeed/runtime/zero/parameter_offload.py

@@ -136,6 +136,7 @@ def __init__(
         zero_quantized_nontrainable_weights=False,
         zero_module_granularity_threshold=0,
         log_trace_cache_warnings=False,
+        retain_graph_checker=None,
     ):
 
         see_memory_usage("DeepSpeedZeRoOffload initialize [begin]", force=False)
@@ -153,6 +154,7 @@ def __init__(
         self.zero_quantized_weights = zero_quantized_weights
         self.zero_quantized_nontrainable_weights = zero_quantized_nontrainable_weights
         self.log_trace_cache_warnings = log_trace_cache_warnings
+        self.retain_graph_checker = retain_graph_checker
 
         if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none:
             self.offload_device = offload_param_config.device
@@ -563,7 +565,11 @@ 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.
+        retain_graph_backward = bool(self.retain_graph_checker()) if self.retain_graph_checker is not None else 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

@@ -285,6 +285,7 @@ def __init__(
             zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights,
             zero_module_granularity_threshold=zero_module_granularity_threshold,
             log_trace_cache_warnings=log_trace_cache_warnings,
+            retain_graph_checker=lambda: self.retain_graph_on_current_backward,
         )
 
         self.persistent_parameters = self.parameter_offload.persistent_parameters
@@ -571,6 +572,7 @@ def initialize_ds_offload(
         zero_quantized_nontrainable_weights,
         zero_module_granularity_threshold,
         log_trace_cache_warnings,
+        retain_graph_checker=None,
     ):
         return DeepSpeedZeRoOffload(module=module,
                                     timers=timers,
@@ -589,7 +591,8 @@ def initialize_ds_offload(
                                     zero_quantized_weights=zero_quantized_weights,
                                     zero_quantized_nontrainable_weights=zero_quantized_nontrainable_weights,
                                     zero_module_granularity_threshold=zero_module_granularity_threshold,
-                                    log_trace_cache_warnings=log_trace_cache_warnings)
+                                    log_trace_cache_warnings=log_trace_cache_warnings,
+                                    retain_graph_checker=retain_graph_checker)
 
     def _get_param_partition_group(self, param):
         return getattr(param, "ds_process_group", self.dp_process_group)
@@ -3128,8 +3131,7 @@ def state_dict(self):
             torch.save(checkpoint, "saved.pth")
         """
         if self.elastic_checkpoint:
-            raise NotImplementedError(
-                "ZeRO-3 elastic checkpointing is deprecated and unsupported. Use Universal Checkpointing instead.")
+            raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.")
 
         return self._rigid_state_dict()
 
@@ -3289,8 +3291,7 @@ def load_state_dict(self,
         """
 
         if self.elastic_checkpoint:
-            raise NotImplementedError(
-                "ZeRO-3 elastic checkpointing is deprecated and unsupported. Use Universal Checkpointing instead.")
+            raise NotImplementedError("ZeRO-3 does not yet support elastic checkpointing, please disable for now.")
 
         if checkpoint_folder:
             self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights)

tests/unit/v1/zero/test_zero_user_backward.py

@@ -539,6 +539,128 @@ 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"""