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)
| 3011 | optimizer.independent_gradient_partition_epilogue() |
| 3012 | |
| 3013 | def scale(self, loss): |
| 3014 | r"""Apply loss scaler for manual backward pass. |
| 3015 | |
| 3016 | Use this method when calling loss.backward() directly instead of engine.backward(). |
| 3017 | This applies the appropriate loss scaler for mixed precision training, allowing you |
| 3018 | to manually control the backward pass while still benefiting from DeepSpeed's |
| 3019 | gradient scaling functionality. |
| 3020 | |
| 3021 | Example:: |
| 3022 | |
| 3023 | output = engine(input) |
| 3024 | loss = criterion(output, target) |
| 3025 | scaled_loss = engine.scale(loss) |
| 3026 | scaled_loss.backward() # Manual backward call |
| 3027 | engine.step() |
| 3028 | |
| 3029 | Arguments: |
| 3030 | loss: Scalar loss tensor to be scaled |
| 3031 | |
| 3032 | Returns: |
| 3033 | Scaled loss tensor ready for .backward() call |
| 3034 | |
| 3035 | Raises: |
| 3036 | RuntimeError: If AMP (NVIDIA Apex) is enabled. AMP requires using engine.backward() |
| 3037 | directly as it uses a context manager that cannot be separated from |
| 3038 | the backward call. |
| 3039 | AssertionError: If loss is not a scalar tensor with grad_fn, or if no optimizer |
| 3040 | is configured. |
| 3041 | """ |
| 3042 | assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ |
| 3043 | "must provide optimizer during init in order to use scale" |
| 3044 | assert maybe_loss_for_backward(loss), \ |
| 3045 | "loss must be a scalar tensor with grad_fn. For non-scalar tensors, use tensor.backward(grad)" |
| 3046 | |
| 3047 | # AMP (NVIDIA Apex) uses a context manager that wraps both scaling and backward, |
| 3048 | # so it cannot be used with manual backward calls |
| 3049 | if self.amp_enabled(): |
| 3050 | raise RuntimeError("engine.scale() is not compatible with AMP (NVIDIA Apex). " |
| 3051 | "When using AMP, you must call engine.backward(loss) instead of manual backward.") |
| 3052 | |
| 3053 | # Apply loss scaler based on optimizer type |
| 3054 | scaled_loss = loss |
| 3055 | if isinstance(self.optimizer, ZeROOptimizer): |
| 3056 | scaled_loss = self.optimizer.scale_if_loss(scaled_loss) |
| 3057 | elif self.torch_autocast_z0_gradscaler: |
| 3058 | scaled_loss = self.torch_autocast_z0_gradscaler.scale(scaled_loss) |
| 3059 | |
| 3060 | # Mark that scale() was called for validation in backward hook |
| 3061 | self._manual_backward_expected = True |
| 3062 | |
| 3063 | return scaled_loss |
| 3064 | |
| 3065 | @instrument_w_nvtx |
| 3066 | def backward(self, loss, retain_graph=False, scale_wrt_gas=True): |
no test coverage detected