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

Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com> Co-authored-by: Loay Rashid <42599591+loaydatrain@users.noreply.github.com> Co-authored-by: Kaiqin Kong <k1kong@ucsd.edu> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
2026-06-17 15:18:52 -07:00
parent eed9c4bfbf
commit e60601df7f
6 changed files with 265 additions and 2 deletions
@@ -0,0 +1,119 @@
+# MixKit training data (QAD 5090 recipe)
+
+The QAD 5090 models are distilled from Wan2.1-T2V-1.3B on a MixKit subset at
+**480×832, 77 frames, 16 fps**. FastVideo training consumes **Parquet** shards of
+precomputed VAE latents + text embeddings (no text encoder / VAE needed at train
+time).
+
+## Option A — download the preprocessed data (recommended)
+
+The encoded dataset is published on the Hugging Face Hub, ready to train:
+
+```bash
+# from the repo root
+bash examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh
+```
+
+This pulls [`weizhou03/HD-Mixkit-Finetune-Wan`](https://huggingface.co/datasets/weizhou03/HD-Mixkit-Finetune-Wan)
+into `data/HD-Mixkit-Finetune-Wan/`:
+
+```
+data/HD-Mixkit-Finetune-Wan/
+├── combined_parquet_dataset/      # training shards  -> point --data_path here
+│   └── worker_0/data_chunk_*.parquet
+└── validation_parquet_dataset/    # validation shards
+    └── worker_0/data_chunk_0.parquet
+```
+
+Each Parquet row holds the VAE latent bytes + text-embedding bytes (plus
+shape/dtype metadata), matching FastVideo's standard preprocessing output.
+
+## Option B — build the Parquet from raw videos
+
+If you want to reproduce the encoding from your own MixKit videos, arrange them as
+a `merged` dataset (videos + a captions JSON), then run FastVideo's standard
+preprocessing to VAE-encode and text-embed them into Parquet:
+
+```bash
+GPU_NUM=2
+torchrun --nproc_per_node=$GPU_NUM \
+    -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+    --model_path "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" \
+    --mode preprocess \
+    --workload_type t2v \
+    --preprocess.video_loader_type torchvision \
+    --preprocess.dataset_type merged \
+    --preprocess.dataset_path "data/mixkit_raw/" \
+    --preprocess.dataset_output_dir "data/HD-Mixkit-Finetune-Wan/" \
+    --preprocess.max_height 480 \
+    --preprocess.max_width 832 \
+    --preprocess.num_frames 77 \
+    --preprocess.train_fps 16 \
+    --preprocess.samples_per_file 8
+```
+
+The raw videos are full-HD MixKit clips (≈1080p/30fps); preprocessing resizes to
+480×832, resamples to 16 fps, and extracts 77 frames per clip. See
+[`docs/training/data_preprocess.md`](../../../../../docs/training/data_preprocess.md)
+for the full parameter reference.
+
+## Train (QAT finetune)
+
+With the data in place, run the quantization-aware finetune. The 4-bit attention
+path is **config-driven** — selected purely by an env var, no monkey-patching:
+
+```bash
+bash examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh
+# or point at your own parquet dir / GPU count:
+NUM_GPUS=4 bash .../mixkit/finetune_qat.sh data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset/
+```
+
+`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` routes attention through the
+fake-quantized Triton kernel (straight-through estimator), so the DiT learns to
+absorb FP4 attention error. This kernel is Triton, so it runs on both `sm_100`
+(B200/GB200) and `sm_120` (RTX 5090).
+
+## Train stage 2 (QAT DMD distillation to 3 steps)
+
+Distill the QAT-finetuned generator down to **3 sampling steps**. Only the
+generator is quantized (Attn-QAT); the teacher (`real_score`) and critic
+(`fake_score`) stay full precision. This is enforced in the loader
+(`component_loader.py`, via the `_loading_teacher_critic_model` flag), so the
+same global `ATTN_QAT_TRAIN` env reaches **only** the generator — no per-model
+flags or monkey-patching.
+
+```bash
+# generator init = the stage-1 finetune checkpoint
+bash examples/training/finetune/wan_t2v_1.3B/mixkit/distill_dmd_qat.sh \
+    data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset/ \
+    checkpoints/wan_t2v_qat_finetune/checkpoint-2000/transformer/diffusion_pytorch_model.safetensors
+```
+
+DMD runs a double loop (critic every step, generator every
+`generator_update_interval`), and validation samples the distilled student at
+3 steps — the final 4-bit-attention model.
+
+## Inference (NVFP4 4-bit linear)
+
+For Wan-2.1, enable the FP4 linear layers with the **`nvfp4_qat`** quantization
+config (it matches Wan's `to_q/k/v/out` + `ffn` layers; the plain `NVFP4` config
+is LTX2-specific and will not quantize Wan):
+
+```python
+from fastvideo import VideoGenerator
+from fastvideo.layers.quantization import get_quantization_config
+
+gen = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", num_gpus=1,
+    transformer_quant=get_quantization_config("nvfp4_qat")(),  # a config instance, not the string
+    use_fsdp_inference=False,
+)
+gen.generate(request={"prompt": "...", "output": {"save_video": True}})
+```
+
+The loader converts the tagged linear weights to FP4 at load time
+(`_maybe_convert_model_to_nvfp4`). Combine with
+`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER` on an RTX 5090 (`sm_120`) for the
+full 4-bit path; on other GPUs the attention falls back to Flash while the FP4
+linear layers still run. `flashinfer` (and a host C++ compiler for its FP4
+kernel JIT) are required.
@@ -0,0 +1,51 @@
+#!/bin/bash
+# QAD recipe stage 2 — quantization-aware DMD distillation of Wan2.1-T2V-1.3B
+# down to 3 sampling steps, with the GENERATOR in fake-quant Attn-QAT and the
+# teacher (real_score) + critic (fake_score) at full precision.
+#
+# Generator-only QAT is config-driven: FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN
+# is applied to the generator only, because the loader masks it (and the
+# nvfp4_qat quant) for the teacher/critic via the `_loading_teacher_critic_model`
+# flag (see fastvideo/models/loader/component_loader.py). No monkey-patching.
+#
+# Init the generator from the stage-1 finetune checkpoint (finetune_qat.sh).
+# Data: run download_mixkit_data.sh first.
+#
+# Verified end-to-end on Blackwell (GB200/sm_100): generator loads with
+# ATTN_QAT_TRAIN while teacher/critic load full-precision; the DMD double loop
+# runs (generator updates every generator_update_interval steps, critic every
+# step), 3-step validation generates videos, checkpoint saved.
+set -euo pipefail
+
+export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN   # generator-only (loader-gated)
+export WANDB_MODE=${WANDB_MODE:-online}
+export TOKENIZERS_PARALLELISM=false
+
+BASE="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR=${1:-"data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset/"}
+# Generator init weights = the stage-1 QAT-finetune checkpoint.
+INIT_WEIGHTS=${2:-"checkpoints/wan_t2v_qat_finetune/checkpoint-2000/transformer/diffusion_pytorch_model.safetensors"}
+VALIDATION_FILE="$(dirname "$0")/../crush_smol/validation.json"
+NUM_GPUS=${NUM_GPUS:-4}
+
+torchrun --nnodes 1 --nproc_per_node "${NUM_GPUS}" \
+    fastvideo/training/wan_distillation_pipeline.py \
+    --num_gpus "${NUM_GPUS}" --sp_size 1 --tp_size 1 \
+    --hsdp_replicate_dim "${NUM_GPUS}" --hsdp_shard_dim 1 \
+    --model_path "${BASE}" --pretrained_model_name_or_path "${BASE}" \
+    --real_score_model_path "${BASE}" --fake_score_model_path "${BASE}" \
+    --init_weights_from_safetensors "${INIT_WEIGHTS}" \
+    --data_path "${DATA_DIR}" --dataloader_num_workers 4 \
+    --max_train_steps 2000 --train_batch_size 1 --train_sp_batch_size 1 \
+    --gradient_accumulation_steps 1 \
+    --num_latent_t 20 --num_height 480 --num_width 832 --num_frames 77 \
+    --enable_gradient_checkpointing_type full \
+    --log_validation --validation_dataset_file "${VALIDATION_FILE}" \
+    --validation_steps 200 --validation_sampling_steps 3 --validation_guidance_scale 6.0 \
+    --learning_rate 2e-6 --mixed_precision bf16 --weight_decay 0.01 --max_grad_norm 1.0 \
+    --weight_only_checkpointing_steps 500 --training_state_checkpointing_steps 500 \
+    --tracker_project_name wan_t2v_distill_dmd_qat \
+    --output_dir checkpoints/wan_t2v_distill_dmd_qat \
+    --inference_mode False --dit_precision fp32 --ema_start_step 0 --training_cfg_rate 0.0 \
+    --generator_update_interval 5 --real_score_guidance_scale 2.0 \
+    --dmd_denoising_steps '1000,757,522' --min_timestep_ratio 0.02 --max_timestep_ratio 0.98
@@ -0,0 +1,22 @@
+#!/bin/bash
+# Download the preprocessed MixKit finetune dataset used for the QAD 5090 recipe.
+#
+# This is the MixKit subset already VAE-encoded (Wan2.1-T2V-1.3B) and text-embedded
+# into Parquet shards, so it can be fed straight to training with no further
+# preprocessing. To build the Parquet from raw videos yourself, see README.md.
+#
+# Usage (run from the repo root):
+#   bash examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh [DATA_ROOT]
+set -euo pipefail
+
+DATA_ROOT=${1:-data/HD-Mixkit-Finetune-Wan}
+
+python scripts/huggingface/download_hf.py \
+    --repo_id "weizhou03/HD-Mixkit-Finetune-Wan" \
+    --local_dir "${DATA_ROOT}" \
+    --repo_type "dataset"
+
+echo "Done."
+echo "  Train data:      ${DATA_ROOT}/combined_parquet_dataset"
+echo "  Validation data: ${DATA_ROOT}/validation_parquet_dataset"
+echo "Point your training script's data path at the combined_parquet_dataset directory."
@@ -0,0 +1,44 @@
+#!/bin/bash
+# QAD recipe — quantization-aware finetune of Wan2.1-T2V-1.3B with fake-quant
+# (Attn-QAT) attention.
+#
+# The 4-bit attention path is selected purely by env var (config-driven, no
+# monkey-patching): FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN routes attention
+# through the fake-quantized Triton kernel (straight-through estimator), so the
+# DiT learns to absorb FP4 attention error instead of fighting it.
+#
+# Data: run download_mixkit_data.sh first (preprocessed Parquet).
+#
+# Verified end-to-end on Blackwell (GB200/sm_100): the ATTN_QAT_TRAIN backend is
+# selected (not a fallback), forward+backward run, loss/grad are healthy, and
+# validation generates videos. The kernel is Triton so it runs on sm_100 and
+# sm_120 alike (the FP4 inference kernel, by contrast, is sm_120-only).
+set -euo pipefail
+
+export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN   # <-- enables Attn-QAT training
+export WANDB_MODE=${WANDB_MODE:-online}
+export TOKENIZERS_PARALLELISM=false
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR=${1:-"data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset/"}
+VALIDATION_FILE="$(dirname "$0")/../crush_smol/validation.json"
+NUM_GPUS=${NUM_GPUS:-4}
+
+torchrun --nnodes 1 --nproc_per_node "${NUM_GPUS}" \
+    fastvideo/training/wan_training_pipeline.py \
+    --num_gpus "${NUM_GPUS}" --sp_size "${NUM_GPUS}" --tp_size 1 \
+    --hsdp_replicate_dim 1 --hsdp_shard_dim "${NUM_GPUS}" \
+    --model_path "${MODEL_PATH}" --pretrained_model_name_or_path "${MODEL_PATH}" \
+    --data_path "${DATA_DIR}" --dataloader_num_workers 1 \
+    --max_train_steps 2000 --train_batch_size 1 --train_sp_batch_size 1 \
+    --gradient_accumulation_steps 1 \
+    --num_latent_t 20 --num_height 480 --num_width 832 --num_frames 77 \
+    --enable_gradient_checkpointing_type full \
+    --log_validation --validation_dataset_file "${VALIDATION_FILE}" \
+    --validation_steps 200 --validation_sampling_steps 50 --validation_guidance_scale 3.0 \
+    --learning_rate 5e-5 --mixed_precision bf16 --weight_decay 1e-4 --max_grad_norm 1.0 \
+    --weight_only_checkpointing_steps 500 --training_state_checkpointing_steps 500 \
+    --tracker_project_name wan_t2v_qat_finetune --output_dir checkpoints/wan_t2v_qat_finetune \
+    --inference_mode False --training_cfg_rate 0.1 --not_apply_cfg_solver \
+    --dit_precision fp32 --num_euler_timesteps 50 --ema_start_step 0 \
+    --multi_phased_distill_schedule "4000-1"
@@ -942,6 +942,20 @@ class TransformerLoader(ComponentLoader):
        dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config)
        dit_config.update_model_arch(config)

+        # Generator-only QAT for DMD distillation: the teacher (real_score) and
+        # critic (fake_score) transformers load with this flag set and must stay
+        # full precision. Drop the nvfp4_qat quant from their copied config, and
+        # mask the global ATTN_QAT_TRAIN env so their attention falls back to dense
+        # (the backend is read globally at build time). The generator loads without
+        # the flag and keeps both.
+        _qat_generator_only = hasattr(fastvideo_args, "_loading_teacher_critic_model")
+        _qat_prev_attn_env = None
+        if _qat_generator_only:
+            dit_config.quant_config = None
+            from fastvideo.attention.selector import _cached_get_attn_backend
+            _qat_prev_attn_env = os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None)
+            _cached_get_attn_backend.cache_clear()
+
        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)

        # Find all safetensors files
@@ -1028,6 +1042,12 @@ class TransformerLoader(ComponentLoader):
            torch_compile_kwargs=fastvideo_args.torch_compile_kwargs,
        )

+        if _qat_generator_only:
+            from fastvideo.attention.selector import _cached_get_attn_backend
+            if _qat_prev_attn_env is not None:
+                os.environ["FASTVIDEO_ATTENTION_BACKEND"] = _qat_prev_attn_env
+            _cached_get_attn_backend.cache_clear()
+
        total_params = sum(p.numel() for p in model.parameters())
        logger.info("Loaded model with %.2fB parameters", total_params / 1e9)

@@ -50,13 +50,20 @@ def _maybe_convert_model_to_nvfp4(model: nn.Module) -> None:
    from fastvideo.layers.quantization.nvfp4_config import (
        NVFP4QuantizeMethod, convert_model_to_nvfp4,
    )
+    from fastvideo.layers.quantization.nvfp4_qat_config import (
+        NVFP4QATQuantizeMethod, convert_model_to_fp4,
+    )

    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


 # TODO(PY): move this to utils elsewhere