diff --git a/docs/attention/attn_qat/index.md b/docs/attention/attn_qat/index.md index 4b6913030..2fffe6749 100644 --- a/docs/attention/attn_qat/index.md +++ b/docs/attention/attn_qat/index.md @@ -1,59 +1,33 @@ -# Attn QAT +# Attention QAT -Attn QAT in FastVideo currently has two related backends: +Attention QAT in FastVideo covers two related, but different, backends: -- `ATTN_QAT_TRAIN`: the training-oriented Triton attention path. -- `ATTN_QAT_INFER`: the inference-oriented CUDA kernel path. +- `ATTN_QAT_INFER`: the inference-oriented CUDA kernel path +- `ATTN_QAT_TRAIN`: the training-oriented Triton attention path -Both are selected with `FASTVIDEO_ATTENTION_BACKEND`, but they come from -different codepaths and serve different purposes. +Both are selected with `FASTVIDEO_ATTENTION_BACKEND`, but they are not +interchangeable. The main practical split is: -## Backend Overview +- use `ATTN_QAT_INFER` for standalone inference with the dedicated inference + kernel +- use `ATTN_QAT_TRAIN` for finetuning, validation during training, or when you + specifically want to reproduce the training-side attention path -| Backend | Best for | Package requirement | Primary kernel location | -|---------|----------|---------------------|-------------------------| -| `ATTN_QAT_TRAIN` | finetuning, validation during training, reproducing the training attention path | `fastvideo_kernel` | `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` | -| `ATTN_QAT_INFER` | standalone inference with the `attn_qat_infer` CUDA kernel | `attn_qat_infer` from the in-repo `fastvideo-kernel` checkout | `fastvideo-kernel/attn_qat_infer/` | +## Quick Start -In FastVideo itself, backend selection is routed through: +If your goal is "run Wan 2.1 14B with Attention QAT inference weights", this is +the shortest path: -- `fastvideo/envs.py` -- `fastvideo/platforms/cuda.py` -- `fastvideo/attention/backends/attn_qat_train.py` -- `fastvideo/attention/backends/attn_qat_infer.py` +1. Build the in-repo kernel package so FastVideo can import `attn_qat_infer`. +2. Download the Wan 2.1 14B QAT checkpoint. +3. Edit the provided inference example to point at the 14B base model and the + downloaded QAT safetensors. +4. Run the example with `ATTN_QAT_INFER`. -The legacy training pipeline also contains explicit Attn QAT integration: +### Step 1. Build the kernel package -- `fastvideo/training/training_pipeline.py` - -That pipeline forces generator loading through `ATTN_QAT_TRAIN` when -`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` or -`--generator_4bit_attn` is enabled. - -## Where The Code Lives - -Use these paths when you want to trace or modify the Attn QAT flow: - -| Location | Purpose | -|----------|---------| -| `fastvideo/attention/backends/attn_qat_train.py` | FastVideo wrapper that imports and calls the Triton training kernel | -| `fastvideo/attention/backends/attn_qat_infer.py` | FastVideo wrapper that imports and calls the inference kernel | -| `fastvideo-kernel/CMakeLists.txt` | Kernel build definition that compiles the `attn_qat_infer` inference extensions | -| `fastvideo/platforms/cuda.py` | Chooses the concrete attention backend at runtime | -| `fastvideo/envs.py` | Documents supported `FASTVIDEO_ATTENTION_BACKEND` values | -| `fastvideo/training/training_pipeline.py` | Training-time forcing logic for the generator attention backend | -| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` | Triton implementation for `ATTN_QAT_TRAIN` | -| `fastvideo-kernel/attn_qat_infer/api.py` | Python API entrypoint for the inference kernel | -| `fastvideo-kernel/benchmarks/benchmark_*.py` | Kernel-side attention benchmark scripts for FlashAttn2, SageAttn3, FP4, and comparison plots | -| `fastvideo-kernel/attn_qat_infer/blackwell/api.cu` | CUDA implementation behind `ATTN_QAT_INFER` | -| `fastvideo-kernel/tests/test_attn_qat_train.py` | Kernel-level test coverage for the training path | -| `examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 1.3B Attn QAT finetune launcher | -| `examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 14B Attn QAT finetune launcher | - -## Build And Install - -Before using either backend, make sure the in-repo kernel package is built from -source. +Before using either Attention QAT backend, build the in-repo +`fastvideo-kernel` package from source: ```bash git submodule update --init --recursive @@ -61,80 +35,125 @@ cd fastvideo-kernel ./build.sh ``` -For local development, the important result is: +After a successful build: -- `ATTN_QAT_TRAIN` can import `fastvideo_kernel` -- `ATTN_QAT_INFER` can import `attn_qat_infer` +- `ATTN_QAT_TRAIN` should be able to import `fastvideo_kernel` +- `ATTN_QAT_INFER` should be able to import `attn_qat_infer` -Backend-specific notes: +`ATTN_QAT_INFER` currently targets the Blackwell CUDA path under +`fastvideo-kernel/attn_qat_infer/` and requires CUDA 12.8+. -- `ATTN_QAT_TRAIN` is the Triton training path shipped through - `fastvideo-kernel`. -- `ATTN_QAT_INFER` currently targets the Blackwell CUDA path in - `fastvideo-kernel/attn_qat_infer/`, is built by `fastvideo-kernel/CMakeLists.txt`, - and requires CUDA 12.8+. +### Step 2. Download the Wan 2.1 14B QAT checkpoint -For the full kernel build requirements and CUDA notes, see: +FastVideo includes a helper script: -- [Attention Overview](../index.md) -- `fastvideo-kernel/README.md` +- `examples/inference/optimizations/download_14B_qat.sh` -## Running Training +By default it downloads: -Today the checked-in Attn QAT training launchers use the legacy training -pipeline in `fastvideo/training/wan_training_pipeline.py`. +- Hugging Face repo: `FastVideo/14B_qat_400` +- local directory: `checkpoints/14B_qat_400` -### Ready-made launchers +Prerequisites: -Use the provided SLURM scripts directly: +- `huggingface_hub` installed, for example: + `uv pip install huggingface_hub` +- access to the model repo if it is private or gated: + `huggingface-cli login` + +Run the downloader: ```bash -sbatch examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh -sbatch examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh +bash examples/inference/optimizations/download_14B_qat.sh ``` -Both scripts already set: +To download into a custom directory: ```bash -export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN +bash examples/inference/optimizations/download_14B_qat.sh /path/to/14B_qat_400 ``` -Before launching, update the script-local values that are environment-specific: +The script prints a ready-to-copy `init_weights_from_safetensors=...` value at +the end. -- `WANDB_API_KEY` -- `MODEL_PATH` -- `DATA_DIR` -- `VALIDATION_DATASET_FILE` -- output directory and SLURM resource requests +### Step 3. Edit the provided inference example -### What the launchers run +The example to start from is: -The training scripts eventually invoke: +- `examples/inference/optimizations/attn_qat_inference_example.py` + +Open that file and update these two values: + +1. Change the base model from `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` to + `Wan-AI/Wan2.1-T2V-14B-Diffusers` +2. Replace + `init_weights_from_safetensors="safetensors_path"` with the directory that + contains the downloaded `.safetensors` files + +Example: + +```python +import os + +from fastvideo import VideoGenerator + +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "ATTN_QAT_INFER" + +generator = VideoGenerator.from_pretrained( + "Wan-AI/Wan2.1-T2V-14B-Diffusers", + num_gpus=1, + use_fsdp_inference=True, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=False, + init_weights_from_safetensors="checkpoints/14B_qat_400", +) +``` + +Important: + +- the checked-in example currently uses the `1.3B` base model until you edit it +- do not load the 14B QAT weights on top of the `1.3B` base model; the weights + and model config will not match + +### Step 4. Run the inference example ```bash -torchrun fastvideo/training/wan_training_pipeline.py ... +python examples/inference/optimizations/attn_qat_inference_example.py ``` -If you are adapting the workflow to your own cluster or running outside SLURM, -the key Attn QAT requirement is still the same: +Generated videos are written to `video_samples/` by default. -```bash -export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN -``` +## Backend Overview -Then launch the normal Wan training pipeline with your preferred `torchrun` -arguments and training flags. +| Backend | Best for | Package requirement | Primary kernel location | +|---------|----------|---------------------|-------------------------| +| `ATTN_QAT_TRAIN` | finetuning, training-time validation, reproducing the training path | `fastvideo_kernel` | `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` | +| `ATTN_QAT_INFER` | standalone inference with the dedicated CUDA kernel | `attn_qat_infer` from the in-repo `fastvideo-kernel` checkout | `fastvideo-kernel/attn_qat_infer/` | -## Running Inference +FastVideo routes backend selection through: -For standalone inference, prefer `ATTN_QAT_INFER` when that kernel is available. -Use `ATTN_QAT_TRAIN` only if you specifically want to exercise the training-side -attention path during validation or debugging. +- `fastvideo/envs.py` +- `fastvideo/platforms/cuda.py` +- `fastvideo/attention/backends/attn_qat_train.py` +- `fastvideo/attention/backends/attn_qat_infer.py` -### Python API +The legacy training pipeline also contains explicit Attention QAT integration: -This follows the same pattern as `examples/inference/basic/basic.py`, but with -the Attn QAT backend forced through the environment: +- `fastvideo/training/training_pipeline.py` + +That pipeline forces generator loading through `ATTN_QAT_TRAIN` when +`FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN` or `--generator_4bit_attn` is +enabled. + +## Inference Workflows + +For standalone inference, prefer `ATTN_QAT_INFER` when the CUDA kernel is +available. Use `ATTN_QAT_TRAIN` for inference only if you intentionally want to +exercise the training-side attention path for debugging or parity checks. + +### Minimal Python example ```python import os @@ -155,17 +174,35 @@ generator.generate_video( ) ``` -To run inference with a finetuned checkpoint instead of the base model, replace -the model ID with your checkpoint path. +### Loading custom safetensors during inference -### CLI +FastVideo supports loading custom transformer weights through +`init_weights_from_safetensors`. -You can also force the backend when using the FastVideo CLI: +This value can point to either: + +- a directory containing one or more `.safetensors` files +- a single `.safetensors` file + +For Wan 2.1 14B QAT inference, the common pattern is: + +```python +generator = VideoGenerator.from_pretrained( + "Wan-AI/Wan2.1-T2V-14B-Diffusers", + num_gpus=1, + use_fsdp_inference=True, + init_weights_from_safetensors="checkpoints/14B_qat_400", +) +``` + +### CLI example + +You can also force the backend from the command line: ```bash FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER \ fastvideo generate \ - --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \ --num-gpus 1 \ --sp-size 1 \ --tp-size 1 \ @@ -178,8 +215,87 @@ fastvideo generate \ --output-path outputs_video/ ``` -This is the same basic flow used by `scripts/inference/v1_inference_wan.sh`, -just with `FASTVIDEO_ATTENTION_BACKEND` set explicitly. +If you want to use custom QAT transformer weights from the CLI, pass the same +custom weight override that the Python API uses: + +```bash +FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_INFER \ +fastvideo generate \ + --model-path Wan-AI/Wan2.1-T2V-14B-Diffusers \ + --init-weights-from-safetensors checkpoints/14B_qat_400 \ + --num-gpus 1 \ + --output-path outputs_video/ \ + --prompt "A cinematic close-up of rain on a neon street at night." +``` + +## Training Workflows + +Today the checked-in Attention QAT training launchers use the legacy training +pipeline in `fastvideo/training/wan_training_pipeline.py`. + +### Ready-made launchers + +Use the provided SLURM scripts directly: + +```bash +sbatch examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh +sbatch examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh +``` + +Both scripts already set: + +```bash +export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN +``` + +Before launching, update the script-local values that depend on your +environment: + +- `WANDB_API_KEY` +- `MODEL_PATH` +- `DATA_DIR` +- `VALIDATION_DATASET_FILE` +- output directory and SLURM resource requests + +### What the launchers run + +The training scripts eventually invoke: + +```bash +torchrun fastvideo/training/wan_training_pipeline.py ... +``` + +If you are adapting the workflow to your own cluster or running outside SLURM, +the main Attention QAT requirement is still: + +```bash +export FASTVIDEO_ATTENTION_BACKEND=ATTN_QAT_TRAIN +``` + +Then launch the normal Wan training pipeline with your preferred `torchrun` +arguments and training flags. + +## Where The Code Lives + +Use these paths when you want to trace or modify the Attention QAT flow: + +| Location | Purpose | +|----------|---------| +| `fastvideo/attention/backends/attn_qat_train.py` | FastVideo wrapper that imports and calls the Triton training kernel | +| `fastvideo/attention/backends/attn_qat_infer.py` | FastVideo wrapper that imports and calls the inference kernel | +| `fastvideo-kernel/CMakeLists.txt` | Kernel build definition that compiles the `attn_qat_infer` inference extensions | +| `fastvideo/platforms/cuda.py` | Chooses the concrete attention backend at runtime | +| `fastvideo/envs.py` | Documents supported `FASTVIDEO_ATTENTION_BACKEND` values | +| `fastvideo/training/training_pipeline.py` | Training-time forcing logic for the generator attention backend | +| `fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py` | Triton implementation for `ATTN_QAT_TRAIN` | +| `fastvideo-kernel/attn_qat_infer/api.py` | Python API entrypoint for the inference kernel | +| `fastvideo-kernel/benchmarks/benchmark_*.py` | Kernel-side benchmark scripts for FlashAttn2, SageAttention3, FP4, and comparison plots | +| `fastvideo-kernel/attn_qat_infer/blackwell/api.cu` | CUDA implementation behind `ATTN_QAT_INFER` | +| `fastvideo-kernel/tests/test_attn_qat_train.py` | Kernel-level test coverage for the training path | +| `examples/inference/optimizations/attn_qat_inference_example.py` | Ready-to-edit inference example for custom Attention QAT weights | +| `examples/inference/optimizations/download_14B_qat.sh` | Helper script for downloading the Wan 2.1 14B QAT checkpoint | +| `examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 1.3B Attention QAT finetune launcher | +| `examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh` | Ready-to-run Wan 14B Attention QAT finetune launcher | ## Troubleshooting @@ -187,10 +303,16 @@ just with `FASTVIDEO_ATTENTION_BACKEND` set explicitly. successfully and exposes `fastvideo_kernel`. - If `ATTN_QAT_INFER` fails to import, verify that the local build exposes the `attn_qat_infer` package. +- If the Wan 2.1 14B example fails after you changed only the checkpoint path, + make sure you also changed the base model to + `Wan-AI/Wan2.1-T2V-14B-Diffusers`. +- If you hit issues with CPU memory pressure or obscure CUDA argument errors in + the example script, try setting `pin_cpu_memory=False`. - If you want a known-safe fallback for debugging, use `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`. -See also: +## Related Pages +- [Attention Overview](../index.md) - [Inference Optimizations](../../inference/optimizations.md) - [Debugging](../../utilities/debugging.md) diff --git a/docs/attention/index.md b/docs/attention/index.md index 2b37b413e..e274962fa 100644 --- a/docs/attention/index.md +++ b/docs/attention/index.md @@ -5,7 +5,8 @@ FastVideo provides highly optimized custom attention kernels to accelerate video ## Supported Kernels * **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks. -* **[Attn QAT](attn_qat/index.md)**: Training and inference backends for the Attn QAT workflow. +* **[Attention QAT](attn_qat/index.md)**: Dedicated guide for Attention QAT + inference, training, checkpoint loading, and troubleshooting. * **[Sliding Tile Attention (STA)](sta/index.md)**: STA kernel support is kept in `fastvideo-kernel`; full FastVideo STA pipeline workflow is archived in `sta_do_not_delete`. diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index a1e78ab1f..22d6450ef 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -136,6 +136,10 @@ This backend uses the `attn_qat_infer` implementation that lives in the Use this backend when you want to run the dedicated FP4 inference CUDA kernel directly during inference. +For the full Attention QAT guide, including Wan 2.1 14B checkpoint download, +example editing steps, training launchers, and troubleshooting, see +[Attention QAT](../attention/attn_qat/index.md). + This backend currently assumes access to the in-repo `fastvideo-kernel` checkout or an equivalent editable/source install that exposes: @@ -156,6 +160,10 @@ This backend uses the FastVideoKernel Triton attention implementation from want the training-oriented Triton attention path rather than the `attn_qat_infer` CUDA kernel path. +The dedicated [Attention QAT](../attention/attn_qat/index.md) page covers when +to use `ATTN_QAT_TRAIN` versus `ATTN_QAT_INFER`, the ready-made training +launchers, and the end-to-end Wan 2.1 14B inference workflow. + This backend currently assumes access to an install that exposes: - `fastvideo_kernel` diff --git a/mkdocs.yml b/mkdocs.yml index f7b490e98..a9b4b1f1c 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -152,7 +152,7 @@ nav: - DMD: distillation/dmd.md - Attention: - Overview: attention/index.md - - Attn QAT: attention/attn_qat/index.md + - Attention QAT: attention/attn_qat/index.md - Video Sparse Attention: attention/vsa/index.md - Sliding Tile Attention (Archived): attention/sta/index.md - Backend Development: contributing/attention_backend.md