r"""Execute the weight update step after forward and backward propagation on effective_train_batch. In managed gradient accumulation (default), the optimizer update is applied only on the accumulation boundary tracked by the internal micro-step counter. With ``manage
(self, lr_kwargs=None)
| 3513 | self.global_samples += self.train_batch_size() |
| 3514 | |
| 3515 | def step(self, lr_kwargs=None): |
| 3516 | r"""Execute the weight update step after forward and backward propagation |
| 3517 | on effective_train_batch. |
| 3518 | |
| 3519 | In managed gradient accumulation (default), the optimizer update is applied only |
| 3520 | on the accumulation boundary tracked by the internal micro-step counter. With |
| 3521 | ``managed_gradient_accumulation=false``, every ``step()`` is the accumulation |
| 3522 | boundary: it finalizes the locally accumulated gradients and applies an update. |
| 3523 | """ |
| 3524 | assert not self.inside_no_sync_ctxt, \ |
| 3525 | "It is illegal to call Engine.step() inside no_sync context manager" |
| 3526 | |
| 3527 | see_memory_usage("Engine before step", force=self.memory_breakdown()) |
| 3528 | |
| 3529 | # Check early because self.global_steps is incremented at some point here. |
| 3530 | # TODO: Delay self.global_steps increment until very end of this function. |
| 3531 | flops_profiler_active = self.flops_profiler_enabled( |
| 3532 | ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 |
| 3533 | |
| 3534 | self._start_timers(self.engine_timers.step_timers) |
| 3535 | |
| 3536 | assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ |
| 3537 | "must provide optimizer during init in order to use step" |
| 3538 | |
| 3539 | report_progress = False |
| 3540 | |
| 3541 | self._step_applied = False # assume False, will flip to True |
| 3542 | |
| 3543 | # Unmanaged mode: step() is the accumulation boundary. |
| 3544 | self._running_engine_step = True |
| 3545 | |
| 3546 | # Unmanaged boundary: stage 2/3 already reduced/partitioned per backward so only finalize (incl. offload); stage 0/1/DDP reduce here. |
| 3547 | if not self.managed_gradient_accumulation(): |
| 3548 | if self.zero_optimization_partition_gradients(): |
| 3549 | self.optimizer.finalize_gradient_accumulation_boundary() |
| 3550 | elif self.enable_backward_allreduce and not self.inside_no_sync_ctxt: |
| 3551 | self.allreduce_gradients() |
| 3552 | |
| 3553 | if self.zenflow: |
| 3554 | self.optimizer._sync_selective_optimizer_lr() |
| 3555 | if self.auto_update: |
| 3556 | self.update_interval += 1 |
| 3557 | |
| 3558 | # Update the model when we reach gradient accumulation boundaries |
| 3559 | if self.is_gradient_accumulation_boundary(): |
| 3560 | self.gas_boundary_ctr += 1 |
| 3561 | |
| 3562 | if self.checkpoint_engine.is_decoupled(): |
| 3563 | self._commit_decoupled_checkpoint() |
| 3564 | |
| 3565 | if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) |
| 3566 | and self.quantizer.any_precision_switch()): |
| 3567 | log_dist("computing eigenvalue...", ranks=[0]) |
| 3568 | loss_scale = self._get_optimizer_loss_scale() or 1.0 |
| 3569 | self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, loss_scale) |
| 3570 | |
| 3571 | if self.progressive_layer_drop: |
| 3572 | self.progressive_layer_drop.update_state(self.global_steps) |
no test coverage detected