[feat] Add NVFP4 QAT support to modular trainer for LTX-2 fine-tuning

Resolve string pipeline.dit_config.quant_config values (e.g.
nvfp4_qat_train) into registered QuantizationConfig instances during
config parsing. The resolver lives in _parse_pipeline_config so it
covers training, resume, and dcp_to_diffusers export (which rebuilds
from checkpoint metadata and bypasses load_run_config).

nvfp4_qat_train runs attention/FFN linear forwards through real NVFP4
GEMMs (flashinfer cutlass) with a straight-through-estimator backward
into dense fp32 master weights; no extra params or buffers, so FSDP2
sharding, DCP checkpoints, and export are unchanged. Attention stays
dense, matching LTX-2 NVFP4 deployment.

Add an NVFP4 QAT overfit config, a QAT smoke test asserting the quant
method attaches to all 1344 matched linears (GB200 grad-norm ref
seeded), and an ltx2_nvfp4_qat nightly e2e case.

Verified on 4x GB200: 300-step cat-sample overfit, ~2x dense step
time, train loss -> 0.022, validation mean MS-SSIM vs the training
clip 0.14 (step 0) -> 0.955 (step 300, nightly e2e PASSED; bf16 dense
twin reaches 0.964).

examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml

@@ -0,0 +1,99 @@
+# LTX-2 T2V overfitting test config with NVFP4 quantization-aware
+# training (QAT).
+#
+# Same recipe as overfit_ltx2_t2v.yaml, but attention/FFN linears run
+# their forward through real NVFP4 GEMMs (flashinfer cutlass) with a
+# straight-through-estimator backward into the dense fp32 master
+# weights. No extra parameters or buffers are added, so FSDP sharding,
+# checkpointing, and export are identical to dense training. Attention
+# stays dense (LTX-2 NVFP4 deployment does not quantize attention math).
+#
+# Requires flashinfer with FP4 support and an sm_100-class GPU
+# (developed on GB200).
+#
+# Preprocess data first (same data as the bf16 overfit):
+#   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
+#
+# Run:
+#   NUM_GPUS=4 bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX2-Distilled-Diffusers
+    trainable: true
+    enable_gradient_checkpointing_type: full
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/ltx2_overfit_preprocessed
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    # LTX2Model requires 0.0: CFG dropout would zero post-connector
+    # embeddings, which is not the model's unconditional input.
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 11  # (81 - 1) / 8 + 1
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 300
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/ltx2_overfit_nvfp4_qat
+    # A full training-state checkpoint is ~150GB for the 13B trainable
+    # video branch; disable saves for the overfit smoke run.
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 1
+
+  tracker:
+    trackers: [wandb]
+    project_name: fastvideo_ltx2
+    run_name: ltx2_overfit_nvfp4_qat
+
+  model:
+    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
+    # so the default noise-minus-clean target reproduces the official
+    # unweighted masked-MSE (mask is all-ones for plain T2V).
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
+    dataset_file: data/ltx2_overfit_preprocessed/validation_prompts.json
+    every_steps: 50
+    sampling_steps: [8]
+    guidance_scale: 1.0
+    num_frames: 81
+
+# quant_config selects NVFP4 QAT: attention/FFN linears (to_q/k/v/out,
+# ffn.fc_in/fc_out) fake-quantize through real FP4 GEMMs with an STE
+# backward. The string resolves to NVFP4QATTrainConfig at config parse.
+pipeline:
+  dit_config:
+    quant_config: nvfp4_qat_train

fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py

@@ -52,6 +52,15 @@ _CASES = {
         "prep_dir": Path(DATA_DIR) / "ltx2_3_overfit_preprocessed",
         "out_dir": Path(DATA_DIR) / "outputs_ltx2_3_overfit",
     },
+    # NVFP4 quantization-aware training: same LTX-2.0 data, attention/FFN
+    # linears forward through real FP4 GEMMs with an STE backward.
+    # Requires flashinfer FP4 kernels (sm_100-class GPU).
+    "ltx2_nvfp4_qat": {
+        "model": "FastVideo/LTX2-Distilled-Diffusers",
+        "config": "examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml",
+        "prep_dir": Path(DATA_DIR) / "ltx2_overfit_preprocessed",
+        "out_dir": Path(DATA_DIR) / "outputs_ltx2_nvfp4_qat_overfit",
+    },
 }
 
 

fastvideo/tests/train/fixtures/ltx2_t2v_qat_finetune_min.yaml

@@ -0,0 +1,50 @@
+# Minimum config to run a single NVFP4-QAT training step of
+# FineTuneMethod on LTX2Model for the per-method smoke test. Identical
+# to ltx2_t2v_finetune_min.yaml except pipeline.dit_config.quant_config
+# selects the NVFP4 QAT-train quantization (STE through real FP4
+# GEMMs). Requires flashinfer FP4 kernels (sm_100-class GPU); the test
+# additionally skips below 60GB GPU memory.
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX2-Distilled-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 3
+    num_height: 256
+    num_width: 256
+    num_frames: 17
+
+  optimizer:
+    learning_rate: 1.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1
+    gradient_accumulation_steps: 1
+
+  model:
+    precondition_outputs: false
+
+pipeline:
+  dit_config:
+    quant_config: nvfp4_qat_train

fastvideo/tests/train/methods/grad_norm_refs.json

@@ -9,6 +9,9 @@
   "test_ltx2_finetune": {
     "GB200": 0.3127
   },
+  "test_ltx2_nvfp4_qat_finetune": {
+    "GB200": 0.3713
+  },
   "test_matrixgame2_finetune": {
     "H200": 4.8247
   },

fastvideo/tests/train/methods/test_ltx2_finetune.py

@@ -157,3 +157,93 @@ def test_ltx2_finetune_single_train_step(
     # Device-keyed grad-norm regression on top of the same harness.
     # Skips when the current GPU has no seeded reference.
     check_grad_norm_regression(ref_name, model.transformer.model)
+
+
+def _flashinfer_fp4_available() -> bool:
+    try:
+        from flashinfer import mm_fp4, nvfp4_quantize  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+@pytest.mark.usefixtures("distributed_setup")
+def test_ltx2_nvfp4_qat_finetune_single_train_step(
+        monkeypatch: pytest.MonkeyPatch) -> None:
+    """Same harness as the dense test, with NVFP4 QAT enabled via
+    ``pipeline.dit_config.quant_config: nvfp4_qat_train``. Additionally
+    asserts the QAT quant method is actually attached to the expected
+    attention/FFN linears."""
+    if _gpu_too_small():
+        pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB "
+                    "memory (LTX-2 DiT is 18.9B params)")
+    if not _flashinfer_fp4_available():
+        pytest.skip("requires flashinfer with FP4 kernels")
+
+    from fastvideo.layers.quantization.nvfp4_qat_train_config import (
+        NVFP4QATTrainConfig,
+        NVFP4QATTrainQuantizeMethod,
+    )
+
+    cfg = load_run_config(str(_FIXTURE_DIR / "ltx2_t2v_qat_finetune_min.yaml"))
+    assert isinstance(cfg.training.pipeline_config.dit_config.quant_config,
+                      NVFP4QATTrainConfig), (
+                          "quant_config string was not resolved to "
+                          "NVFP4QATTrainConfig by config parsing")
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+
+    monkeypatch.setattr(
+        "fastvideo.train.utils.dataloader."
+        "build_parquet_t2v_train_dataloader",
+        lambda *args, **kwargs: None,
+    )
+
+    model = LTX2Model(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=True,
+    )
+    model.transformer = model.transformer.to(device=device, dtype=dtype)
+
+    # 48 audio+video blocks x 28 substring-matched linears (attn1 4 +
+    # attn2 4 + ffn 2, video and audio branches, + a2v 4 + v2a 4).
+    quantized = sum(
+        isinstance(getattr(m, "quant_method", None),
+                   NVFP4QATTrainQuantizeMethod)
+        for m in model.transformer.modules())
+    assert quantized == 1344, (
+        f"expected 1344 NVFP4-QAT linears on the LTX-2 DiT, found "
+        f"{quantized}")
+
+    method = FineTuneMethod(
+        cfg=cfg,
+        role_models={"student": model},
+    )
+    method.on_train_start()
+
+    batch = _build_synthetic_batch(device, dtype, text_dim=3840)
+    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
+
+    loss = loss_map["total_loss"]
+    assert torch.isfinite(loss).item(), (
+        f"total_loss is not finite: {loss.item()}")
+
+    method.backward(loss_map, outputs, grad_accum_rounds=1)
+
+    blocks = resolve_blocks(model.transformer.model)
+    assert blocks is not None and len(blocks) > 0
+    trainable = [p for p in blocks[0].parameters() if p.requires_grad]
+    assert trainable, "layer 0 has no trainable parameters"
+    for i, p in enumerate(trainable):
+        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
+        assert torch.isfinite(p.grad).all().item(), (
+            f"layer 0 param[{i}] grad contains NaN/Inf")
+    assert any(
+        p.grad.detach().float().norm().item() > 0.0 for p in trainable), (
+            "all layer-0 grads are exactly zero; STE backward did not "
+            "reach the first transformer block")
+
+    check_grad_norm_regression("test_ltx2_nvfp4_qat_finetune",
+                               model.transformer.model)

fastvideo/train/utils/config.py

@@ -282,9 +282,28 @@ def _parse_pipeline_config(
 
     pipeline_config = PipelineConfig.from_kwargs(kwargs)
     _apply_training_dit_arch_overrides(pipeline_config, dit_arch_overrides)
+    _resolve_dit_quant_config(pipeline_config)
     return pipeline_config
 
 
+def _resolve_dit_quant_config(pipeline_config: Any) -> None:
+    """Resolve a string ``pipeline.dit_config.quant_config`` (e.g.
+    ``nvfp4_qat_train``) into its registered QuantizationConfig instance.
+
+    Model construction calls ``quant_config.get_quant_method()`` inside
+    ``LinearBase.__init__``, so a bare YAML string would crash there.
+    This must live in ``_parse_pipeline_config`` (not ``load_run_config``)
+    because ``dcp_to_diffusers`` rebuilds models from a checkpoint's raw
+    config by calling ``_parse_pipeline_config`` directly.
+    """
+    dit_config = getattr(pipeline_config, "dit_config", None)
+    quant = getattr(dit_config, "quant_config", None)
+    if isinstance(quant, str):
+        from fastvideo.layers.quantization import (
+            get_quantization_config, )
+        dit_config.quant_config = get_quantization_config(quant)()
+
+
 def _split_training_dit_arch_overrides(pipeline_raw: Any) -> tuple[Any, dict[str, Any]]:
     """Remove train-only DiT arch overrides before generic config parsing."""
     if not isinstance(pipeline_raw, dict) or not isinstance(pipeline_raw.get("dit_config"), dict):