[feat] QAD 5090: QAT training recipe — finetune + DMD distillation (12/12)

[alexzms]

Purpose

Finale of the QAD 5090 stack (#1225). The full quantization-aware training
recipe for Wan2.1-T2V-1.3B: reproduce the MixKit data, run QAT finetune, then
QAT DMD distillation down to 3 steps — plus the loader fix that makes Wan NVFP4
inference work. Every stage was verified end-to-end on Blackwell (GB200/sm_100).

Changes

• mixkit/download_mixkit_data.sh + README.md: one-line download of the
  preprocessed Parquet dataset, and the two-stage recipe how-to.
• mixkit/finetune_qat.sh: stage 1 — QAT finetune. Attn-QAT attention is
  config-driven (FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN).
• mixkit/distill_dmd_qat.sh: stage 2 — QAT DMD distillation to 3 steps,
  generator init from the stage-1 checkpoint.
• fastvideo/models/loader/component_loader.py: generator-only QAT for DMD
  distillation. Teacher (real_score) + critic (fake_score) load with the
  _loading_teacher_critic_model flag; mask the nvfp4_qat quant and the global
  ATTN_QAT_TRAIN env for them, so only the generator runs fake-quant. No
  monkey-patching.
• fastvideo/models/loader/fsdp_load.py: bugfix — convert
  NVFP4QATQuantizeMethod layers to FP4 at load time, so Wan nvfp4_qat
  inference doesn't crash on a missing _fp4_weight.

Test Results (Blackwell GB200 / sm_100)

• ✅ Stage 1 (finetune): generator selects ATTN_QAT_TRAIN, fwd+bwd run,
  loss/grad healthy, validation generates videos, checkpoint saved.
• ✅ Stage 2 (DMD distillation): generator loads ATTN_QAT_TRAIN while
  teacher/critic load full precision (per-model gating verified in the logs); the
  DMD double loop runs (generator updates every generator_update_interval,
  critic every step, healthy loss), 3-step validation generates videos,
  checkpoint saved.
• ✅ Inference: with the loader fix, nvfp4_qat runs real FP4 linear
  inference end-to-end.

Notes

• The QAT scripts use ATTN_QAT_TRAIN, so the recipe depends on [feat] QAD 5090: Wire the Attn-QAT training attention backend (10/12) #1459 (wiring)
  • [kernel] QAD 5090: Add Attn-QAT training Triton kernels (11/12) #1460 (train kernel).
• This is attention QAT; wiring the FP4 linear STE into training (so linear
  is also 4-bit during training) is a follow-up.

Part of #1225.

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🔴 PR merge requirements

Waiting for

• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• check-success=fastcheck-passed
• check-success=full-suite-passed
• #approved-reviews-by>=1
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces the MixKit training data recipe (QAD 5090) for finetuning Wan2.1-T2V-1.3B, including documentation, a data download script, and a quantization-aware finetuning script. It also updates the FSDP model loader to support converting loaded model weights for NVFP4-QAT linear layers. Feedback on the loader implementation suggests scanning the model for all present quantization methods before converting, rather than returning early on the first match, to prevent silently skipping conversions in mixed-method models.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The current implementation returns early on the first match of either NVFP4QuantizeMethod or NVFP4QATQuantizeMethod. If a model contains a mix of both quantization methods, the second method will be silently skipped, which could lead to runtime errors. Scanning the model first to identify all present quantization methods and then running the conversions sequentially is more robust.

Suggested change

	for mod in model.modules():
	if isinstance(getattr(mod, "quant_method", None),
	NVFP4QuantizeMethod):
	qm = getattr(mod, "quant_method", None)
	if isinstance(qm, NVFP4QuantizeMethod):
	logger.info("Converting loaded model weights for NVFP4 linear layers")
	convert_model_to_nvfp4(model)
	return
	if isinstance(qm, NVFP4QATQuantizeMethod):
	logger.info("Converting loaded model weights for NVFP4-QAT linear layers")
	convert_model_to_fp4(model)
	return
	has_nvfp4 = False
	has_nvfp4_qat = False
	for mod in model.modules():
	qm = getattr(mod, "quant_method", None)
	if isinstance(qm, NVFP4QuantizeMethod):
	has_nvfp4 = True
	elif isinstance(qm, NVFP4QATQuantizeMethod):
	has_nvfp4_qat = True
	if has_nvfp4:
	logger.info("Converting loaded model weights for NVFP4 linear layers")
	convert_model_to_nvfp4(model)
	if has_nvfp4_qat:
	logger.info("Converting loaded model weights for NVFP4-QAT linear layers")
	convert_model_to_fp4(model)

[SolitaryThinker]

/merge