diff --git a/docs/training/minimax_h3_fp4_qad.md b/docs/training/minimax_h3_fp4_qad.md index aeb93884a..14f2688ea 100644 --- a/docs/training/minimax_h3_fp4_qad.md +++ b/docs/training/minimax_h3_fp4_qad.md @@ -221,7 +221,6 @@ python -m fastvideo.train.entrypoint.dcp_to_diffusers \ | `fastvideo/train/methods/knowledge_distillation/pdd_qad_metrics.py` | Rung aggregation and score | | `fastvideo/train/methods/knowledge_distillation/pdd_qad_recon.py` | Layer-wise reconstruction hooks | | `fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py` | Step-0 baseline cache | -| `fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py` | Audio guard statistics | | `fastvideo/train/models/minimax_h3/omniref_pdd.py` | OmniRef Ref2VA profile | | `fastvideo/train/models/minimax_h3/v2_t2v.py` | V2/V2.1 T2VA profile | | `fastvideo/layers/quantization/nvfp4_qad.py` | NVFP4 linear emulation and STE | diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py index a6b6b6f44..ed42b9977 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_reference.py @@ -260,9 +260,15 @@ def fp4_vsa_attention_reference(q: torch.Tensor, if position == 0 and first_block_max_floor is not None: group_max = group_max.clamp(min=first_block_max_floor) new_max = torch.maximum(row_max, group_max.amax(dim=-1)) - m_scaled = new_max * scale_log2 + offset - p_tilde = torch.exp2(scores * scale_log2 - m_scaled[:, None]) - alpha = torch.exp2((row_max - new_max) * scale_log2) + # A listed block can have no selected keys for this query + # half; both maxima are -inf and the usual exp2 terms are + # NaN, so emit zeros and keep the running state unchanged. + empty = torch.isinf(new_max) & (new_max < 0) + m_scaled = torch.where(empty, torch.zeros_like(new_max), new_max * scale_log2 + offset) + p_tilde = torch.where(empty[:, None], torch.zeros_like(scores), + torch.exp2(scores * scale_log2 - m_scaled[:, None])) + alpha = torch.where(empty, torch.zeros_like(new_max), + torch.exp2((row_max - new_max) * scale_log2)) alpha = torch.where(torch.isinf(row_max), torch.zeros_like(alpha), alpha) row_sum = row_sum * alpha + p_tilde.sum(dim=-1) if quantize: @@ -279,7 +285,9 @@ def fp4_vsa_attention_reference(q: torch.Tensor, out[b, h, rows] = acc * inv[:, None] if weights is not None: for n, p_hat, block_max in visited: - decay = torch.exp2((block_max - row_max) * scale_log2) * inv + finite = torch.isfinite(block_max) & torch.isfinite(row_max) + decay = torch.where(finite, torch.exp2((block_max - row_max) * scale_log2) * inv, + torch.zeros_like(inv)) weights[b, h, rows, n * BLOCK:(n + 1) * BLOCK] = p_hat * decay[:, None] if weights is not None: return out, weights diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py index f540b3091..a089f5da9 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_vsa_train.py @@ -175,13 +175,20 @@ def _fwd_kernel(Q, K, V, O, OHP, LSE, MRUN, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_QUAD if FLOOR: group_max = tl.where(it == 0, tl.maximum(group_max, FLOOR_VALUE), group_max) m_new = tl.maximum(m_i, tl.max(group_max, axis=1)) + # A valid mask can list a block that has no selected keys for one query + # half (tile 64 pairs two halves per block). Both maxima are then -inf; + # exp2(-inf - -inf) is NaN, so keep the running state unchanged and + # emit zero probabilities for that half. + empty = m_new == float("-inf") m_scaled = m_new * scale_log2 + OFFSET - alpha = tl.math.exp2((m_i - m_new) * scale_log2) + alpha = tl.where(empty, 0.0, tl.math.exp2((m_i - m_new) * scale_log2)) if QUANT: p_tilde, p_hat = _quantized_p(s, m_scaled, it == 0, scale_log2, FLOOR, FLOOR_VALUE, ROWS, BLOCK_N) else: p_tilde = tl.math.exp2(s * scale_log2 - m_scaled[:, None]) p_hat = p_tilde + p_tilde = tl.where(empty[:, None], 0.0, p_tilde) + p_hat = tl.where(empty[:, None], 0.0, p_hat) l_i = l_i * alpha + tl.sum(p_tilde, axis=1) v = tl.load(V + base + keys[:, None] * D + offs_d[None, :]) acc = acc * alpha[:, None] + tl.dot(p_hat.to(tl.bfloat16), v) @@ -228,7 +235,10 @@ def _bwd_dq_kernel(Q, K, V, DO, LSE, DELTA, DQ, Q2K_IDX, Q2K_NUM, KV_VALID, Q2K_ quad = tl.load(Q2K_QUAD + meta * list_width + (it // halves)).to(tl.int32) bit = (quad >> ((pid % halves) * 2 + (it % halves))) & 1 ok = ok & (bit != 0) - p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse[:, None]), 0.0) + # A row with no valid keys at all has lse = -inf; clamp it so an empty + # half contributes zero instead of exp2(-inf - -inf) = NaN. + lse_safe = tl.where(lse == float("-inf"), 0.0, lse) + p = tl.where(ok[None, :], tl.math.exp2(s * scale_log2 - lse_safe[:, None]), 0.0) dp = tl.dot(do, tl.trans(v)) ds = p * (dp - delta[:, None]) dq += tl.dot(ds.to(tl.bfloat16), k) @@ -274,14 +284,16 @@ def _bwd_dkdv_kernel(Q, K, V, DO, LSE, DELTA, MRUN, DK, DV, K2Q_IDX, K2Q_SLOT, K s = tl.where(ok[None, :] & (bit != 0), s, float("-inf")) else: s = tl.where(ok[None, :], s, float("-inf")) - p = tl.math.exp2(s * scale_log2 - lse[:, None]) + lse_safe = tl.where(lse == float("-inf"), 0.0, lse) + p = tl.math.exp2(s * scale_log2 - lse_safe[:, None]) if QUANT: qmeta = bh * n_qblocks + m_block first_visit = slot == tl.load(Q2K_NUM + qmeta) - 1 m_run = tl.load(MRUN + (qmeta * list_width + slot) * BLOCK_N + row0 + offs_r) - m_scaled = m_run * scale_log2 + OFFSET + m_ok = m_run != float("-inf") + m_scaled = tl.where(m_ok, m_run * scale_log2 + OFFSET, 0.0) _, p_hat = _quantized_p(s, m_scaled, first_visit, scale_log2, FLOOR, FLOOR_VALUE, ROWS, ROWS) - weight = p_hat * tl.math.exp2(m_scaled - lse)[:, None] + weight = tl.where(m_ok[:, None], p_hat * tl.math.exp2(m_scaled - lse_safe)[:, None], 0.0) else: weight = p dv += tl.dot(tl.trans(weight.to(tl.bfloat16)), do) diff --git a/fastvideo-kernel/tests/test_attn_qat_vsa_train.py b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py index 09fc08308..dc7385395 100644 --- a/fastvideo-kernel/tests/test_attn_qat_vsa_train.py +++ b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py @@ -205,6 +205,30 @@ def test_tile128_route_is_the_general_route_without_quadrants(): assert torch.equal(q2k_idx, old_idx) and torch.equal(q2k_num, old_num) and torch.equal(kv_valid, old_valid) +def test_tile64_first_listed_block_can_be_empty_for_one_half(): + """Regression: an empty half in the first visited block must not poison later blocks.""" + n_tiles = 4 + gen = torch.Generator().manual_seed(7) + seq_len = (n_tiles // 2) * ref.BLOCK + q, k, v = (torch.randn(1, 1, seq_len, 128, generator=gen).bfloat16() for _ in range(3)) + tile_mask = torch.zeros(1, 1, n_tiles, n_tiles, dtype=torch.bool) + tile_mask[0, 0, 1, 2] = True # query tile 1 lists block 1 (empty for query tile 0) + tile_mask[0, 0, 0, 0] = True # query tile 0 lists block 0 (valid) + tile_mask[0, 0, 2, 2] = True + tile_mask[0, 0, 3, 3] = True + tile_valid = torch.full((n_tiles, ), 64, dtype=torch.int32) + q2k_idx, q2k_num, kv_valid, q2k_quad = ref.tile_mask_to_fp4_blocks(tile_mask, 64, tile_valid) + assert int(q2k_num[0, 0, 0]) == 2 and int(q2k_idx[0, 0, 0, 0]) == 1 + assert int(q2k_quad[0, 0, 0, 0]) & 3 == 0 # first slot is block 1, empty for query half 0 + + out = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize=False) + assert torch.isfinite(out).all() + mask = _dense_quad_mask(q2k_idx, q2k_num, q2k_quad, seq_len) + scores = (q.float() @ k.float().transpose(-1, -2)) / math.sqrt(128) + expected = torch.softmax(scores.masked_fill(~mask, -math.inf), dim=-1) @ v.float() + torch.testing.assert_close(out, expected, rtol=1e-4, atol=1e-5) + + def test_tile64_quadrant_mask_is_exact_masked_softmax(): q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, _, _ = _problem64(5, seed=5) out = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, quantize=False) diff --git a/fastvideo/models/dits/minimax_h3_token_compression.py b/fastvideo/models/dits/minimax_h3_token_compression.py index fe74fd0d8..f0b4d5439 100644 --- a/fastvideo/models/dits/minimax_h3_token_compression.py +++ b/fastvideo/models/dits/minimax_h3_token_compression.py @@ -6,8 +6,8 @@ original RoPE coordinates and AdaLN row; after the last compressed block every v its nearest kept token: ``H_out = H + expand(F(P H) - P H)`` where ``P`` is the compression operator. With compression off nothing changes (no hook acts), so the student equals the teacher bit for bit. -Unlike ``scripts/benchmarks/minimax_h3_realtime/h3_levers.py`` (which replaces the block list), this installs -forward hooks, so module names, FSDP sharding, activation checkpointing and QAD linears are untouched: +Rather than replacing the block list, this installs forward hooks, so module names, FSDP sharding, activation +checkpointing and QAD linears are untouched: * block ``start`` pre-hook compresses the hidden states and swaps AdaLN rows / RoPE / sequence length; * every block in the range gets the compressed AdaLN rows / RoPE; diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index afd76ba5e..23481e67f 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -343,7 +343,10 @@ def maybe_load_fsdp_model( has_mixed_parameter_dtypes = callable(dtype_selector) and any( dtype_selector(_strip_checkpoint_wrapper_prefix(name), param_dtype) != param_dtype for name, _ in model.named_parameters()) - if training_mode and has_mixed_parameter_dtypes: + # Models that explicitly replicate their frozen FP32 boundary modules keep + # those parameters out of the sharded update, so mixed dtypes are safe. + frozen_fp32_ok = bool(getattr(getattr(model, "config", None), "replicate_fp32_modules", False)) + if training_mode and has_mixed_parameter_dtypes and not frozen_fp32_ok: raise NotImplementedError("FSDP training with model-selected mixed parameter dtypes requires " "separate gradient synchronization for replicated parameters.") diff --git a/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py b/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py index aca155a81..02e133046 100644 --- a/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py +++ b/fastvideo/tests/train/models/test_minimax_h3_omniref_qad.py @@ -104,7 +104,13 @@ def _reference_plan(groups: dict[tuple[str, str], list[dict]], seed: int, per_gr for index in range(max(len(entries) for entries in picked)): for entries in picked: if index < len(entries): - plan.append(entries[index]) + entry = entries[index] + source = entry["id"].split(":", 1)[-1] + plan.append({ + **entry, + "id": f"omniref-{entry['case']}-{source}-{entry['resolution']}", + "source": source, + }) return plan @@ -512,27 +518,6 @@ def test_skip_blocks_keep_whole_blocks_bf16(cpu_loader, single_process_group): assert methods and all(name.startswith("transformer_blocks.0.") for name in methods) -def test_audio_stats_and_seed_floor_gate(): - from fastvideo.train.methods.knowledge_distillation.pdd_qad_audio_stats import (audio_stats, log_mel_distance, - seed_floor_gate) - rate = 32000 - t = torch.arange(rate) / rate - low, high = torch.sin(2 * math.pi * 1000 * t), torch.sin(2 * math.pi * 10000 * t) - stats_low, stats_high = audio_stats(low, rate), audio_stats(high[:, None].repeat(1, 2), rate) - assert stats_low["spectral_centroid_hz"] == pytest.approx(1000, rel=0.05) - assert stats_low["share_above_8khz"] < 0.01 and stats_high["share_above_8khz"] > 0.99 - assert log_mel_distance(low, low, rate) == 0.0 and log_mel_distance(low, high, rate) > 1.0 - teacher = [{"spectral_centroid_hz": 950.0, "share_above_8khz": 0.1, "dynamic_range_db": 30.0, - "loudness_iqr_db": 6.0}] * 3 - other = [{**row, "spectral_centroid_hz": 900.0, "dynamic_range_db": 28.0, "loudness_iqr_db": 5.0, - "share_above_8khz": 0.09} for row in teacher] - good = [{**row, "spectral_centroid_hz": 930.0, "share_above_8khz": 0.095} for row in teacher] - dull = [{**row, "spectral_centroid_hz": 530.0} for row in teacher] - assert seed_floor_gate(good, teacher, other)["pass"] - result = seed_floor_gate(dull, teacher, other) - assert not result["pass"] and not result["spectral_centroid_hz/pass"] - - def test_fp32_boundary_modules_match_inference_and_stay_frozen(cpu_loader, single_process_group): # noqa: F811 """The QAD loader keeps H3's FP32 boundary modules FP32 (as inference) and refuses to train them.""" from types import SimpleNamespace diff --git a/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py b/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py index d3fb89dea..f03528dab 100644 --- a/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py +++ b/fastvideo/tests/train/utils/test_dcp_to_diffusers_h3.py @@ -7,6 +7,7 @@ from types import SimpleNamespace from diffusers import ConfigMixin, ModelMixin from diffusers.configuration_utils import register_to_config +import pytest from safetensors.torch import load_file import torch from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper @@ -16,6 +17,7 @@ from fastvideo.models.loader.weight_utils import safetensors_weights_iterator from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Ref2VAModularPipeline from fastvideo.train.entrypoint.dcp_to_diffusers import ( _native_export_state_dict, + _resolve_export_checkpoint, _save_diffusers_safetensors, _save_role_pretrained, _strict_reload_verify, @@ -23,6 +25,20 @@ from fastvideo.train.entrypoint.dcp_to_diffusers import ( from fastvideo.train.utils.lora import enable_lora_training, finalize_lora_training +def test_resolve_export_checkpoint_accepts_tagged_weights_only_dir(tmp_path: Path) -> None: + tagged = tmp_path / "best" + (tagged / "dcp").mkdir(parents=True) + (tagged / "dcp" / ".metadata").write_text("") + assert _resolve_export_checkpoint(str(tagged)) == tagged.resolve() + + +def test_resolve_export_checkpoint_rejects_plain_dir(tmp_path: Path) -> None: + plain = tmp_path / "best" + plain.mkdir() + with pytest.raises(ValueError, match="best"): + _resolve_export_checkpoint(str(plain)) + + class _ArchConfig: exclude_lora_layers: list[str] = [] diff --git a/fastvideo/train/entrypoint/dcp_to_diffusers.py b/fastvideo/train/entrypoint/dcp_to_diffusers.py index 5795f20bb..8c3d56769 100644 --- a/fastvideo/train/entrypoint/dcp_to_diffusers.py +++ b/fastvideo/train/entrypoint/dcp_to_diffusers.py @@ -429,6 +429,15 @@ def _strict_reload_verify( logger.info("Strict reload verification passed.") +def _resolve_export_checkpoint(checkpoint_dir: str) -> Path: + """A DCP image directory, including a tagged weights-only one, or a regular checkpoint path.""" + raw_path = Path(os.path.expanduser(checkpoint_dir)).resolve() + if (raw_path / "dcp" / ".metadata").is_file(): + return raw_path + from fastvideo.train.utils.checkpoint import _resolve_resume_checkpoint + return _resolve_resume_checkpoint(checkpoint_dir, output_dir=checkpoint_dir) + + def convert( *, checkpoint_dir: str, @@ -447,10 +456,7 @@ def convert( from fastvideo.distributed import ( maybe_init_distributed_environment_and_model_parallel, ) from fastvideo.train.utils.builder import build_from_config - from fastvideo.train.utils.checkpoint import ( - CheckpointManager, - _resolve_resume_checkpoint, - ) + from fastvideo.train.utils.checkpoint import CheckpointManager from fastvideo.train.utils.config import ( RunConfig, load_run_config, @@ -459,13 +465,7 @@ def convert( import torch.distributed.checkpoint as dcp # -- Resolve checkpoint directory -- - # A tagged (weights-only) checkpoint is not named checkpoint-, so it - # bypasses the resume resolver; any directory carrying a DCP image works. - raw_path = Path(os.path.expanduser(checkpoint_dir)).resolve() - resolved = raw_path if (raw_path / "dcp" / ".metadata").is_file() else _resolve_resume_checkpoint( - checkpoint_dir, - output_dir=checkpoint_dir, - ) + resolved = _resolve_export_checkpoint(checkpoint_dir) dcp_dir = resolved / "dcp" if not dcp_dir.is_dir(): raise FileNotFoundError(f"Missing dcp/ under {resolved}") diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py deleted file mode 100644 index b50bd6abe..000000000 --- a/fastvideo/train/methods/knowledge_distillation/pdd_qad_audio_stats.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Decoded-audio statistics and the teacher-seed-floor gate for QAD checkpoints. - -Latent-space audio error does not show what the Trim DMD post-mortem found mattered: decoded audio drifted dull -(spectral centroid 952 Hz teacher -> 530 Hz late checkpoint) while losses looked fine. These statistics, compared -with how much the bf16 teacher itself varies between two seeds, make that drift a selection gate. -""" -from __future__ import annotations - -import math - -import torch - -STATS = ("spectral_centroid_hz", "share_above_8khz", "dynamic_range_db", "loudness_iqr_db") - - -def _spectrogram(wave: torch.Tensor, n_fft: int = 2048, hop: int = 512) -> torch.Tensor: - """Power spectrogram ``[freq, frames]`` of a mono or stereo (``[samples, 2]``) waveform.""" - mono = wave.float().mean(dim=-1) if wave.ndim == 2 else wave.float() - window = torch.hann_window(n_fft) - spec = torch.stft(mono, n_fft=n_fft, hop_length=hop, window=window, return_complex=True) - return spec.abs().pow(2) - - -def audio_stats(wave: torch.Tensor, sample_rate: int) -> dict[str, float]: - power = _spectrogram(wave) - freqs = torch.linspace(0, sample_rate / 2, power.shape[0]) - energy = power.sum(dim=0) - voiced = energy > energy.max() * 1e-6 - centroid = (power * freqs[:, None]).sum(0) / energy.clamp_min(1e-20) - frame_db = 10 * torch.log10(energy.clamp_min(1e-20)) - voiced_db = frame_db[voiced] if bool(voiced.any()) else frame_db - q = torch.quantile(voiced_db, torch.tensor([0.05, 0.25, 0.75, 0.95])) - return { - "spectral_centroid_hz": float(centroid[voiced].median()) if bool(voiced.any()) else 0.0, - "share_above_8khz": float(power[freqs > 8000].sum() / power.sum().clamp_min(1e-20)), - "dynamic_range_db": float(q[3] - q[0]), - "loudness_iqr_db": float(q[2] - q[1]), - } - - -def log_mel_distance(a: torch.Tensor, b: torch.Tensor, sample_rate: int, n_mels: int = 80) -> float: - """Mean absolute difference of log-mel spectrograms (triangular mel filterbank, no torchaudio needed).""" - pa, pb = _spectrogram(a), _spectrogram(b) - frames = min(pa.shape[1], pb.shape[1]) - n_freq = pa.shape[0] - - def mel(f: torch.Tensor) -> torch.Tensor: - return 2595 * torch.log10(1 + f / 700) - - edges = 700 * (10**(torch.linspace(0, float(mel(torch.tensor(sample_rate / 2))), n_mels + 2) / 2595) - 1) - freqs = torch.linspace(0, sample_rate / 2, n_freq) - lo, mid, hi = edges[:-2, None], edges[1:-1, None], edges[2:, None] - bank = torch.clamp(torch.minimum((freqs - lo) / (mid - lo), (hi - freqs) / (hi - mid)), min=0) - la = torch.log(bank @ pa[:, :frames] + 1e-8) - lb = torch.log(bank @ pb[:, :frames] + 1e-8) - return float((la - lb).abs().mean()) - - -def seed_floor_gate(student: list[dict[str, float]], - teacher: list[dict[str, float]], - teacher_other_seed: list[dict[str, float]], - slack: float = 1.0) -> dict[str, float | bool]: - """Per statistic: median |student - teacher| (same seed) must be <= slack * median |teacher(s+1) - teacher(s)|. - - Rows are aligned lists. Returns per-stat student gap, floor and a ``pass`` flag over all statistics. - """ - result: dict[str, float | bool] = {} - passed = True - for name in STATS: - gap = torch.tensor([abs(s[name] - t[name]) for s, t in zip(student, teacher, strict=True)]) - floor = torch.tensor([abs(o[name] - t[name]) for o, t in zip(teacher_other_seed, teacher, strict=True)]) - gap_m, floor_m = float(gap.median()), float(floor.median()) - result[f"{name}/student_gap"] = gap_m - result[f"{name}/teacher_floor"] = floor_m - ok = gap_m <= slack * floor_m or math.isclose(gap_m, 0.0, abs_tol=1e-9) - result[f"{name}/pass"] = ok - passed &= ok - result["pass"] = passed - return result - - -__all__ = ["STATS", "audio_stats", "log_mel_distance", "seed_floor_gate"] diff --git a/fastvideo/train/models/minimax_h3/v2_t2v.py b/fastvideo/train/models/minimax_h3/v2_t2v.py index 5f527cd21..98a093fd6 100644 --- a/fastvideo/train/models/minimax_h3/v2_t2v.py +++ b/fastvideo/train/models/minimax_h3/v2_t2v.py @@ -5,12 +5,12 @@ Same contract as ``omniref_pdd.MiniMaxH3OmniRefPDDModel`` so ``OmniRefPDDQADMeth ("rungs", one per DMD step), ``rung_scope`` / ``forward_rung`` / ``schedulers`` / ``prepare`` / ``quantization``. Differences from OmniRef: plain ``transformer`` module (no PDD head fusion), V2 sampling contract (video shift 10, -audio shift 3, ``dmd_denoising_steps``), VSA tile 64 with sparsity 0.8 in bf16 (native Triton fwd+bwd, no FP4 -attention emulation), text-only rows (``v2_data``), and the optional middle-block spatial token compression -(``v2_compression``). +audio shift 3, ``dmd_denoising_steps``), VSA tile 64 with sparsity 0.8 (BF16 Triton kernels by default; +``qad.fp4_attention`` installs the same FP4 emulator as the shared recipe), text-only rows (``v2_data``), and the +optional middle-block spatial token compression (``v2_compression``). -``quantization(linears=, attention=)``: ``linears`` toggles the NVFP4 QAD linears; ``attention=False`` also switches -off the token compression (the student's only non-linear deviation), so ``quantization(False, False)`` is the teacher. +``quantization(linears=, attention=)``: ``linears`` toggles the NVFP4 QAD linears; ``attention=False`` switches off +the FP4 attention emulator and the token compression, so ``quantization(False, False)`` is the BF16 teacher. """ from __future__ import annotations diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index b96bacafe..c7da738ca 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -103,6 +103,21 @@ class Trainer: if self.global_rank == 0 and validation_metrics: self.tracker.log(validation_metrics, iteration) + def _consume_method_requests(self, method: TrainingMethod, checkpoint_manager: Any | None, step: int) -> bool: + """Save a method-requested tagged checkpoint and report whether training should stop. + + Runs immediately after every evaluation, including the pre-loop one, so + a request always refers to the weights that were just evaluated. + """ + request_hook = getattr(method, "consume_checkpoint_request", None) + tag = request_hook() if callable(request_hook) else None + if tag and checkpoint_manager is not None: + checkpoint_manager.save_tagged(str(tag), step) + stop = bool(getattr(method, "stop_requested", False)) + if stop and self.global_rank == 0: + self.tracker.log({"stopped_early_at": float(step)}, step) + return stop + def run( self, method: TrainingMethod, @@ -144,6 +159,7 @@ class Trainer: iteration=start_step, ) self._run_method_validation(method, start_step) + stop_early = self._consume_method_requests(method, checkpoint_manager, start_step) method.optimizers_zero_grad(start_step) data_stream = self._iter_dataloader(dataloader) @@ -167,7 +183,10 @@ class Trainer: ) # Allow method-specific optimization flow (e.g. DiffusionNFT). method_manages_optimization = bool(method.manages_optimization()) + last_step = start_step for step in progress: + if stop_early: + break t0 = time.perf_counter() if performance_enabled: self.performance_monitor.reset() @@ -283,15 +302,7 @@ class Trainer: ) last_step = step - # Opt-in method hooks: a tagged checkpoint (e.g. the best evaluation - # so far) and an early stop (e.g. guardrail patience exhausted). - request_hook = getattr(method, "consume_checkpoint_request", None) - tag = request_hook() if callable(request_hook) else None - if tag and checkpoint_manager is not None: - checkpoint_manager.save_tagged(str(tag), step) - if bool(getattr(method, "stop_requested", False)): - if self.global_rank == 0: - self.tracker.log({"stopped_early_at": float(step)}, step) + if self._consume_method_requests(method, checkpoint_manager, step): break else: last_step = max_steps