From 21e749156c1f48cae130c4b75f735ff2ee8a59df Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Mon, 19 Oct 2026 12:16:45 -0700 Subject: [PATCH] [fix]: address QAD review findings - Include V2 width/height/num_frames and the train manifest in the step-0 fingerprint so a dimension change cannot reuse a stale baseline. - Honor CompressionSpec.first_steps in the V2 rung scope. - Mark tagged checkpoints weights-only in metadata and teach the exporter to resolve and load them, so best/ can be exported. - Add a tile-64 backward comparison against the STE reference. --- docs/training/minimax_h3_fp4_qad.md | 3 ++- .../tests/test_attn_qat_vsa_train.py | 20 +++++++++++++++++++ .../train/entrypoint/dcp_to_diffusers.py | 11 ++++++++-- .../knowledge_distillation/pdd_qad_step0.py | 4 ++-- fastvideo/train/models/minimax_h3/v2_t2v.py | 3 ++- fastvideo/train/utils/checkpoint.py | 6 +++++- 6 files changed, 40 insertions(+), 7 deletions(-) diff --git a/docs/training/minimax_h3_fp4_qad.md b/docs/training/minimax_h3_fp4_qad.md index 1727faebd..3f2657521 100644 --- a/docs/training/minimax_h3_fp4_qad.md +++ b/docs/training/minimax_h3_fp4_qad.md @@ -131,7 +131,8 @@ directory: The saved state is the bf16 master plus the quantization plan, not packed NVFP4 weights; inference applies the same converter used for the released model. Use `fastvideo/train/entrypoint/dcp_to_diffusers.py` to export a checkpoint if your -deployment path expects Diffusers weights: +deployment path expects Diffusers weights. It detects the weights-only marker +on a tagged checkpoint, so `best/` can be exported directly: ```bash python -m fastvideo.train.entrypoint.dcp_to_diffusers \ diff --git a/fastvideo-kernel/tests/test_attn_qat_vsa_train.py b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py index 4d13bfa91..09fc08308 100644 --- a/fastvideo-kernel/tests/test_attn_qat_vsa_train.py +++ b/fastvideo-kernel/tests/test_attn_qat_vsa_train.py @@ -237,6 +237,26 @@ def test_triton_tile64_forward_matches_reference(two_level_p, smooth_k): assert _rel(out, expected) < 2e-3 < 0.1 * _rel(expected, exact) +@needs_gpu +@pytest.mark.parametrize("two_level_p", [False, True]) +def test_triton_tile64_backward_matches_ste_reference(two_level_p): + q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, _, _ = _problem64(6, heads=2, seed=11, device="cuda") + q, k, v = (t.clone().requires_grad_(True) for t in (q, k, v)) + out = _triton().fp4_vsa_attn_qat(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad=q2k_quad, + two_level_p=two_level_p) + grad_out = torch.randn_like(out) + dq, dk, dv = torch.autograd.grad(out, (q, k, v), grad_out) + _, weights = ref.fp4_vsa_attention_reference(q, k, v, q2k_idx, q2k_num, kv_valid, q2k_quad, + two_level_p=two_level_p, return_dv_weights=True) + q_hat, k_hat, v_hat = ref.fake_quant_qkv(q.detach(), k.detach(), v.detach()) + token_mask = _dense_quad_mask(q2k_idx, q2k_num, q2k_quad, q.shape[2]) + valid = ref.key_valid_mask(kv_valid, q2k_idx.shape[-1], q.device).reshape(-1) + token_mask = token_mask & valid[None, None, None, :] + expected = ref.ste_gradients_reference(q_hat, k_hat, v_hat, token_mask, weights, grad_out, 1 / math.sqrt(128)) + for name, got, want in zip("qkv", (dq, dk, dv), expected, strict=True): + assert _rel(got, want) < 2e-2, name + + @needs_gpu @pytest.mark.parametrize("two_level_p,smooth_k", [(False, False), (True, False), (False, True), (True, True)]) def test_triton_forward_matches_reference(two_level_p, smooth_k): diff --git a/fastvideo/train/entrypoint/dcp_to_diffusers.py b/fastvideo/train/entrypoint/dcp_to_diffusers.py index 0eb3587e0..116c7bb4f 100644 --- a/fastvideo/train/entrypoint/dcp_to_diffusers.py +++ b/fastvideo/train/entrypoint/dcp_to_diffusers.py @@ -40,6 +40,7 @@ import gc import json import os import sys +from pathlib import Path from typing import Any import fastvideo.envs as envs @@ -458,20 +459,24 @@ def convert( import torch.distributed.checkpoint as dcp # -- Resolve checkpoint directory -- - resolved = _resolve_resume_checkpoint( + # 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, ) dcp_dir = resolved / "dcp" if not dcp_dir.is_dir(): raise FileNotFoundError(f"Missing dcp/ under {resolved}") + metadata = CheckpointManager.load_metadata(resolved) + weights_only = bool(metadata.get("weights_only", False)) # -- Obtain config -- cfg: RunConfig if config_path is not None: cfg = load_run_config(config_path) else: - metadata = CheckpointManager.load_metadata(resolved) raw_config = metadata.get("config") if raw_config is None: raise ValueError("Checkpoint metadata.json does not " @@ -500,6 +505,8 @@ def convert( # -- Load DCP weights into the model -- states = method.checkpoint_state() + if weights_only: + states = {key: value for key, value in states.items() if key.startswith("roles.")} logger.info( "Loading DCP checkpoint from %s", resolved, diff --git a/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py b/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py index 7587c0b33..f98f3c2fe 100644 --- a/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py +++ b/fastvideo/train/methods/knowledge_distillation/pdd_qad_step0.py @@ -23,9 +23,9 @@ METHOD_KEYS = ("init_student_dcp", "max_eval_rows", "max_rollout_rows", "noise_f "eligibility_keys", "fidelity_tolerance", "objective", "recon_normalize", "recon_norm_rows", "recon_group_weights", "rungs_per_row") DATA_KEYS = ("manifests", "cases", "max_frames", "eval_manifest", "eval_resolutions", "eval_per_group", - "eval_max_frames", "eval_seed") + "eval_max_frames", "eval_seed", "width", "height", "num_frames") # The training-row plan; it only reaches step 0 through provisional FFN calibration (no amax_json). -PROVISIONAL_DATA_KEYS = ("seed", "per_group", "max_train_rows") +PROVISIONAL_DATA_KEYS = ("seed", "per_group", "max_train_rows", "train_manifest") QAD_KEYS = ("linears", "fp4_attention") # Logged per evaluation but not part of the step-0 result. VOLATILE = ("eval/seconds", "eval/score_vs_ptq") diff --git a/fastvideo/train/models/minimax_h3/v2_t2v.py b/fastvideo/train/models/minimax_h3/v2_t2v.py index 53049715e..5f527cd21 100644 --- a/fastvideo/train/models/minimax_h3/v2_t2v.py +++ b/fastvideo/train/models/minimax_h3/v2_t2v.py @@ -240,7 +240,8 @@ class MiniMaxH3V2T2VModel(MiniMaxH3OmniRefPDDModel): with contextlib.ExitStack() as stack: stack.enter_context(set_forward_context(current_timestep=rung, attn_metadata=self._metadata(prepared, rung))) - if self.compression is not None and self._compression_on: + first_steps = self.compression.spec.first_steps if self.compression is not None else None + if self.compression is not None and self._compression_on and (first_steps is None or rung < first_steps): grid = self._token_grid(prepared) sub = compressed_grid(self.compression.spec.mode, grid) stack.enter_context(self.compression.scope(grid, self._metadata(prepared, rung, sub))) diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index 1d520fa4c..72b4b4c38 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/fastvideo/train/utils/checkpoint.py @@ -301,7 +301,7 @@ class CheckpointManager: if tagged_dir.exists(): shutil.rmtree(tagged_dir) os.makedirs(dcp_dir, exist_ok=True) - self._write_metadata(tagged_dir, step) + self._write_metadata(tagged_dir, step, weights_only=True) logger.info("Saving %s weights (step %d) to %s", tag, step, tagged_dir) _barrier() states = {key: value for key, value in self.method.checkpoint_state().items() if key.startswith("roles.")} @@ -315,8 +315,12 @@ class CheckpointManager: self, checkpoint_dir: Path, step: int, + *, + weights_only: bool = False, ) -> None: metadata: dict[str, Any] = {"step": step} + if weights_only: + metadata["weights_only"] = True if self._raw_config is not None: metadata["config"] = self._raw_config meta_path = checkpoint_dir / "metadata.json"