[docs] Document enable_torch_compile (+ A/B example) (#1366)

2026-05-22 00:25:42 -07:00
parent 2dd5760291
commit 68179cd752
2 changed files with 183 additions and 0 deletions
@@ -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.

+<a id="torch-compile"></a>
+
+## 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:
@@ -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()