From 18ca8c97012e31ac4f3d1bbc5b31d44bdceeba52 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Thu, 9 Jul 2026 11:44:49 +0000 Subject: [PATCH] [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). --- .../configs/overfit_ltx2_t2v_nvfp4_qat.yaml | 99 +++++++++++++++++++ .../test_e2e_ltx2_overfit_new_stack.py | 9 ++ .../fixtures/ltx2_t2v_qat_finetune_min.yaml | 50 ++++++++++ .../tests/train/methods/grad_norm_refs.json | 3 + .../tests/train/methods/test_ltx2_finetune.py | 90 +++++++++++++++++ fastvideo/train/utils/config.py | 19 ++++ 6 files changed, 270 insertions(+) create mode 100644 examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml create mode 100644 fastvideo/tests/train/fixtures/ltx2_t2v_qat_finetune_min.yaml diff --git a/examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml b/examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml new file mode 100644 index 000000000..75007168f --- /dev/null +++ b/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 diff --git a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py b/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py index 180d707d2..b63f0bfaf 100644 --- a/fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py +++ b/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", + }, } diff --git a/fastvideo/tests/train/fixtures/ltx2_t2v_qat_finetune_min.yaml b/fastvideo/tests/train/fixtures/ltx2_t2v_qat_finetune_min.yaml new file mode 100644 index 000000000..9ac4697b4 --- /dev/null +++ b/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 diff --git a/fastvideo/tests/train/methods/grad_norm_refs.json b/fastvideo/tests/train/methods/grad_norm_refs.json index ca07d583b..e7b241c3c 100644 --- a/fastvideo/tests/train/methods/grad_norm_refs.json +++ b/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 }, diff --git a/fastvideo/tests/train/methods/test_ltx2_finetune.py b/fastvideo/tests/train/methods/test_ltx2_finetune.py index 29b041e65..2e5376fa7 100644 --- a/fastvideo/tests/train/methods/test_ltx2_finetune.py +++ b/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) diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index 3bebe61e0..99363f95c 100644 --- a/fastvideo/train/utils/config.py +++ b/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):