r"""Execute the weight update step after forward and backward propagation on effective_train_batch.
(self, lr_kwargs=None)
| 3239 | self.global_samples += self.train_batch_size() |
| 3240 | |
| 3241 | def step(self, lr_kwargs=None): |
| 3242 | r"""Execute the weight update step after forward and backward propagation |
| 3243 | on effective_train_batch. |
| 3244 | """ |
| 3245 | assert not self.inside_no_sync_ctxt, \ |
| 3246 | "It is illegal to call Engine.step() inside no_sync context manager" |
| 3247 | |
| 3248 | see_memory_usage("Engine before step", force=self.memory_breakdown()) |
| 3249 | |
| 3250 | # Check early because self.global_steps is incremented at some point here. |
| 3251 | # TODO: Delay self.global_steps increment until very end of this function. |
| 3252 | flops_profiler_active = self.flops_profiler_enabled( |
| 3253 | ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 |
| 3254 | |
| 3255 | self._start_timers(self.engine_timers.step_timers) |
| 3256 | |
| 3257 | assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ |
| 3258 | "must provide optimizer during init in order to use step" |
| 3259 | |
| 3260 | report_progress = False |
| 3261 | |
| 3262 | self._step_applied = False # assume False, will flip to True |
| 3263 | |
| 3264 | if self.zenflow: |
| 3265 | self.optimizer._sync_selective_optimizer_lr() |
| 3266 | if self.auto_update: |
| 3267 | self.update_interval += 1 |
| 3268 | |
| 3269 | # Update the model when we reach gradient accumulation boundaries |
| 3270 | if self.is_gradient_accumulation_boundary(): |
| 3271 | self.gas_boundary_ctr += 1 |
| 3272 | |
| 3273 | if self.checkpoint_engine.is_decoupled(): |
| 3274 | self._commit_decoupled_checkpoint() |
| 3275 | |
| 3276 | if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) |
| 3277 | and self.quantizer.any_precision_switch()): |
| 3278 | log_dist("computing eigenvalue...", ranks=[0]) |
| 3279 | loss_scale = self._get_optimizer_loss_scale() or 1.0 |
| 3280 | self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, loss_scale) |
| 3281 | |
| 3282 | if self.progressive_layer_drop: |
| 3283 | self.progressive_layer_drop.update_state(self.global_steps) |
| 3284 | |
| 3285 | if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() |
| 3286 | and self.quantizer.any_precision_switch()): |
| 3287 | self._take_model_step(lr_kwargs, self.block_eigenvalue) |
| 3288 | else: |
| 3289 | self._take_model_step(lr_kwargs) |
| 3290 | |
| 3291 | report_progress = self.global_rank == 0 if self.global_rank else True |
| 3292 | |
| 3293 | if self.zenflow: |
| 3294 | self._zenflow_step(lr_kwargs) |
| 3295 | |
| 3296 | self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress) |
| 3297 | |
| 3298 | self._stop_timers(self.engine_timers.step_timers) |
no test coverage detected