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)
| 3149 | |
| 3150 | @instrument_w_nvtx |
| 3151 | def backward(self, loss, retain_graph=False, scale_wrt_gas=True): |
| 3152 | r"""Execute backward pass on the loss |
| 3153 | Arguments: |
| 3154 | loss: Torch tensor on which to execute backward propagation |
| 3155 | retain_graph: bool, default: false |
| 3156 | forward on user defined choice of retain_graph |
| 3157 | scale_wrt_gas: bool, default: true |
| 3158 | whether to scale gradients and return value by gradient accumulation steps |
| 3159 | |
| 3160 | With ``managed_gradient_accumulation=false``, ``backward()`` only accumulates |
| 3161 | gradients locally (it does not trigger the accumulation-boundary reduction); the |
| 3162 | reduction and optimizer update happen in ``step()``. |
| 3163 | """ |
| 3164 | assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ |
| 3165 | "must provide optimizer during init in order to use backward" |
| 3166 | assert maybe_loss_for_backward( |
| 3167 | loss), "loss must be a scalar tensor. If you need to pass output gradients, backward() of output tensors" |
| 3168 | |
| 3169 | self._running_engine_backward = True |
| 3170 | # Store scale_wrt_gas so the hook can respect it |
| 3171 | self._scale_wrt_gas = scale_wrt_gas |
| 3172 | |
| 3173 | # Unmanaged mode: count this backward so step() can advance global_samples by the actual micro-batch count. |
| 3174 | if not self.managed_gradient_accumulation(): |
| 3175 | self._unmanaged_backward_count += 1 |
| 3176 | |
| 3177 | # Set flag to prevent hooks from firing (we'll manually call prologue/epilogue) |
| 3178 | backward_kwargs = {"retain_graph": retain_graph} |
| 3179 | if self.eigenvalue_enabled(): |
| 3180 | backward_kwargs["create_graph"] = True |
| 3181 | backward_kwargs["retain_graph"] = True |
| 3182 | |
| 3183 | # Used only for return value |
| 3184 | gas_scaled_loss = loss / self.gradient_accumulation_steps() if scale_wrt_gas else loss |
| 3185 | |
| 3186 | # TODO: handle these scaling with direct calls to loss.backward() |
| 3187 | if isinstance(self.optimizer, ZeROOptimizer): |
| 3188 | loss = self.optimizer.scale_if_loss(loss) |
| 3189 | elif self.torch_autocast_z0_gradscaler: |
| 3190 | loss = self.torch_autocast_z0_gradscaler.scale(loss) |
| 3191 | |
| 3192 | with compiled_autograd(self._is_compiled_autograd_enabled, self._compile_kwargs): |
| 3193 | if self.zero_optimization() or not self.amp_enabled(): |
| 3194 | loss.backward(**backward_kwargs) |
| 3195 | elif self.amp_enabled(): |
| 3196 | # AMP requires delaying unscale when inside gradient accumulation boundaries |
| 3197 | # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations |
| 3198 | delay_unscale = not self.is_gradient_accumulation_boundary() |
| 3199 | with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss: |
| 3200 | scaled_loss.backward(**backward_kwargs) |
| 3201 | |
| 3202 | # backward_epilogue is not called in a hook when self._support_torch_style_backward is False |
| 3203 | self._backward_epilogue() |
| 3204 | |
| 3205 | self._running_engine_backward = False |
| 3206 | |
| 3207 | return gas_scaled_loss |
| 3208 |
nothing calls this directly
no test coverage detected