Method backward

deepspeed/runtime/engine.py:2852–2900

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
        """
        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"

        self._running_engine_backward = True
        # Store scale_wrt_gas so the hook can respect it
        self._scale_wrt_gas = scale_wrt_gas

        # 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

        # Used only for return value
        gas_scaled_loss = loss / self.gradient_accumulation_steps() if scale_wrt_gas else 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)

        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

    def is_gradient_accumulation_boundary(self):
        """