diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 9ef9eff31..fd50a6b62 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -12,6 +12,8 @@ This page describes the various options for speeding up generation times in Fast - [Sage Attention](#sage-attention) - [Sage Attention 3](#sage-attention-3) +- [torch.compile](#torch-compile) + ## Attention Backends ### Available Backends @@ -175,6 +177,106 @@ These backends are model-specific and require the corresponding kernels and dependencies. Use the support matrix and model examples to confirm compatibility before enabling them. + + +## torch.compile + +FastVideo can `torch.compile` the DiT (transformer) for a substantial +end-to-end speedup. It is **off by default** and enabled per-run. + +### Enabling + +```python +generator = VideoGenerator.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + enable_torch_compile=True, +) +``` + +A complete A/B example (eager vs compiled, warmup excluded) is in +[`examples/inference/optimizations/torch_compile_example.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/optimizations/torch_compile_example.py). + +`fastvideo generate` is config-file driven; to enable `torch.compile` +from the CLI, set the relevant field in your run config and pass it via +`fastvideo generate --config run.yaml`. There is no top-level +`--enable-torch-compile` flag on the subcommand. + +Only DiT submodules that declare `_compile_conditions` are compiled +(most shipped models). The text encoder and VAE are not compiled by this +flag. + +### What to expect + +| Config | Effect | +|---|---| +| Wan2.1-T2V-1.3B, A100-80GB, 480×832×81f, 50 steps | end-to-end **259.7s → 198.1s (−23.7%)**; per-step **4.91 → 3.78 s/it** | + +The speedup is **configuration-dependent** — it varies with model, +resolution, step count, and GPU. Treat the number above as one measured +data point, not a guarantee; benchmark your own config (recipe below). + +There is a **one-time graph-build cost** on the first generation (tens of +seconds to minutes, model-dependent). It amortizes over subsequent +generations with the same input shapes. Always exclude the first +(warmup) generation when measuring steady-state latency — measuring the +warmup is the most common way to wrongly conclude "compile is slower". + +**Numerics.** Inductor's lowering is designed to preserve eager +semantics within floating-point tolerance, but per-model equivalence is +not asserted by any standing SSIM regression here — the SSIM tests in +[`fastvideo/tests/ssim/`](https://github.com/hao-ai-lab/FastVideo/tree/main/fastvideo/tests/ssim) +run with `enable_torch_compile` disabled. If you depend on compile +output staying close to eager (or your previous compiled run), run an +MS-SSIM gate on *your* config, especially when combining +`enable_torch_compile=True` with other numerics-affecting flags +(quantized attention backends, FP4, layerwise offload edge cases). + +### Known interactions + +- **Layerwise CPU offload** (`dit_layerwise_offload=True`, the default): + the offload hook previously caused an implicit graph break once per + transformer layer, fragmenting the compiled region. Addressed in + hao-ai-lab/FastVideo#1365 — keep that fix to get a clean compiled + region under the default offload path. +- **`mode="reduce-overhead"` / CUDA graphs**: not yet supported + end-to-end. The attention dispatch is an untraceable custom op and + still breaks the graph, which CUDA-graph trees cannot span. Use the + default inductor mode (shown above) until that is resolved. + +Extra `torch.compile` options are passed through `torch_compile_kwargs` +(a dict), accepted by `VideoGenerator.from_pretrained(...)` and by the +CLI as a JSON string via `--torch-compile-kwargs`. Example (currently +**not** recommended — see the CUDA-graphs caveat above): + +```python +VideoGenerator.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + enable_torch_compile=True, + torch_compile_kwargs={"mode": "reduce-overhead"}, # may error today +) +``` + +### Benchmarking torch.compile + +Same discipline as attention backends — same prompt, same seed, same +config; **discard the first generation** (graph build): + +```python +import time +from fastvideo import VideoGenerator + +gen = VideoGenerator.from_pretrained("your-model-id", enable_torch_compile=True) +req = {"prompt": "Your prompt", "sampling": {"seed": 1024}, + "output": {"save_video": False}} +gen.generate(req) # warmup: graph build, discard +t0 = time.perf_counter() +gen.generate(req) # measured: shapes reused +print(f"compiled steady-state: {time.perf_counter() - t0:.2f}s") +``` + +See `examples/inference/optimizations/torch_compile_example.py` for a +baseline-vs-compile A/B with the warmup correctly excluded. + ## Benchmarking different optimizations To benchmark backend performance, generate the same prompt with the same seed and compare end-to-end generation times: diff --git a/examples/inference/optimizations/torch_compile_example.py b/examples/inference/optimizations/torch_compile_example.py new file mode 100644 index 000000000..2a7cfbf13 --- /dev/null +++ b/examples/inference/optimizations/torch_compile_example.py @@ -0,0 +1,81 @@ +"""torch.compile A/B example for FastVideo. + +`enable_torch_compile=True` compiles the DiT submodules that declare +`_compile_conditions` for a substantial end-to-end speedup (e.g. +Wan2.1-T2V-1.3B on A100: ~-24% e2e). It is off by default. + +The first compiled generation pays a one-time graph-build cost; it +amortizes over later generations with the same input shapes. This +script does one un-measured warmup then a measured run so the reported +number is steady-state, not graph-build — measuring the warmup is the +most common way to wrongly conclude compile is slower. + +Usage: + # baseline (eager) + python torch_compile_example.py + # compiled + python torch_compile_example.py --compile +""" + +import argparse +import os +import time + +from fastvideo import VideoGenerator + +PROMPT = ( + "A high-definition video of a robotic arm welding a metal structure, " + "bright sparks and smoke, industrial setting." +) + + +def main() -> None: + parser = argparse.ArgumentParser(description="torch.compile A/B") + parser.add_argument("--model", default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + parser.add_argument("--compile", action="store_true", + help="Enable torch.compile for the DiT") + parser.add_argument("--num_gpus", type=int, default=1) + args = parser.parse_args() + + mode = "COMPILE" if args.compile else "BASELINE" + print(f"Mode: {mode} (enable_torch_compile={args.compile})") + + os.makedirs("video_samples", exist_ok=True) + + generator = VideoGenerator.from_pretrained( + args.model, + num_gpus=args.num_gpus, + enable_torch_compile=args.compile, + ) + + def _run(tag: str) -> float: + save = tag == "measured" + # Modern typed-request API (generate_video is deprecated). Same + # prompt/seed/shapes both runs so the compiled graph is reused. + request: dict = { + "prompt": PROMPT, + "sampling": {"seed": 1024}, + "output": {"save_video": save}, + } + if save: + request["output"]["output_path"] = ( + f"video_samples/torch_compile_{tag}.mp4") + t0 = time.perf_counter() + generator.generate(request) + return time.perf_counter() - t0 + + try: + # Warmup: pays the one-time graph build when --compile. Discarded. + w = _run("warmup") + print(f"warmup: {w:.2f}s " + f"({'incl. graph build' if args.compile else 'cold start'})") + + # Measured: steady state, compiled graph reused (same shapes/seed). + m = _run("measured") + print(f"=== {mode} steady-state e2e: {m:.2f}s ===") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main()