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)
| 3348 | self.global_samples += self.train_batch_size() |
| 3349 | |
| 3350 | def step(self, lr_kwargs=None): |
| 3351 | r"""Execute the weight update step after forward and backward propagation |
| 3352 | on effective_train_batch. |
| 3353 | |
| 3354 | In managed gradient accumulation (default), the optimizer update is applied only |
| 3355 | on the accumulation boundary tracked by the internal micro-step counter. With |
| 3356 | ``managed_gradient_accumulation=false``, every ``step()`` is the accumulation |
| 3357 | boundary: it finalizes the locally accumulated gradients and applies an update. |
| 3358 | """ |
| 3359 | assert not self.inside_no_sync_ctxt, \ |
| 3360 | "It is illegal to call Engine.step() inside no_sync context manager" |
| 3361 | |
| 3362 | see_memory_usage("Engine before step", force=self.memory_breakdown()) |
| 3363 | |
| 3364 | # Check early because self.global_steps is incremented at some point here. |
| 3365 | # TODO: Delay self.global_steps increment until very end of this function. |
| 3366 | flops_profiler_active = self.flops_profiler_enabled( |
| 3367 | ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 |
| 3368 | |
| 3369 | self._start_timers(self.engine_timers.step_timers) |
| 3370 | |
| 3371 | assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ |
| 3372 | "must provide optimizer during init in order to use step" |
| 3373 | |
| 3374 | report_progress = False |
| 3375 | |
| 3376 | self._step_applied = False # assume False, will flip to True |
| 3377 | |
| 3378 | # Unmanaged mode: step() is the accumulation boundary. |
| 3379 | self._running_engine_step = True |
| 3380 | |
| 3381 | # Unmanaged boundary: stage 2/3 already reduced/partitioned per backward so only finalize; stage 0/1/DDP reduce here. |
| 3382 | if not self.managed_gradient_accumulation(): |
| 3383 | if self.zero_optimization_partition_gradients(): |
| 3384 | self.optimizer.finalize_gradient_accumulation_boundary() |
| 3385 | elif self.enable_backward_allreduce and not self.inside_no_sync_ctxt: |
| 3386 | self.allreduce_gradients() |
| 3387 | |
| 3388 | if self.zenflow: |
| 3389 | self.optimizer._sync_selective_optimizer_lr() |
| 3390 | if self.auto_update: |
| 3391 | self.update_interval += 1 |
| 3392 | |
| 3393 | # Update the model when we reach gradient accumulation boundaries |
| 3394 | if self.is_gradient_accumulation_boundary(): |
| 3395 | self.gas_boundary_ctr += 1 |
| 3396 | |
| 3397 | if self.checkpoint_engine.is_decoupled(): |
| 3398 | self._commit_decoupled_checkpoint() |
| 3399 | |
| 3400 | if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) |
| 3401 | and self.quantizer.any_precision_switch()): |
| 3402 | log_dist("computing eigenvalue...", ranks=[0]) |
| 3403 | loss_scale = self._get_optimizer_loss_scale() or 1.0 |
| 3404 | self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, loss_scale) |
| 3405 | |
| 3406 | if self.progressive_layer_drop: |
| 3407 | self.progressive_layer_drop.update_state(self.global_steps) |
no test coverage detected