Attn-QAT Training

Attn-QAT simulates low-bit attention during training while keeping the rest of the training method unchanged. In the modular fastvideo/train framework it is a per-role model option, not a separate training method: supervised fine-tuning and DMD2 still own their losses and optimizer cadence.

This guide covers the QAD Wan2.1-T2V-1.3B MixKit workflow:

1. run a 4,000-step supervised Attn-QAT fine-tune;
2. export the stage-1 DCP checkpoint to Diffusers format; and
3. distill the student to three denoising steps with DMD2.

The ready-to-run configs and wrappers are in examples/train/scenario/qad_wan2_1_mixkit/.

Role-local attention backends

A DMD2 run owns three independent model roles. Configure the attention backend on each role so fake quantization is applied only to the student:

models:
  student:
    attention_backend: ATTN_QAT_TRAIN
  teacher:
    attention_backend: FLASH_ATTN
  critic:
    attention_backend: FLASH_ATTN

The override is active only while that role’s transformer is constructed, then the previous process-wide backend is restored. This lets student, teacher, and critic use different implementations in one process. Invalid role-level names fail during configuration instead of silently selecting another backend.

See Training Infrastructure for the complete model-role configuration reference.

Prerequisites

• Install FastVideo and make the fastvideo-kernel Python package importable. ATTN_QAT_TRAIN intentionally fails instead of falling back to dense attention when its kernel cannot be loaded.
• Prepare the precomputed MixKit VAE latents and text embeddings.
• Run the commands below from the repository root. The supplied recipe expects four GPUs by default; set NUM_GPUS to override it.

Download the published preprocessed dataset:

bash examples/datasets/mixkit/download_dataset.sh

Stage 1: supervised Attn-QAT fine-tuning

The stage-1 config uses ATTN_QAT_TRAIN on the student, sequence parallelism across four GPUs, FP32 master weights, and 4,000 optimizer steps:

NUM_GPUS=4 \
  bash examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh

Pass a dataset directory as the first positional argument when it differs from the default:

NUM_GPUS=4 \
  bash examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh \
  /path/to/combined_parquet_dataset

The wrapper calls examples/train/run.sh; the YAML file remains the source of truth for optimizer, validation, checkpointing, and distributed settings.

Export the stage-1 checkpoint

Modular training checkpoints use Distributed Checkpoint (DCP) format. Export the student before using it to initialize stage 2:

bash examples/train/scenario/qad_wan2_1_mixkit/export_stage1.sh \
  checkpoints/wan_t2v_qat_finetune/checkpoint-4000 \
  checkpoints/wan_t2v_qat_finetune/diffusers

Both arguments are optional; the command above shows their defaults.

Stage 2: three-step DMD2 distillation

Stage 2 loads the exported student weights, keeps Attn-QAT on the student, and uses Flash Attention for the teacher and critic:

NUM_GPUS=4 \
  bash examples/train/scenario/qad_wan2_1_mixkit/run_stage2.sh \
  data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset \
  checkpoints/wan_t2v_qat_finetune/diffusers/transformer/model.safetensors

The migrated recipe preserves these behaviors:

The timestep ratios apply to randomly sampled teacher and critic score timesteps; dmd_denoising_steps separately controls the student rollout. See DMD Distillation for general DMD concepts.

Architecture-specific Triton routing

The training kernel is runtime-JIT-compiled Triton code and selects its route on every call. It supports different query and key/value sequence lengths for cross-attention; key and value must have the same sequence length.

Warp specialization is disabled automatically on SM100 and SM120 because the Triton 3.7 NVWS compiler pass aborts for this kernel on Blackwell. No user setting is required.

The available tuning and comparison controls are:

The first invocation JIT-compiles the selected configuration; later calls reuse the Triton cache. To measure the production shape, run python benchmarks/benchmark_attn_qat_train.py from fastvideo-kernel/.

On two B200 GPUs with Wan2.1 1.3B, 31,200 tokens, full activation checkpointing, and AdamW, the default long-sequence SM100 route measured about 10% lower full-step latency than the previous SM100 optimized path in guarded SP2 acceptance runs.

For import and backend-selection failures, see Debugging.