# SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import time from collections.abc import Iterator from dataclasses import dataclass from typing import Any, TYPE_CHECKING import torch from tqdm.auto import tqdm from fastvideo.distributed import get_sp_group, get_world_group from fastvideo.train.callbacks.callback import CallbackDict from fastvideo.train.methods.base import LogScalar, TrainingMethod from fastvideo.train.utils.tracking import build_tracker if TYPE_CHECKING: from fastvideo.train.utils.training_config import ( TrainingConfig, ) def _coerce_log_scalar( value: Any, *, where: str, ) -> float | torch.Tensor: """Coerce *value* to a loggable scalar. GPU tensors stay on device so we avoid a ``cudaDeviceSynchronize`` per accumulation step. The caller must materialize them to float when logging. """ if isinstance(value, torch.Tensor): if value.numel() != 1: raise ValueError(f"Expected scalar tensor at {where}, " f"got shape={tuple(value.shape)}") return value.detach() if isinstance(value, float | int): return float(value) raise TypeError(f"Expected a scalar (float/int/Tensor) at " f"{where}, got {type(value).__name__}") @dataclass(slots=True) class TrainLoopState: step: int accum_iter: int class Trainer: def __init__( self, training_config: TrainingConfig, *, config: dict[str, Any] | None = None, callback_configs: dict[str, dict[str, Any]] | None = None, ) -> None: self.training_config = training_config self.world_group = get_world_group() self.sp_group = get_sp_group() self.global_rank = self.world_group.rank self.local_rank = self.world_group.local_rank self.tracker = build_tracker( training_config.tracker, training_config.checkpoint, config=config, ) self.callbacks = CallbackDict( callback_configs or {}, training_config, ) def _iter_dataloader(self, dataloader: Any) -> Iterator[dict[str, Any]]: data_iter = iter(dataloader) while True: batch = next(data_iter, None) if batch is None: data_iter = iter(dataloader) batch = next(data_iter) yield batch def _run_method_validation( self, method: TrainingMethod, iteration: int, ) -> None: hook = getattr(method, "on_validation_begin", None) if hook is None: return validation_metrics: dict[str, LogScalar] = hook(iteration) validation_metrics = { k: float(_coerce_log_scalar(v, where=(f"method.on_validation_begin().metrics[{k!r}]"))) for k, v in validation_metrics.items() } if self.global_rank == 0 and validation_metrics: self.tracker.log(validation_metrics, iteration) def run( self, method: TrainingMethod, *, dataloader: Any, max_steps: int, start_step: int = 0, checkpoint_manager: Any | None = None, ) -> None: tc = self.training_config grad_accum = max( 1, int(tc.loop.gradient_accumulation_steps or 1), ) method.set_tracker(self.tracker) method.on_train_start() self.callbacks.on_train_start( method, iteration=start_step, ) resume_from_checkpoint = (tc.checkpoint.resume_from_checkpoint or "") if checkpoint_manager is not None: if resume_from_checkpoint: method.seed_optimizer_state_for_resume() resumed_step = (checkpoint_manager.maybe_resume(resume_from_checkpoint=(resume_from_checkpoint))) if resumed_step is not None: start_step = int(resumed_step) self.callbacks.on_validation_begin( method, iteration=start_step, ) self._run_method_validation(method, start_step) method.optimizers_zero_grad(start_step) data_stream = self._iter_dataloader(dataloader) # Restore the RNG snapshot LAST — after dcp.load, # after iter(dataloader), after everything that may # have advanced the RNG as a side-effect. if (checkpoint_manager is not None and resume_from_checkpoint): checkpoint_manager.load_rng_snapshot(resume_from_checkpoint, ) progress = tqdm( range(start_step + 1, max_steps + 1), initial=start_step, desc="Steps", disable=self.local_rank > 0, ) # Allow method-specific optimization flow (e.g. DiffusionNFT). method_manages_optimization = bool(method.manages_optimization()) for step in progress: t0 = time.perf_counter() # Accumulate on GPU during grad-accum; materialise # to CPU once per step right before logging. loss_sums: dict[str, float | torch.Tensor] = {} metric_sums: dict[str, float | torch.Tensor] = {} if method_manages_optimization: loss_map, outputs, step_metrics = method.managed_train_step( data_stream, step, ) for k, v in loss_map.items(): if isinstance(v, torch.Tensor): loss_sums[k] = v.detach() for k, v in step_metrics.items(): if k in loss_sums: raise ValueError(f"Metric key {k!r} collides " "with loss key. Use a " "different name (e.g. prefix " "with 'train/').") metric_sums[k] = _coerce_log_scalar( v, where=("method.managed_train_step()" f".metrics[{k!r}]"), ) else: for accum_iter in range(grad_accum): batch = next(data_stream) loss_map, outputs, step_metrics = (method.single_train_step( batch, step, )) method.backward( loss_map, outputs, grad_accum_rounds=grad_accum, ) for k, v in loss_map.items(): if isinstance(v, torch.Tensor): prev = loss_sums.get(k, 0.0) loss_sums[k] = prev + v.detach() for k, v in step_metrics.items(): if k in loss_sums: raise ValueError(f"Metric key {k!r} collides " "with loss key. Use a " "different name (e.g. prefix " "with 'train/').") prev = metric_sums.get(k, 0.0) metric_sums[k] = (prev + _coerce_log_scalar( v, where=("method.single_train_step()" f".metrics[{k!r}]"), )) if not method_manages_optimization: self.callbacks.on_before_optimizer_step( method, iteration=step, ) method.optimizers_schedulers_step(step) method.optimizers_zero_grad(step) # Single CPU sync point: materialise GPU tensors # to float right before logging. divisor = 1 if method_manages_optimization else grad_accum metrics = {k: float(v) / divisor for k, v in loss_sums.items()} metrics.update({k: float(v) / divisor for k, v in metric_sums.items()}) metrics["step_time_sec"] = (time.perf_counter() - t0) metrics["vsa_sparsity"] = float(tc.vsa_sparsity) if self.global_rank == 0 and metrics: self.tracker.log(metrics, step) self.callbacks.on_training_step_end( method, metrics, iteration=step, ) if checkpoint_manager is not None: checkpoint_manager.maybe_save(step) self.callbacks.on_validation_begin( method, iteration=step, ) self._run_method_validation(method, step) self.callbacks.on_validation_end( method, iteration=step, ) self.callbacks.on_train_end( method, iteration=max_steps, ) if checkpoint_manager is not None: checkpoint_manager.save_final(max_steps) self.tracker.finish()