From e60601df7ffcd6bcd1a9d1c298de282015bac6b2 Mon Sep 17 00:00:00 2001 From: alexzms <3036648523@qq.com> Date: Wed, 17 Jun 2026 15:18:52 -0700 Subject: [PATCH] =?UTF-8?q?[feat]=20QAD=205090:=20QAT=20training=20recipe?= =?UTF-8?q?=20=E2=80=94=20finetune=20+=20DMD=20distillation=20(12/12)=20(#?= =?UTF-8?q?1462)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: William Lin Co-authored-by: Loay Rashid <42599591+loaydatrain@users.noreply.github.com> Co-authored-by: Kaiqin Kong Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../finetune/wan_t2v_1.3B/mixkit/README.md | 119 ++++++++++++++++++ .../wan_t2v_1.3B/mixkit/distill_dmd_qat.sh | 51 ++++++++ .../mixkit/download_mixkit_data.sh | 22 ++++ .../wan_t2v_1.3B/mixkit/finetune_qat.sh | 44 +++++++ fastvideo/models/loader/component_loader.py | 20 +++ fastvideo/models/loader/fsdp_load.py | 11 +- 6 files changed, 265 insertions(+), 2 deletions(-) create mode 100644 examples/training/finetune/wan_t2v_1.3B/mixkit/README.md create mode 100755 examples/training/finetune/wan_t2v_1.3B/mixkit/distill_dmd_qat.sh create mode 100755 examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh create mode 100755 examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh diff --git a/examples/training/finetune/wan_t2v_1.3B/mixkit/README.md b/examples/training/finetune/wan_t2v_1.3B/mixkit/README.md new file mode 100644 index 000000000..d6102161c --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/mixkit/README.md @@ -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. diff --git a/examples/training/finetune/wan_t2v_1.3B/mixkit/distill_dmd_qat.sh b/examples/training/finetune/wan_t2v_1.3B/mixkit/distill_dmd_qat.sh new file mode 100755 index 000000000..fb9dc0b7b --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/mixkit/distill_dmd_qat.sh @@ -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 diff --git a/examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh b/examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh new file mode 100755 index 000000000..09fae3259 --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/mixkit/download_mixkit_data.sh @@ -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." diff --git a/examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh b/examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh new file mode 100755 index 000000000..07ae06781 --- /dev/null +++ b/examples/training/finetune/wan_t2v_1.3B/mixkit/finetune_qat.sh @@ -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" diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 64d6bb842..0152d33c0 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -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) diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index dd90d966f..09a834cc7 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -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