r"""Execute the weight update step after forward and backward propagation on effective_train_batch.
(self, lr_kwargs=None)
| 2021 | self.global_samples += self.train_batch_size() |
| 2022 | |
| 2023 | def step(self, lr_kwargs=None): |
| 2024 | r"""Execute the weight update step after forward and backward propagation |
| 2025 | on effective_train_batch. |
| 2026 | """ |
| 2027 | see_memory_usage("Engine before step", force=self.memory_breakdown()) |
| 2028 | |
| 2029 | # Check early because self.global_steps is incremented at some point here. |
| 2030 | # TODO: Delay self.global_steps increment until very end of this function. |
| 2031 | flops_profiler_active = self.flops_profiler_enabled( |
| 2032 | ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0 |
| 2033 | |
| 2034 | self._start_timers(self.engine_timers.step_timers) |
| 2035 | |
| 2036 | assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \ |
| 2037 | "must provide optimizer during init in order to use step" |
| 2038 | |
| 2039 | report_progress = False |
| 2040 | |
| 2041 | self._step_applied = False # assume False, will flip to True |
| 2042 | |
| 2043 | # Update the model when we reach gradient accumulation boundaries |
| 2044 | if self.is_gradient_accumulation_boundary(): |
| 2045 | self.gas_boundary_ctr += 1 |
| 2046 | |
| 2047 | if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0) |
| 2048 | and self.quantizer.any_precision_switch()): |
| 2049 | log_dist(f"computing eigenvalue...", ranks=[0]) |
| 2050 | self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device, |
| 2051 | self.optimizer.cur_scale) |
| 2052 | |
| 2053 | if self.progressive_layer_drop: |
| 2054 | self.progressive_layer_drop.update_state(self.global_steps) |
| 2055 | |
| 2056 | if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() |
| 2057 | and self.quantizer.any_precision_switch()): |
| 2058 | self._take_model_step(lr_kwargs, self.block_eigenvalue) |
| 2059 | else: |
| 2060 | self._take_model_step(lr_kwargs) |
| 2061 | |
| 2062 | report_progress = self.global_rank == 0 if self.global_rank else True |
| 2063 | |
| 2064 | self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress) |
| 2065 | |
| 2066 | self._stop_timers(self.engine_timers.step_timers) |
| 2067 | |
| 2068 | # Log learning rate |
| 2069 | if self.monitor.enabled: |
| 2070 | if self.is_gradient_accumulation_boundary(): |
| 2071 | if self.global_rank == 0: |
| 2072 | self.summary_events = [(f"Train/Samples/lr", self.get_lr()[0], self.global_samples)] |
| 2073 | |
| 2074 | if self.fp16_enabled() and hasattr(self.optimizer, "cur_scale"): |
| 2075 | self.summary_events.append(( |
| 2076 | f"Train/Samples/loss_scale", |
| 2077 | self.optimizer.cur_scale, |
| 2078 | self.global_samples, |
| 2079 | )) |
| 2080 |
no test coverage detected