Method backward

deepspeed/runtime/engine.py:3303–3371

r"""Execute backward pass on the loss Arguments: loss: Torch tensor on which to execute backward propagation retain_graph: bool, default: false forward on user defined choice of retain_graph scale_wrt_gas: bool, default: true

(self, loss, retain_graph=False, scale_wrt_gas=True)

    @instrument_w_nvtx
    def backward(self, loss, retain_graph=False, scale_wrt_gas=True):
        r"""Execute backward pass on the loss
        Arguments:
            loss: Torch tensor on which to execute backward propagation
            retain_graph: bool, default: false
                forward on user defined choice of retain_graph
            scale_wrt_gas: bool, default: true
                whether to scale gradients and return value by gradient accumulation steps

        With ``managed_gradient_accumulation=false``, ``backward()`` only accumulates
        gradients locally (it does not trigger the accumulation-boundary reduction); the
        reduction and optimizer update happen in ``step()``.
        """
        assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \
            "must provide optimizer during init in order to use backward"
        assert maybe_loss_for_backward(
            loss), "loss must be a scalar tensor. If you need to pass output gradients, backward() of output tensors"

        with self._running_engine_backward_lock:
            self._running_engine_backward_count += 1
            self._running_engine_backward = True
        engine_backward_graph_state = _ENGINE_BACKWARD_GRAPH_TRACKER.new_state()
        engine_backward_graph_hook = None
        try:
            # Unmanaged mode: count this backward so step() can advance global_samples by the actual micro-batch count.
            if not self.managed_gradient_accumulation():
                self._unmanaged_backward_count += 1

            # Set flag to prevent hooks from firing (we'll manually call prologue/epilogue)
            backward_kwargs = {"retain_graph": retain_graph}
            if self.eigenvalue_enabled():
                backward_kwargs["create_graph"] = True
                backward_kwargs["retain_graph"] = True

            loss = loss / self.gradient_accumulation_steps() if scale_wrt_gas else loss
            gas_scaled_loss = loss

            # TODO: handle these scaling with direct calls to loss.backward()
            if isinstance(self.optimizer, ZeROOptimizer):
                loss = self.optimizer.scale_if_loss(loss)
            elif self.torch_autocast_z0_gradscaler:
                loss = self.torch_autocast_z0_gradscaler.scale(loss)

            if loss.requires_grad:
                engine_backward_graph_hook = loss.register_hook(
                    lambda grad: _ENGINE_BACKWARD_GRAPH_TRACKER.register_current_graph(
                        grad, engine_backward_graph_state))

            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)