diff --git a/docs/training/train_infra.md b/docs/training/train_infra.md index 766d261d2..a37689636 100644 --- a/docs/training/train_infra.md +++ b/docs/training/train_infra.md @@ -151,6 +151,10 @@ training: project_name: my_project run_name: my_run + performance: + enabled: true + peak_tflops_per_gpu: null # optional dense BF16 peak for one GPU + model: weighting_scheme: uniform # uniform, logit_normal, mode precondition_outputs: false @@ -177,6 +181,81 @@ The repeat count duplicates that dataset's parquet file list before shuffling/sa See [Training Trackers](trackers.md) to configure Weights & Biases or SwanLab, including SwanLab installation and authentication. +### Training performance metrics + +The modular trainer logs low-overhead performance metrics on every optimizer +step. The Transformer boundary is instrumented separately for each model role, +so the accounting includes all forwards performed by the selected method: + +- bidirectional SFT records its dense or VSA student forward; +- causal SFT records block-causal attention geometry; +- streaming/self-forcing records every rollout chunk and KV-cache update; +- DMD records student, critic, and teacher forwards independently, including + repeated rollouts and the conditional/unconditional teacher passes. + +The common metrics are: + +| Metric | Meaning | +|--------|---------| +| `step_time_sec` | Training wall time for one optimizer step, excluding checkpoint and validation callbacks | +| `perf/steps_per_sec` | Reciprocal of `step_time_sec` | +| `perf/samples_per_sec` | Configured global samples processed per second, including gradient accumulation and data-parallel replicas | +| `perf/model_forward_calls` | Actual Transformer invocations during the optimizer step | +| `perf/causal_chunks` | Block-causal chunks represented by those invocations | +| `perf/query_latent_frames_per_sec` | Latent query frames processed across all roles and rollouts | +| `perf/query_tokens_per_sec` | Patch tokens processed across all roles and rollouts | +| `perf/attention_density` | Effective self-attention pairs divided by dense attention pairs | +| `perf/estimated_tflops_per_gpu` | Estimated useful model FLOP/s per GPU | +| `perf/estimated_mfu` | Estimated useful FLOP/s divided by the dense BF16 peak (ratio, not percent) | +| `perf/peak_tflops_per_gpu` | Dense BF16 peak used for the MFU ratio (configured or inferred) | +| `perf/role//*` | Forward count, grad-carrying forward count, causal chunks, and estimated TFLOP/s for one role | + +For DMD2, combine these metrics with the existing `update_student` metric to +separate the expensive generator-update steps from critic-only steps. + +`perf/samples_per_sec` uses the configured `train_batch_size`, gradient +accumulation, and data-parallel replica count. Methods that manage their own +optimization (for example DiffusionNFT, whose outer step consumes +`num_batches_per_epoch` dataloader batches) report the throughput of a single +batch in the denominator, so use the method's own sample counter +(`nft/num_sampled`) for those runs. The FLOP-derived metrics still aggregate +every forward the method performs. + +MFU is an analytic Transformer-core estimate. Each no-grad forward contributes +`1F`; a forward whose result carries autograd contributes `3F` (forward plus an +approximately `2F` backward). Activation-checkpoint recomputation is excluded, +as expected for model FLOPs utilization. Wan VSA uses the kernel's clamped +tile-top-k density and includes its gate projection and pooled-attention +overhead. VMOBA runs are modeled as dense self-attention: the MoBA top-k +selection is not estimated, so `perf/attention_density` stays at `1.0` and the +attention-dependent FLOPs are an upper bound. Causal Wan uses the configured +chunk size, local window, and actual streaming cache position; when +`local_attn_size` is unset, the modeled window is the transformer's 21-frame +compatibility cap, and `pipeline.dit_config.sliding_window_num_frames` only +sizes the streaming KV cache. MatrixGame2's 15-frame compatibility window is +not modeled, so its attention density is an upper bound. Embedding, +normalization, optimizer, communication, MatrixGame action modules, and other +non-core work are not included, so MFU is an estimate rather than a +hardware-profiler measurement. + +`perf/estimated_tflops_per_gpu` and `perf/estimated_mfu` aggregate every role +(student, teacher, critic, EMA) that ran during the step, so they are not +directly comparable to a single-model MFU number. Replica scaling assumes +`world_size = data_parallel x sp_size`; `tp_size > 1` is not modeled and would +inflate the per-GPU estimates. + +Forward counts and wall-clock throughput work for every modular model. The +FLOP-derived metrics currently require a Wan-style architecture exposing +`hidden_size`, `ffn_dim`, `num_layers`, and a three-axis `patch_size`; they are +omitted for other architectures instead of reporting a misleading estimate. + +Known NVIDIA accelerators use an inferred dense BF16 peak. Set +`training.performance.peak_tflops_per_gpu` explicitly for a different board +form factor or clock. If the device is unknown and no peak is configured, the +trainer still logs throughput, call counts, attention density, and estimated +TFLOP/s, but omits `perf/estimated_mfu`. Set `enabled: false` to disable the +forward hooks and all `perf/*` metrics. + ### `callbacks` — Pluggable hooks Callbacks run at specific points in the training loop (before/after optimizer diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index 51dc03e66..af64fd45a 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -130,6 +130,13 @@ training: project_name: my_project # default: "fastvideo" run_name: my_run # default: "" + # --- training.performance [TYPED] -> PerformanceConfig --- + performance: + enabled: true # default: true + # Optional dense BF16 peak for one GPU. Known NVIDIA device names are + # inferred automatically; set this for non-standard clocks/form factors. + peak_tflops_per_gpu: null # e.g. 989.5 for H100 SXM + # --- training.vsa [TYPED] -> VSAConfig --- vsa: sparsity: 0.0 # default: 0.0 (0.0 = disabled) diff --git a/fastvideo/tests/train/trainer/test_performance_wiring.py b/fastvideo/tests/train/trainer/test_performance_wiring.py new file mode 100644 index 000000000..5fd3ff3b7 --- /dev/null +++ b/fastvideo/tests/train/trainer/test_performance_wiring.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for Trainer performance-metric wiring.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import torch + +from fastvideo.train.trainer import Trainer +from fastvideo.train.utils.performance import TrainingPerformanceMonitor +from fastvideo.train.utils.training_config import TrainingConfig + + +class _FakeTransformer(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace( + arch_config=SimpleNamespace( + hidden_size=8, + ffn_dim=16, + num_layers=2, + patch_size=(1, 2, 2), + )) + self.weight = torch.nn.Parameter(torch.ones(())) + + def forward( + self, + *, + hidden_states: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + del kwargs + return hidden_states * self.weight + + +class _FakeRoleModel: + + def __init__(self) -> None: + self.transformer = _FakeTransformer() + + +class _DummyTracker: + + def __init__(self) -> None: + self.logs: list[tuple[dict[str, float], int]] = [] + + def log(self, metrics: dict[str, float], step: int) -> None: + self.logs.append((metrics, step)) + + def finish(self) -> None: + pass + + +class _PerformanceMethod: + + def __init__(self) -> None: + self._role_models = {"student": _FakeRoleModel()} + self.weight = torch.nn.Parameter(torch.tensor(1.0)) + self.tracker = None + self.pre_loop_forwards = 0 + + def set_tracker(self, tracker: Any) -> None: + self.tracker = tracker + + def on_train_start(self) -> None: + # Stands in for teacher-cache generation or step-0 validation. + self._forward() + self.pre_loop_forwards += 1 + + def manages_optimization(self) -> bool: + return False + + def _forward(self) -> torch.Tensor: + hidden = torch.ones(1, 3, 2, 4, 4) + return self._role_models["student"].transformer(hidden_states=hidden) + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, float]]: + del batch, iteration + output = self._forward() + return {"total_loss": output.sum() * 0.0 + 1.0}, {}, {} + + def backward( + self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int, + ) -> None: + del outputs + (loss_map["total_loss"] / grad_accum_rounds).backward() + + def optimizers_schedulers_step(self, iteration: int) -> None: + del iteration + + def optimizers_zero_grad(self, iteration: int) -> None: + del iteration + self.weight.grad = None + + +def _run_trainer( + monkeypatch: Any, + *, + performance_enabled: bool, +) -> tuple[_PerformanceMethod, _DummyTracker]: + tracker = _DummyTracker() + group = SimpleNamespace(rank=0, local_rank=0, rank_in_group=0, world_size=1) + + monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: group) + monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: group) + monkeypatch.setattr( + "fastvideo.train.trainer.build_tracker", + lambda *args, **kwargs: tracker, + ) + + cfg = TrainingConfig() + cfg.tracker.project_name = "" + cfg.loop.gradient_accumulation_steps = 1 + cfg.performance.enabled = performance_enabled + trainer = Trainer(cfg) + method = _PerformanceMethod() + + trainer.run( + method, + dataloader=[{ + "sample": "x" + }], + max_steps=1, + ) + return method, tracker + + +def test_pre_loop_forwards_are_not_counted(monkeypatch) -> None: + method, tracker = _run_trainer(monkeypatch, performance_enabled=True) + + assert method.pre_loop_forwards == 1 + metrics = tracker.logs[-1][0] + # Only the forward inside the training step is measured; the forward run + # during on_train_start() was dropped by the pre-loop reset. + assert metrics["perf/model_forward_calls"] == 1.0 + assert metrics["perf/role/student/forward_calls"] == 1.0 + + +def test_pre_loop_forwards_are_dropped_before_the_first_step(monkeypatch) -> None: + observed: list[int] = [] + original_reset = TrainingPerformanceMonitor.reset + + def recording_reset(self: TrainingPerformanceMonitor) -> None: + observed.append(len(self._work)) + original_reset(self) + + monkeypatch.setattr(TrainingPerformanceMonitor, "reset", recording_reset) + method, _ = _run_trainer(monkeypatch, performance_enabled=True) + + assert method.pre_loop_forwards == 1 + # The monitor must already be empty when the first training step starts, so + # a long pre-loop phase (teacher cache generation, step-0 validation) + # cannot accumulate per-forward records. + assert observed[-1] == 0 + + +def test_disabled_performance_logs_no_perf_metrics(monkeypatch) -> None: + _, tracker = _run_trainer(monkeypatch, performance_enabled=False) + + metrics = tracker.logs[-1][0] + assert not any(key.startswith("perf/") for key in metrics) diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index 8471eea8c..a15c52027 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -81,6 +81,9 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None: assert t.tracker.project_name == "fastvideo" assert t.tracker.run_name == "" + assert t.performance.enabled is True + assert t.performance.peak_tflops_per_gpu is None + assert t.model.weighting_scheme == "uniform" assert t.model.precondition_outputs is False assert t.model.moba_config == {} @@ -134,6 +137,10 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: "project_name": "myproj", "run_name": "myrun", }, + "performance": { + "enabled": True, + "peak_tflops_per_gpu": 989.5, + }, "vsa": { "sparsity": 0.5 }, @@ -172,6 +179,9 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None: assert t.tracker.trackers == ["wandb"] assert t.tracker.project_name == "myproj" + assert t.performance.enabled is True + assert t.performance.peak_tflops_per_gpu == pytest.approx(989.5) + assert t.vsa_sparsity == pytest.approx(0.5) assert t.model.weighting_scheme == "logit_normal" assert t.model.precondition_outputs is True @@ -223,6 +233,13 @@ def test_missing_config_file_raises(tmp_path: Path) -> None: load_run_config(str(tmp_path / "does_not_exist.yaml")) +def test_non_positive_peak_tflops_raises(tmp_path: Path) -> None: + data = _minimal_yaml() + data["training"] = {"performance": {"peak_tflops_per_gpu": 0}} + with pytest.raises(ValueError, match="peak_tflops_per_gpu"): + load_run_config(_write_yaml(tmp_path, data)) + + # --------------------------------------------------------------------------- # Special parsing # --------------------------------------------------------------------------- diff --git a/fastvideo/tests/train/utils/test_performance.py b/fastvideo/tests/train/utils/test_performance.py new file mode 100644 index 000000000..0d9d08022 --- /dev/null +++ b/fastvideo/tests/train/utils/test_performance.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only tests for modular training performance accounting.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from fastvideo.forward_context import set_forward_context +from fastvideo.train.utils.performance import ( + TrainingPerformanceMonitor, + _blockwise_causal_frame_pairs, + _teacher_forcing_frame_pairs, + estimate_transformer_forward, + infer_peak_bf16_tflops, +) + + +def _arch() -> SimpleNamespace: + return SimpleNamespace( + hidden_size=8, + ffn_dim=16, + num_layers=2, + patch_size=(1, 2, 2), + ) + + +class _FakeTransformer(torch.nn.Module): + + def __init__(self, *, causal: bool = False) -> None: + super().__init__() + self.config = SimpleNamespace(arch_config=_arch()) + self.weight = torch.nn.Parameter(torch.ones(())) + if causal: + self.num_frame_per_block = 2 + self.local_attn_size = 3 + self.text_len = 4 + + def forward( + self, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + del encoder_hidden_states, kwargs + return hidden_states * self.weight + + +class _UnsupportedTransformer(torch.nn.Module): + + def forward(self, *, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + +def _kwargs(*, frames: int = 4) -> dict[str, torch.Tensor]: + return { + "hidden_states": torch.ones(2, 3, frames, 4, 4), + "encoder_hidden_states": torch.ones(2, 3, 8), + } + + +def test_dense_forward_flops_and_backward_factor() -> None: + module = _FakeTransformer() + output = torch.ones(2, 3, 4, 4, 4, requires_grad=True) + + work = estimate_transformer_forward( + module, + _kwargs(), + output, + role="student", + attention_metadata=None, + ) + + assert work is not None + # B=2, L=16, C=3, d=8, ffn=16, layers=2. + expected_per_layer = ( + 8 * 2 * 16 * 8**2 + + 4 * 2 * 16 * 8 * 16 + + 4 * 2 * 16 * 8**2 + + 4 * 2 * 3 * 8**2 + + 4 * 2 * 16 * 3 * 8 + + 4 * 2 * 16**2 * 8 + ) + assert work.forward_flops == expected_per_layer * 2 + assert work.useful_flops == work.forward_flops * 3 + assert work.query_frames == 8 + assert work.query_tokens == 32 + assert work.attention_pairs == 2 * 16**2 + + +def test_causal_chunk_and_teacher_forcing_pair_counts() -> None: + assert _blockwise_causal_frame_pairs(5, 2, -1) == 17 + assert _blockwise_causal_frame_pairs(5, 2, 3) == 13 + assert _teacher_forcing_frame_pairs(4, 2) == 24 + + +def test_streaming_causal_forward_uses_cache_window() -> None: + module = _FakeTransformer(causal=True) + kwargs = _kwargs(frames=2) + # Four spatial tokens/frame, starting after two frames. The three-frame + # local window gives 8 query tokens x 12 key tokens. + kwargs.update({"kv_cache": [{}], "current_start": 8}) + + work = estimate_transformer_forward( + module, + kwargs, + torch.ones(2, 3, 2, 4, 4), + role="student", + attention_metadata=None, + ) + + assert work is not None + assert work.causal_chunks == 1 + assert work.attention_pairs == 2 * 8 * 12 + assert work.dense_attention_pairs == 2 * 8 * 16 + + +def test_vsa_uses_kernel_topk_density_and_overhead() -> None: + module = _FakeTransformer() + metadata = SimpleNamespace( + variable_block_sizes=torch.tensor([4, 4, 4, 4]), + VSA_sparsity=0.75, + ) + + work = estimate_transformer_forward( + module, + _kwargs(), + torch.ones(2, 3, 4, 4, 4), + role="student", + attention_metadata=metadata, + ) + + assert work is not None + assert work.attention_pairs == pytest.approx(2 * 16**2 * 0.25) + dense_work = estimate_transformer_forward( + module, + _kwargs(), + torch.ones(2, 3, 4, 4, 4), + role="student", + attention_metadata=None, + ) + assert dense_work is not None + # VSA includes its gate projection and pooled dense attention overhead. + assert work.forward_flops > dense_work.forward_flops - ( + dense_work.attention_pairs - work.attention_pairs) * 4 * 8 * 2 + + +def test_streaming_cross_attention_cache_skips_kv_projections() -> None: + module = _FakeTransformer(causal=True) + kwargs = _kwargs(frames=2) + kwargs.update({"kv_cache": [{}], "current_start": 8}) + output = torch.ones(2, 3, 2, 4, 4) + + uncached = estimate_transformer_forward( + module, + kwargs, + output, + role="student", + attention_metadata=None, + ) + cached = estimate_transformer_forward( + module, + kwargs, + output, + role="student", + attention_metadata=None, + cross_attention_cached=True, + ) + + assert uncached is not None and cached is not None + expected_saved = 4 * 2 * 4 * 8**2 * 2 + assert uncached.forward_flops - cached.forward_flops == expected_saved + + +def test_monitor_counts_multiple_rollouts_and_roles_without_double_counting_accumulation() -> None: + models = { + role: SimpleNamespace(transformer=_FakeTransformer()) + for role in ("student", "teacher", "critic") + } + monitor = TrainingPerformanceMonitor() + monitor.attach(models) + monitor.reset() + + with set_forward_context(current_timestep=0, attn_metadata=None): + models["student"].transformer(**_kwargs()) + models["student"].transformer(**_kwargs()) + models["critic"].transformer(**_kwargs()) + with torch.no_grad(): + models["teacher"].transformer(**_kwargs()) + models["teacher"].transformer(**_kwargs()) + + metrics = monitor.metrics( + step_time_sec=2.0, + local_batch_size=2, + grad_accum=4, + world_size=8, + sp_size=4, + peak_tflops_per_gpu=100.0, + ) + monitor.close() + + assert metrics["perf/model_forward_calls"] == 5 + assert metrics["perf/causal_chunks"] == 0 + assert metrics["perf/role/student/forward_calls"] == 2 + assert metrics["perf/role/teacher/forward_calls"] == 2 + assert metrics["perf/role/teacher/backward_forwards"] == 0 + assert metrics["perf/role/critic/backward_forwards"] == 1 + # 2 local samples x DP=2 x grad-accum=4 / 2 seconds. + assert metrics["perf/samples_per_sec"] == 8 + assert metrics["perf/estimated_mfu"] > 0 + + +def test_monitor_counts_forward_when_flop_estimator_does_not_support_model() -> None: + monitor = TrainingPerformanceMonitor() + model = SimpleNamespace(transformer=_UnsupportedTransformer()) + monitor.attach({"student": model}) + monitor.reset() + + model.transformer(hidden_states=torch.ones(1)) + metrics = monitor.metrics( + step_time_sec=1.0, + local_batch_size=1, + grad_accum=1, + world_size=1, + sp_size=1, + peak_tflops_per_gpu=None, + ) + monitor.close() + + assert metrics["perf/model_forward_calls"] == 1 + assert metrics["perf/role/student/forward_calls"] == 1 + assert "perf/estimated_tflops_per_gpu" not in metrics + + +@pytest.mark.parametrize( + ("device_name", "expected"), + [ + ("NVIDIA A40", 149.7), + ("NVIDIA L40S", 362.05), + ("NVIDIA H100 80GB HBM3", 989.5), + ("NVIDIA H200", 989.5), + ("NVIDIA B200", 2250.0), + ("NVIDIA GB200", 2500.0), + ("Some Future GPU", None), + ], +) +def test_infer_peak_bf16_tflops(device_name: str, expected: float | None) -> None: + assert infer_peak_bf16_tflops(device_name) == expected diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index be6054f77..adb5cd7e5 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -13,6 +13,10 @@ 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.performance import ( + TrainingPerformanceMonitor, + infer_peak_bf16_tflops, +) from fastvideo.train.utils.tracking import build_tracker if TYPE_CHECKING: @@ -72,6 +76,7 @@ class Trainer: callback_configs or {}, training_config, ) + self.performance_monitor = TrainingPerformanceMonitor() def _iter_dataloader(self, dataloader: Any) -> Iterator[dict[str, Any]]: data_iter = iter(dataloader) @@ -113,6 +118,13 @@ class Trainer: int(tc.loop.gradient_accumulation_steps or 1), ) + performance_enabled = bool(tc.performance.enabled) + peak_tflops = tc.performance.peak_tflops_per_gpu + if performance_enabled: + if peak_tflops is None and torch.cuda.is_available(): + peak_tflops = infer_peak_bf16_tflops(torch.cuda.get_device_name(torch.cuda.current_device())) + self.performance_monitor.attach(getattr(method, "_role_models", {}), ) + method.set_tracker(self.tracker) method.on_train_start() self.callbacks.on_train_start( @@ -141,6 +153,11 @@ class Trainer: # 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, ) + # Pre-loop forwards (teacher cache generation, step-0 validation) are + # not part of a training step; drop them so a long cache build cannot + # accumulate per-forward records. + if performance_enabled: + self.performance_monitor.reset() progress = tqdm( range(start_step + 1, max_steps + 1), initial=start_step, @@ -151,6 +168,8 @@ class Trainer: method_manages_optimization = bool(method.manages_optimization()) for step in progress: t0 = time.perf_counter() + if performance_enabled: + self.performance_monitor.reset() # Accumulate on GPU during grad-accum; materialise # to CPU once per step right before logging. @@ -219,8 +238,27 @@ class Trainer: 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) + step_time_sec = time.perf_counter() - t0 + metrics["step_time_sec"] = step_time_sec metrics["vsa_sparsity"] = float(tc.vsa_sparsity) + if performance_enabled: + metrics.update( + self.performance_monitor.metrics( + step_time_sec=step_time_sec, + local_batch_size=tc.data.train_batch_size, + grad_accum=(1 if method_manages_optimization else grad_accum), + world_size=getattr( + self.world_group, + "world_size", + tc.distributed.num_gpus, + ), + sp_size=getattr( + self.sp_group, + "world_size", + tc.distributed.sp_size, + ), + peak_tflops_per_gpu=peak_tflops, + )) if self.global_rank == 0 and metrics: self.tracker.log(metrics, step) @@ -251,4 +289,5 @@ class Trainer: if checkpoint_manager is not None: checkpoint_manager.save_final(max_steps) + self.performance_monitor.close() self.tracker.finish() diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index d396b0258..c34fcb0a0 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -18,6 +18,7 @@ from fastvideo.train.utils.training_config import ( DistributedConfig, ModelTrainingConfig, OptimizerConfig, + PerformanceConfig, TrackerConfig, TrainingLoopConfig, ) @@ -351,6 +352,7 @@ def _build_training_config( lo = dict(t.get("loop", {}) or {}) ck = dict(t.get("checkpoint", {}) or {}) tr = dict(t.get("tracker", {}) or {}) + perf = dict(t.get("performance", {}) or {}) vs = dict(t.get("vsa", {}) or {}) m = dict(t.get("model", {}) or {}) @@ -382,6 +384,15 @@ def _build_training_config( "{'t2v', 't2va', 'text_only'}, got " f"{preprocessed_data_type!r}") + peak_tflops = get_optional_float( + perf, + "peak_tflops_per_gpu", + where="training.performance.peak_tflops_per_gpu", + ) + if peak_tflops is not None and peak_tflops <= 0.0: + raise ValueError("training.performance.peak_tflops_per_gpu must be " + f"> 0, got {peak_tflops}") + return TrainingConfig( distributed=DistributedConfig( num_gpus=num_gpus, @@ -429,6 +440,10 @@ def _build_training_config( project_name=str(tr.get("project_name", "fastvideo") or "fastvideo"), run_name=str(tr.get("run_name", "") or ""), ), + performance=PerformanceConfig( + enabled=bool(perf.get("enabled", True)), + peak_tflops_per_gpu=peak_tflops, + ), vsa_sparsity=float(vs.get("sparsity", 0.0) or 0.0), vsa_cache_tile_buf=bool(vs.get("cache_tile_buf", False) or False), model=ModelTrainingConfig( diff --git a/fastvideo/train/utils/performance.py b/fastvideo/train/utils/performance.py new file mode 100644 index 000000000..bef74566f --- /dev/null +++ b/fastvideo/train/utils/performance.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Low-overhead performance accounting for modular training. + +The monitor counts logical Transformer invocations at the role boundary. This +is important for distillation: one optimizer step may contain several student +rollouts plus critic and teacher forwards. FLOPs are an analytic estimate of +the Transformer core; measured step and sample throughput remain exact wall +clock metrics. +""" + +from __future__ import annotations + +import math +from collections import defaultdict +from dataclasses import dataclass +from typing import Any + +import torch + +from fastvideo.attention.backends.video_sparse_attn import compute_topk +from fastvideo.forward_context import get_forward_context + +_FLOPS_PER_TFLOP = 1.0e12 + + +@dataclass(slots=True) +class _ForwardWork: + role: str + forward_flops: float + useful_flops: float + causal_chunks: int + query_frames: int + query_tokens: int + attention_pairs: float + dense_attention_pairs: float + + +def infer_peak_bf16_tflops(device_name: str) -> float | None: + """Infer dense BF16 tensor-core peak TFLOP/s from a CUDA device name. + + Values intentionally exclude structured sparsity. Explicit configuration + should be used when board form factor or clocks differ from the common + data-center variant. + """ + name = device_name.upper() + # Order matters: GB200 contains B200. + peaks = ( + ("GB200", 2500.0), + ("B200", 2250.0), + ("H200", 989.5), + ("H100", 989.5), + ("L40S", 362.05), + ("A100", 312.0), + ("A40", 149.7), + ) + return next((peak for token, peak in peaks if token in name), None) + + +def _arch_config(module: torch.nn.Module) -> Any | None: + config = getattr(module, "config", None) + return getattr(config, "arch_config", config) + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, list | tuple): + return next((item for item in value if isinstance(item, torch.Tensor)), None) + return None + + +def _output_requires_grad(output: Any) -> bool: + if isinstance(output, torch.Tensor): + return bool(output.requires_grad) + if isinstance(output, dict): + return any(_output_requires_grad(value) for value in output.values()) + if isinstance(output, list | tuple): + return any(_output_requires_grad(value) for value in output) + return False + + +def _blockwise_causal_frame_pairs( + num_frames: int, + frames_per_block: int, + local_attn_size: int, +) -> int: + """Frame-level pairs used by full-sequence block-causal training.""" + total = 0 + block = max(1, frames_per_block) + for start in range(0, num_frames, block): + end = min(start + block, num_frames) + attended_start = 0 if local_attn_size < 0 else max(0, end - local_attn_size) + total += (end - start) * (end - attended_start) + return total + + +def _teacher_forcing_frame_pairs( + num_frames: int, + frames_per_block: int, +) -> int: + """Pairs for causal Wan's concatenated ``[clean | noisy]`` mask. + + Clean and noisy query blocks each attend every frame, which is the same + count as a full block-causal mask with no local window. + """ + return 2 * _blockwise_causal_frame_pairs(num_frames, frames_per_block, -1) + + +def _vsa_attention_pairs( + seq_len: int, + metadata: Any, +) -> tuple[float, int]: + """Return nominal sparse token pairs and tile count. + + The kernel keeps a clamped top-k of key tiles per query tile; the nominal + pair count scales the dense count by that tile fraction. + """ + block_sizes = getattr(metadata, "variable_block_sizes", None) + num_blocks = int(block_sizes.numel()) if block_sizes is not None else 0 + if num_blocks <= 0: + return float(seq_len * seq_len), 0 + sparsity = float(getattr(metadata, "VSA_sparsity", 0.0)) + topk = compute_topk(sparsity, num_blocks) + return float(seq_len * seq_len) * topk / num_blocks, num_blocks + + +def estimate_transformer_forward( + module: torch.nn.Module, + kwargs: dict[str, Any], + output: Any, + *, + role: str, + attention_metadata: Any | None, + cross_attention_cached: bool = False, +) -> _ForwardWork | None: + """Estimate one Wan-style Transformer invocation from its logical input.""" + hidden_states = _first_tensor(kwargs.get("hidden_states")) + if hidden_states is None or hidden_states.ndim != 5: + return None + + arch = _arch_config(module) + if arch is None: + return None + required = ("hidden_size", "ffn_dim", "num_layers", "patch_size") + if any(not hasattr(arch, name) for name in required): + return None + + batch_size = int(hidden_states.shape[0]) + raw_frames = int(hidden_states.shape[2]) + raw_height = int(hidden_states.shape[3]) + raw_width = int(hidden_states.shape[4]) + patch_size = arch.patch_size + patch: tuple[int, + ...] = ((1, patch_size, + patch_size) if isinstance(patch_size, int) else tuple(int(value) for value in patch_size)) + if len(patch) != 3 or any(value <= 0 for value in patch): + return None + + frames = raw_frames // patch[0] + spatial_tokens = (raw_height // patch[1]) * (raw_width // patch[2]) + seq_len = frames * spatial_tokens + if seq_len <= 0: + return None + + hidden_size = int(arch.hidden_size) + ffn_dim = int(arch.ffn_dim) + num_layers = int(arch.num_layers) + + context = _first_tensor(kwargs.get("encoder_hidden_states")) + context_tokens = int(context.shape[-2]) if context is not None and context.ndim >= 2 else 0 + image_context = _first_tensor(kwargs.get("encoder_hidden_states_image")) + if image_context is not None and image_context.ndim >= 2: + context_tokens += int(image_context.shape[-2]) + + is_causal = hasattr(module, "num_frame_per_block") + teacher_forcing = is_causal and _first_tensor(kwargs.get("clean_x")) is not None + query_tokens = seq_len + dense_pairs = float(seq_len * seq_len) + vsa_tiles = 0 + causal_chunks = 0 + + if is_causal: + frames_per_block = int(getattr(module, "num_frame_per_block", 1)) + causal_chunks = math.ceil(frames / max(1, frames_per_block)) + local_attn_size = int(getattr(module, "local_attn_size", -1)) + kv_cache = kwargs.get("kv_cache") + if kv_cache is not None: + current_start = int(kwargs.get("current_start", 0)) + if local_attn_size >= 0: + max_frames = local_attn_size + else: + # Causal Wan caps the KV window at + # GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES (21) frames when + # local_attn_size is unset; sliding_window_num_frames only + # sizes the streaming KV cache. MatrixGame2's 15-frame + # compatibility window is not modeled here. + max_frames = 21 + key_tokens = min(current_start + seq_len, max_frames * spatial_tokens) + attention_pairs = float(seq_len * key_tokens) + dense_pairs = float(seq_len * (current_start + seq_len)) + elif teacher_forcing: + query_tokens = 2 * seq_len + attention_pairs = float(_teacher_forcing_frame_pairs(frames, frames_per_block) * spatial_tokens**2) + dense_pairs = float(query_tokens * query_tokens) + else: + attention_pairs = float( + _blockwise_causal_frame_pairs(frames, frames_per_block, local_attn_size) * spatial_tokens**2) + else: + attention_pairs, vsa_tiles = _vsa_attention_pairs(seq_len, attention_metadata) + + # Causal Wan pads text to ``text_len`` before every cross-attention block. + if is_causal: + context_tokens = max(context_tokens, int(getattr(module, "text_len", 0))) + + b = batch_size + seq_length = query_tokens + dim = hidden_size + # Four self-attention projections, two MLP projections, cross-attention + # Q/K/V/out projections, and the two attention matmuls. + per_layer = (8.0 * b * seq_length * dim * dim + 4.0 * b * seq_length * dim * ffn_dim + + 4.0 * b * seq_length * dim * dim + + (0.0 if cross_attention_cached else 4.0 * b * context_tokens * dim * dim) + + 4.0 * b * seq_length * context_tokens * dim + 4.0 * b * attention_pairs * dim) + + if vsa_tiles: + # Wan VSA adds one gate projection and dense pooled QK/AV matmuls. + per_layer += 2.0 * b * seq_length * dim * dim + per_layer += 4.0 * b * vsa_tiles * vsa_tiles * dim + + forward_flops = per_layer * num_layers + backward_expected = _output_requires_grad(output) + # Standard MFU counts model-algorithm FLOPs, excluding activation + # checkpoint recompute: training is approximately F + 2B = 3F. + useful_flops = forward_flops * (3.0 if backward_expected else 1.0) + return _ForwardWork( + role=role, + forward_flops=forward_flops, + useful_flops=useful_flops, + causal_chunks=causal_chunks, + query_frames=b * (2 * frames if teacher_forcing else frames), + query_tokens=b * query_tokens, + attention_pairs=b * attention_pairs, + dense_attention_pairs=b * dense_pairs, + ) + + +class TrainingPerformanceMonitor: + """Count per-role Transformer work for one optimizer step at a time.""" + + def __init__(self) -> None: + self._handles: list[Any] = [] + self._work: list[_ForwardWork] = [] + self._forward_calls: dict[str, int] = defaultdict(int) + self._backward_forwards: dict[str, int] = defaultdict(int) + + def attach(self, role_models: dict[str, Any]) -> None: + self.close() + for role, model in role_models.items(): + transformer = getattr(model, "transformer", None) + if not isinstance(transformer, torch.nn.Module): + continue + + cached_cross_attention: list[bool] = [] + + def pre_hook( + module: torch.nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + *, + cache_state: list[bool] = cached_cross_attention, + ) -> None: + del module, args + cache = kwargs.get("crossattn_cache") + initialized = (isinstance(cache, list) and bool(cache) + and all(isinstance(item, dict) and bool(item.get("is_init", False)) for item in cache)) + cache_state.append(initialized) + + def hook( + module: torch.nn.Module, + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + *, + role_name: str = role, + cache_state: list[bool] = cached_cross_attention, + ) -> None: + del args + self._forward_calls[role_name] += 1 + if _output_requires_grad(output): + self._backward_forwards[role_name] += 1 + try: + metadata = get_forward_context().attn_metadata + except AssertionError: + metadata = None + try: + work = estimate_transformer_forward( + module, + kwargs, + output, + role=role_name, + attention_metadata=metadata, + cross_attention_cached=(cache_state.pop() if cache_state else False), + ) + except Exception: + # Metrics must never abort a training run. + work = None + if work is not None: + self._work.append(work) + + self._handles.append(transformer.register_forward_pre_hook(pre_hook, with_kwargs=True)) + self._handles.append(transformer.register_forward_hook(hook, with_kwargs=True)) + + def reset(self) -> None: + self._work.clear() + self._forward_calls.clear() + self._backward_forwards.clear() + + def close(self) -> None: + for handle in self._handles: + handle.remove() + self._handles.clear() + + def metrics( + self, + *, + step_time_sec: float, + local_batch_size: int, + grad_accum: int, + world_size: int, + sp_size: int, + peak_tflops_per_gpu: float | None, + ) -> dict[str, float]: + if step_time_sec <= 0.0: + return {} + world_size = max(1, int(world_size)) + sp_size = max(1, int(sp_size)) + data_parallel_size = max(1, world_size // sp_size) + accumulation = max(1, int(grad_accum)) + sample_scale_to_world = data_parallel_size * accumulation + work_scale_to_world = data_parallel_size + + metrics = { + "perf/steps_per_sec": 1.0 / step_time_sec, + "perf/samples_per_sec": (max(0, int(local_batch_size)) * sample_scale_to_world / step_time_sec), + "perf/model_forward_calls": float(sum(self._forward_calls.values())), + } + + for role, calls in self._forward_calls.items(): + prefix = f"perf/role/{role}" + metrics[f"{prefix}/forward_calls"] = float(calls) + metrics[f"{prefix}/backward_forwards"] = float(self._backward_forwards[role]) + + if not self._work: + return metrics + + useful_flops = sum(item.useful_flops for item in self._work) * work_scale_to_world + query_frames = sum(item.query_frames for item in self._work) * work_scale_to_world + query_tokens = sum(item.query_tokens for item in self._work) * work_scale_to_world + attention_pairs = sum(item.attention_pairs for item in self._work) + dense_pairs = sum(item.dense_attention_pairs for item in self._work) + achieved_tflops_per_gpu = useful_flops / step_time_sec / world_size / _FLOPS_PER_TFLOP + + metrics.update({ + "perf/causal_chunks": float(sum(item.causal_chunks for item in self._work)), + "perf/query_latent_frames_per_sec": query_frames / step_time_sec, + "perf/query_tokens_per_sec": query_tokens / step_time_sec, + "perf/attention_density": (attention_pairs / dense_pairs if dense_pairs > 0.0 else 1.0), + "perf/estimated_tflops_per_gpu": achieved_tflops_per_gpu, + }) + if peak_tflops_per_gpu is not None: + metrics["perf/peak_tflops_per_gpu"] = peak_tflops_per_gpu + metrics["perf/estimated_mfu"] = achieved_tflops_per_gpu / peak_tflops_per_gpu + + by_role: dict[str, list[_ForwardWork]] = defaultdict(list) + for item in self._work: + by_role[item.role].append(item) + for role, items in by_role.items(): + prefix = f"perf/role/{role}" + metrics[f"{prefix}/causal_chunks"] = float(sum(item.causal_chunks for item in items)) + role_flops = sum(item.useful_flops for item in items) * work_scale_to_world + metrics[f"{prefix}/estimated_tflops_per_gpu"] = (role_flops / step_time_sec / world_size / _FLOPS_PER_TFLOP) + return metrics diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 25a546cef..85e838ec8 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -68,6 +68,16 @@ class TrackerConfig: run_name: str = "" +@dataclass(slots=True) +class PerformanceConfig: + """Low-overhead per-step training performance metrics.""" + + enabled: bool = True + # Dense BF16 tensor-core peak for one GPU. When unset, known NVIDIA + # accelerators are inferred from their device name. + peak_tflops_per_gpu: float | None = None + + @dataclass(slots=True) class ModelTrainingConfig: weighting_scheme: str = "uniform" @@ -87,6 +97,7 @@ class TrainingConfig: loop: TrainingLoopConfig = field(default_factory=TrainingLoopConfig) checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) tracker: TrackerConfig = field(default_factory=TrackerConfig) + performance: PerformanceConfig = field(default_factory=PerformanceConfig) vsa_sparsity: float = 0.0 # Reuse the per-step padded VSA tile buffer across attention layers. # Defaults to False for training: under full activation checkpointing the