Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| from __future__ import annotations | |
| import time | |
| from typing import TYPE_CHECKING, Callable | |
| import torch | |
| import wandb | |
| from lipforcing.callbacks.callback import Callback | |
| from lipforcing.utils.distributed import is_rank0 | |
| import lipforcing.utils.logging_utils as logger | |
| if TYPE_CHECKING: | |
| from lipforcing.methods import FastGenModel | |
| class TrainProfilerCallback(Callback): | |
| """Callback for profiling training speed and detailed timing breakdowns. | |
| Tracks: | |
| - iter_time: seconds per iteration (wall clock time) | |
| - data_load_time: time spent loading data | |
| - avg_forward_time: average forward pass time across accumulation steps | |
| - backward_time: time spent in backward pass | |
| - optim_step_time: time spent in optimizer step | |
| """ | |
| def __init__(self, every_n: int = 100, detailed: bool = True): | |
| """Initialize the profiler callback. | |
| Args: | |
| every_n: Log metrics every N iterations | |
| detailed: If True, log detailed timing breakdown. If False, only log iter_time. | |
| """ | |
| # For iter_time tracking | |
| self.last_log_time = None | |
| # For detailed profiling | |
| self.detailed = detailed | |
| self.train_step_begin_time = None | |
| self.accum_begin_times = None | |
| self.backward_begin_times = None | |
| self.optimizer_step_begin = None | |
| self.step_end_time = None | |
| self.every_n = every_n | |
| def on_train_begin(self, model: FastGenModel, iteration: int = 0) -> None: | |
| if hasattr(self, "config"): | |
| # overwritten by logging_iter if self.config exists | |
| self.every_n = self.config.trainer.logging_iter | |
| logger.info(f"every_n to profile trainer: {self.every_n}") | |
| def on_training_step_begin( | |
| self, | |
| model: FastGenModel, | |
| iteration: int = 0, | |
| ): | |
| if self.detailed: | |
| self.train_step_begin_time = time.perf_counter() | |
| self.accum_begin_times = [] | |
| self.backward_begin_times = [] | |
| def on_training_accum_step_begin( | |
| self, model: FastGenModel, data_batch: dict[str, torch.Tensor], iteration: int = 0, accum_iter: int = 0 | |
| ): | |
| if self.detailed: | |
| self.accum_begin_times.append(time.perf_counter()) | |
| def on_backward_begin( | |
| self, | |
| model: FastGenModel, | |
| data_batch: dict[str, torch.Tensor], | |
| output_batch: dict[str, torch.Tensor | Callable], | |
| loss_dict: dict[str, torch.Tensor], | |
| iteration: int = 0, | |
| accum_iter: int = 0, | |
| ): | |
| if self.detailed: | |
| self.backward_begin_times.append(time.perf_counter()) | |
| def on_optimizer_step_begin(self, model: FastGenModel, iteration: int = 0): | |
| if self.detailed: | |
| self.optimizer_step_begin = time.perf_counter() | |
| def on_training_step_end( | |
| self, | |
| model: FastGenModel, | |
| data_batch: dict[str, torch.Tensor], | |
| output_batch: dict[str, torch.Tensor | Callable], | |
| loss_dict: dict[str, torch.Tensor], | |
| iteration: int = 0, | |
| ) -> None: | |
| del data_batch, output_batch, loss_dict | |
| if self.detailed: | |
| self.step_end_time = time.perf_counter() | |
| if hasattr(self, "config"): | |
| # only wandb log when config exists | |
| if iteration % self.every_n == 0 and is_rank0(): | |
| metrics = {} | |
| # Calculate iter_time (wall clock time per iteration) | |
| cur_time = time.time() | |
| if self.last_log_time is not None: | |
| iter_time = (cur_time - self.last_log_time) / self.every_n | |
| logger.info(f"{iteration} : avg iteration time {iter_time:.2f} seconds") | |
| metrics["profiler/avg_iteration_time"] = iter_time | |
| self.last_log_time = cur_time | |
| # Calculate detailed timing breakdown | |
| if self.detailed and self.accum_begin_times and self.backward_begin_times: | |
| data_load_time = self.accum_begin_times[0] - self.train_step_begin_time | |
| forward_time = sum( | |
| [b - a for (b, a) in zip(self.backward_begin_times, self.accum_begin_times)] | |
| ) / len(self.accum_begin_times) | |
| backward_time = self.optimizer_step_begin - self.backward_begin_times[-1] | |
| optim_step_time = self.step_end_time - self.optimizer_step_begin | |
| logger.info(f"{iteration} : data loading time {data_load_time:.2f}") | |
| logger.info(f"{iteration} : avg forward pass time {forward_time:.2f}") | |
| logger.info(f"{iteration} : backward pass time {backward_time:.2f}") | |
| logger.info(f"{iteration} : optimizer step time {optim_step_time:.2f}") | |
| metrics.update( | |
| { | |
| "profiler/data_loading_time": data_load_time, | |
| "profiler/avg_forward_pass_time": forward_time, | |
| "profiler/backward_pass_time": backward_time, | |
| "profiler/optimizer_step_time": optim_step_time, | |
| } | |
| ) | |
| if wandb.run and metrics: | |
| wandb.log(metrics, step=iteration) | |