Method scale

deepspeed/runtime/engine.py:2799–2849

r"""Apply loss scaler for manual backward pass. Use this method when calling loss.backward() directly instead of engine.backward(). This applies the appropriate loss scaler for mixed precision training, allowing you to manually control the backward pass while still benefitin

(self, loss)

        optimizer.independent_gradient_partition_epilogue()

    def scale(self, loss):
        r"""Apply loss scaler for manual backward pass.

        Use this method when calling loss.backward() directly instead of engine.backward().
        This applies the appropriate loss scaler for mixed precision training, allowing you
        to manually control the backward pass while still benefiting from DeepSpeed's
        gradient scaling functionality.

        Example::

            output = engine(input)
            loss = criterion(output, target)
            scaled_loss = engine.scale(loss)
            scaled_loss.backward()  # Manual backward call
            engine.step()

        Arguments:
            loss: Scalar loss tensor to be scaled

        Returns:
            Scaled loss tensor ready for .backward() call

        Raises:
            RuntimeError: If AMP (NVIDIA Apex) is enabled. AMP requires using engine.backward()
                         directly as it uses a context manager that cannot be separated from
                         the backward call.
            AssertionError: If loss is not a scalar tensor with grad_fn, or if no optimizer
                           is configured.
        """
        assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \
            "must provide optimizer during init in order to use scale"
        assert maybe_loss_for_backward(loss), \
            "loss must be a scalar tensor with grad_fn. For non-scalar tensors, use tensor.backward(grad)"

        # AMP (NVIDIA Apex) uses a context manager that wraps both scaling and backward,
        # so it cannot be used with manual backward calls
        if self.amp_enabled():
            raise RuntimeError("engine.scale() is not compatible with AMP (NVIDIA Apex). "
                               "When using AMP, you must call engine.backward(loss) instead of manual backward.")

        # Apply loss scaler based on optimizer type
        scaled_loss = loss
        if isinstance(self.optimizer, ZeROOptimizer):
            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)

        # Mark that scale() was called for validation in backward hook
        self._manual_backward_expected = True

        return scaled_loss

    @instrument_w_nvtx
    def backward(self, loss, retain_graph=False, scale_wrt_gas=True):