[feat] Add phase-aware modular training performance metrics (#1826)

Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

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/<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

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)

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)

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
 # ---------------------------------------------------------------------------

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

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()

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(

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

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