r"""Execute the weight update step after forward and backward propagation on effective_train_batch.
(self, lr_kwargs=None)
| 3025 | self.global_samples += self.train_batch_size() |
| 3026 | |
| 3027 | def step(self, lr_kwargs=None): |
| 3028 | r"""Execute the weight update step after forward and backward propagation |
| 3029 | on effective_train_batch. |
| 3030 | """ |
| 3031 | assert not self.inside_no_sync_ctxt, \ |
| 3032 | "It is illegal to call Engine.step() inside no_sync context manager" |
| 3033 | |
| 3034 | see_memory_usage("Engine before step", force=self.memory_breakdown()) |
| 3035 | |
| 3036 | # Check early because self.global_steps is incremented at some point here. |
| 3037 | # TODO: Delay self.global_steps increment until very end of this function. |
| 3038 | flops_profiler_active = self.flops_profiler_enabled( |
| 3039 | ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 |
| 3040 | |
| 3041 | self._start_timers(self.engine_timers.step_timers) |
| 3042 | |
| 3043 | assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ |
| 3044 | "must provide optimizer during init in order to use step" |
| 3045 | |
| 3046 | report_progress = False |
| 3047 | |
| 3048 | self._step_applied = False # assume False, will flip to True |
| 3049 | |
| 3050 | if self.zenflow: |
| 3051 | self.optimizer._sync_selective_optimizer_lr() |
| 3052 | if self.auto_update: |
| 3053 | self.update_interval += 1 |
| 3054 | |
| 3055 | # Update the model when we reach gradient accumulation boundaries |
| 3056 | if self.is_gradient_accumulation_boundary(): |
| 3057 | self.gas_boundary_ctr += 1 |
| 3058 | |
| 3059 | if self.checkpoint_engine.is_decoupled(): |
| 3060 | self._commit_decoupled_checkpoint() |
| 3061 | |
| 3062 | if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) |
| 3063 | and self.quantizer.any_precision_switch()): |
| 3064 | log_dist("computing eigenvalue...", ranks=[0]) |
| 3065 | loss_scale = self._get_optimizer_loss_scale() or 1.0 |
| 3066 | self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, loss_scale) |
| 3067 | |
| 3068 | if self.progressive_layer_drop: |
| 3069 | self.progressive_layer_drop.update_state(self.global_steps) |
| 3070 | |
| 3071 | if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() |
| 3072 | and self.quantizer.any_precision_switch()): |
| 3073 | self._take_model_step(lr_kwargs, self.block_eigenvalue) |
| 3074 | else: |
| 3075 | self._take_model_step(lr_kwargs) |
| 3076 | |
| 3077 | report_progress = self.global_rank == 0 if self.global_rank else True |
| 3078 | |
| 3079 | if self.zenflow: |
| 3080 | self._zenflow_step(lr_kwargs) |
| 3081 | |
| 3082 | self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress) |
| 3083 | |
| 3084 | self._stop_timers(self.engine_timers.step_timers) |
no test coverage detected