PyTorch-compatible backward API (#7665)

Currently DeepSpeed's backward API has more constraints compared to
PyTorch's normal backward API.
Here is the usage as described in the documentation:
```python
    loss = model_engine(batch)
    model_engine.backward(loss)
```

In this example, 
1. Only accepts a (scalar) loss value
1. Need to call engine's backward API

In contrast, in standard PyTorch, you can do:
```python
    output = model(batch)
    output.backward(out_grad)
```

There are several use cases that rely on this flexibility. For example,
combining multiple models or using loss functions defined separately
from the main model.

If you attempt the same pattern with a DeepSpeed engine, some
preprocessing and postprocessing steps will be silently skipped, which
can lead to incorrect results.

The
[document](https://deepspeed.readthedocs.io/en/latest/training.html#jointly-training-models-with-shared-loss)
explains we can call `_backward_epilogue` manually (possibly
`backward_prologue` as well). However, it's easy for users to miss these
calls, and passing a non-scalar gradient is still not supported.

This PR introduces the same `.backward()` behavior as PyTorch, allowing
.backward() to be called directly on tensors and supporting non-scalar
outputs.

To implement post-backward hooks, we had to use some torch internal
APIs. See
[comments](https://github.com/deepspeedai/DeepSpeed/blob/73f7ff1aab9d1387eb7dd4eca7453a25024533f4/deepspeed/runtime/engine.py#L424)
for more details. When the internal APIs are not available, DeepSpeed
engine only accepts the traditional way `model_engine.backward(loss)`.

---------

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Author: tohtana

deepspeed/runtime/base_optimizer.py

@@ -5,11 +5,13 @@
 
 import os
 import torch
+from typing import Any
 
 from deepspeed.utils import logger
 from deepspeed.utils.tensor_fragment import map_to_flat_opt_states
 from deepspeed.runtime.utils import bwc_tensor_model_parallel_rank, see_memory_usage
 from deepspeed.runtime.torch_autocast import get_comm_dtype, is_autocast_initialized
+from deepspeed.runtime.utils import maybe_loss_for_backward
 
 
 class DeepSpeedOptimizer(object):
@@ -18,6 +20,11 @@ class DeepSpeedOptimizer(object):
 
 class ZeROOptimizer(DeepSpeedOptimizer):
 
+    def __init__(self):
+        self._remaining_grad_acc_hooks = 0
+        self._grad_acc_post_hooks = []
+        self._backward_active_depth = 0
+
     def load_hp_checkpoint_state_from_checkpoint_dir(self, lp_groups_name: str, checkpoint_dir: str) -> None:
         checkpoint_dir = os.path.join(checkpoint_dir, "zero")
         optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt")
@@ -79,3 +86,73 @@ def get_param_comm_dtype(self, param):
             return get_comm_dtype(param)
         else:
             return self.communication_data_type
+
+    def needs_scaler(self) -> bool:
+        """
+        Check if this optimizer requires loss scaling for correct backward pass.
+
+        Returns True if any of the following conditions are met:
+        - Custom loss scaler is enabled
+        - torch.autocast gradient scaler is active (fp16 only)
+        - Dynamic loss scaling is enabled (fp16 with DeepSpeed's loss scaler)
+
+        Returns False for bf16 or fp32, which don't require gradient scaling.
+        """
+        return (self.custom_loss_scaler or self.torch_autocast_gradscaler is not None
+                or (hasattr(self, 'dynamic_loss_scale') and self.dynamic_loss_scale))
+
+    def scale_if_loss(self, value: Any) -> Any:
+        """
+        Applies loss scaling to the input value if it is a loss tensor.
+        """
+        if maybe_loss_for_backward(value):
+            if self.custom_loss_scaler:
+                return self.external_loss_scale * value
+            if self.torch_autocast_gradscaler:
+                return self.torch_autocast_gradscaler.scale(value)
+            return self.loss_scaler.scale_loss(value)
+
+        return value
+
+    def backward_prologue(self):
+        pass
+
+    def backward_epilogue(self, **kwargs):
+        pass
+
+    def backward(self, loss, **kwargs):
+        assert maybe_loss_for_backward(loss), "Optimizer's backward() only accepts a scalar tensor"
+
+        scaled_loss = self.backward_prologue(loss)
+        retain_graph = kwargs.pop('retain_graph', False)
+        self.enter_backward()
+        scaled_loss.backward(retain_graph=retain_graph)
+        self.backward_epilogue()
+        self.exit_backward()
+
+    def register_grad_acc_post_hook(self, hook):
+        self._grad_acc_post_hooks.append(hook)
+
+    def unregister_grad_acc_post_hooks(self):
+        self._grad_acc_post_hooks = []
+
+    def run_grad_acc_post_hooks(self):
+        # Custom autograd Functions (e.g., TiledFusedLogitsLoss) can invoke
+        # `torch.autograd.backward()` from their *forward* pass before the user
+        # ever calls `engine.backward(loss)`. Those early backward calls still
+        # trigger ZeRO's grad hooks, but we must not run the engine's
+        # post-backward logic (which reduces/clears grads) until the outer/user
+        # backward is active. The depth guard filters out only those pre-user
+        # invocations while still allowing backward calls that happen during
+        # the real user backward.
+        if self._backward_active_depth == 0:
+            return
+        for hook in self._grad_acc_post_hooks:
+            hook()
+
+    def enter_backward(self):
+        self._backward_active_depth += 1
+
+    def exit_backward(self):
+        if self._backward_active_depth > 0:
+            self._backward_active_depth -= 1

deepspeed/runtime/bf16_optimizer.py

@@ -316,18 +316,10 @@ def step(self, closure=None):
 
         self.clear_hp_grads()
 
-    def backward(self, loss, retain_graph=False, update_hp_grads=True, clear_lp_grads=False, **bwd_kwargs):
-        """Perform a backward pass and copy the low-precision gradients to the
-        high-precision copy.
-
-        We copy/accumulate to the high-precision grads now to prevent accumulating in the
-        bf16 grads after successive backward() calls (i.e., grad accumulation steps > 1)
-
-        The low-precision grads are deallocated during this procedure.
-        """
+    def backward_prologue(self):
         self.clear_lp_grads()
-        loss.backward(retain_graph=retain_graph, **bwd_kwargs)
 
+    def backward_epilogue(self, update_hp_grads=True, clear_lp_grads=False, **bwd_kwargs):
         if update_hp_grads:
             self.update_hp_grads(clear_lp_grads=clear_lp_grads)