FastVideo/examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml at 8e5d6090571d30b13bc34e690da48ab667756fe1 · hao-ai-lab/FastVideo · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# LTX-2 T2V overfitting test config with NVFP4 quantization-aware
# training (QAT).
#
# Same recipe as overfit_ltx2_t2v.yaml, but the curated attention/FFN
# linear set used by LTX-2 NVFP4 deployment runs its 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. Video attention uses ATTN_QAT_TRAIN for its
# quantized forward and STE backward, then ATTN_QAT_INFER during validation.
# Head-dim-64 audio attention and masked text attention remain dense.
#
# Validation-time ATTN_QAT_INFER is arch-aware:
#   * sm_120a/sm_121a: fastvideo-kernel CUTLASS extension -- the exact
#     quantization scheme ATTN_QAT_TRAIN simulates.
#   * sm_100a (GB200) / sm_103a (GB300): FP4 FA4 kernel
#     (github.com/hao-ai-lab/flash-attention-fp4, branch fp4; per-16
#     block-scaled NVFP4 Q/K, BF16 P/V; validated install set:
#     nvidia-cutlass-dsl==4.4.2, quack-kernels==0.4.1,
#     flashinfer-python==0.6.8, FASTVIDEO_FA4=1 — see
#     docs/inference/optimizations.md). This scheme DIFFERS from the
#     CUTLASS one the training simulation matches, so sm_100/sm_103
#     validation and deployment carry a train-sim mismatch -- gate quality
#     by MS-SSIM measurement rather than assuming parity. The resolution
#     receipt ("ATTN_QAT_INFER resolved: ...") records arch + scheme.
#
# Preprocess data first (same data as the bf16 overfit):
#   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
#
# Run on multi-GPU sm_120 hardware:
#   NUM_GPUS=4 \
#     bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml
#
# On GB200 without flash-attention-fp4 installed (or any other arch with no
# ATTN_QAT_INFER kernel), disable only the validation-time swap:
#   NUM_GPUS=4 \
#     bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml \
#       --callbacks.validation.attn_qat_infer false

models:
  student:
    _target_: fastvideo.train.models.ltx2.LTX2Model
    init_from: FastVideo/LTX2-Distilled-Diffusers
    trainable: true
    enable_gradient_checkpointing_type: full
    attention_backend: ATTN_QAT_TRAIN

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_attn_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_attn_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
    # Validation reuses the live transformer and temporarily replaces its
    # ATTN_QAT_TRAIN implementations with the sm120 inference kernel.
    attn_qat_infer: true

# quant_config selects NVFP4 QAT for the same curated linear prefixes as
# LTX-2 NVFP4 deployment. They fake-quantize through real FP4 GEMMs with
# an STE backward. The string resolves to NVFP4QATTrainConfig at parse.
pipeline:
  dit_config:
    quant_config: nvfp4_qat_train