[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

@@ -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 \

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

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-<step>, 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,

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

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

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"