diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index 7620fe6db..2b1f820de 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -345,10 +345,11 @@ training-side regional compile of loader wraps each `_compile_conditions` block in `torch.compile(fullgraph=True)` with inductor `options={"emulate_precision_casts": True}` right after the transformer -loads. Attention backends that cannot be traced end-to-end (VSA, -FLASH_ATTN on flash-attn 3, or the `FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` -escape hatch) degrade the transformer to eager with one warning instead of -failing mid-denoise. +loads. The ordinary compile path keeps its historical compiler-disabled +attention boundary by default; regional compile opts in only the compatible +attention instances owned by this transformer. VSA, FLASH_ATTN on flash-attn +3, and the explicit `FASTVIDEO_DISABLE_ATTENTION_COMPILE=1` escape hatch keep +the transformer eager with one warning instead of failing mid-denoise. ```python generator = VideoGenerator.from_pretrained( @@ -362,7 +363,12 @@ inductor options, and torch.compile forbids mode+options); it is independent of `enable_torch_compile`, and when both are set the regional compile wins for the DiT. -### What to expect +### What to expect from generic compile + +The Wan result below measures the existing generic +`enable_torch_compile=True` path. It is useful evidence that compile can help, +but it is **not** a benchmark or numerical gate for the stricter regional +fullgraph path above. | Config | Effect | |---|---| @@ -378,6 +384,16 @@ 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". +**Regional MiniMax-H3 accuracy caveat (job 2660).** On one GB200 at +768×1344×124, the native 50-point schedule ran exactly 49 transformer +forwards. After one warmup, three fixed-prompt/fixed-seed repeats averaged +**185.08s → 157.01s end to end** and **174.90s → 147.12s denoising**. Each +leg was independently pixel-deterministic, but compiled output did **not** +match eager: mean absolute pixel error **20.247/255**, PSNR **16.67 dB**, +mean SSIM **0.7108**, and mean MS-SSIM **0.6370** across 124 frames. Treat +regional MiniMax-H3 compile as an opt-in performance experiment, not an +eager-parity-safe mode. + **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 @@ -396,9 +412,13 @@ MS-SSIM gate on *your* config, especially when combining 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. + by regional compile because that path injects inductor `options`, and + PyTorch rejects `mode` together with `options`. The generic compile path + can accept `mode`, but CUDA-graph compatibility remains backend- and + shape-dependent. FA2/FA3 inference and FA4 expose traceable custom-op + boundaries; VSA and the FA3 grad-enabled path remain outside the regional + support envelope. Use the default inductor mode shown above unless your + exact configuration has its own gate. Extra `torch.compile` options are passed through `torch_compile_kwargs` (a dict), accepted by `VideoGenerator.from_pretrained(...)` and by the diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 125c77364..ac53492b7 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 import os +from functools import wraps import torch import torch.nn as nn @@ -18,21 +19,42 @@ from fastvideo.layers.rotary_embedding import _apply_rotary_emb def _attention_compile_disabled() -> bool: """Whether to keep attention ``forward`` out of the torch.compile graph. - Attention backends expose traceable custom-op boundaries, so compilation - is enabled by default. Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=1`` for - an explicit eager escape hatch when debugging a backend. + Defaults to ``True`` (the historical behavior: attention runs eager via + ``torch.compiler.disable``). Set ``FASTVIDEO_DISABLE_ATTENTION_COMPILE=0`` + to let attention instances constructed under that environment be traced. """ val = os.environ.get("FASTVIDEO_DISABLE_ATTENTION_COMPILE") if val is None: - return False + return True return val.strip().lower() not in ("0", "false", "no", "off", "") +def _attention_compile_explicitly_disabled() -> bool: + """Whether the environment explicitly requests the eager boundary. + + Regional compile can override the historical default for one loaded + transformer, but it must still honor an explicit debugging escape hatch. + """ + return "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in os.environ and _attention_compile_disabled() + + def _maybe_compiler_disable(fn): - """Apply ``torch.compiler.disable`` unless disabled via env var.""" - if _attention_compile_disabled(): - return torch.compiler.disable(fn) - return fn + """Defer the eager/traceable choice to each attention instance. + + A class-definition-time choice makes a process-wide default the only + option. The deferred wrapper keeps ordinary instances on the historical + eager boundary while allowing the regional loader to opt in only the + attention modules owned by the transformer it is compiling. + """ + disabled_fn = torch.compiler.disable(fn) + + @wraps(fn) + def _dispatch(self, *args, **kwargs): + if self._compile_forward_enabled: + return fn(self, *args, **kwargs) + return disabled_fn(self, *args, **kwargs) + + return _dispatch class DistributedAttention(nn.Module): @@ -77,6 +99,13 @@ class DistributedAttention(nn.Module): self.num_kv_heads = num_kv_heads self.backend = backend_name_to_enum(attn_backend.get_name()) self.dtype = dtype + # Preserve the historical compiler-disabled default. The regional + # inference loader may enable this one instance after validating the + # transformer's resolved backend; no process-global default changes. + self._compile_forward_enabled = not _attention_compile_disabled() + + def _set_compile_forward_enabled(self, enabled: bool) -> None: + self._compile_forward_enabled = enabled @_maybe_compiler_disable def forward( diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index db76bce1a..6a8608186 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -250,6 +250,8 @@ def maybe_load_fsdp_model( if callable(prepare_for_compile): logger.info("Running prepare_for_compile for %s", type(model).__name__) prepare_for_compile() + attention_count = _enable_regional_attention_compile(model) + logger.info("Enabled attention tracing for %d modules in %s", attention_count, type(model).__name__) _compile_model_regions(model, torch_compile_kwargs or {}) return model @@ -269,11 +271,11 @@ def _regional_compile_unsupported_reason(init_params: dict[str, Any]) -> str | N falls back to eager while compile-safe dense loads still compile. """ try: - from fastvideo.attention.layer import _attention_compile_disabled + from fastvideo.attention.layer import _attention_compile_explicitly_disabled except Exception: # pragma: no cover - attention stack not importable pass else: - if _attention_compile_disabled(): + if _attention_compile_explicitly_disabled(): # The escape hatch wraps attention forwards in # torch.compiler.disable, which is a hard dynamo error inside a # fullgraph region ("Skip inlining `torch.compiler.disable()`d @@ -301,10 +303,22 @@ def _regional_compile_unsupported_reason(init_params: dict[str, Any]) -> str | N return ("attention backend resolved to FLASH_ATTN with flash-attn 3, " "whose grad-enabled path graph-breaks (incompatible with " "fullgraph regional compile); use FA2, FA4 (FASTVIDEO_FA4=1), " - "or TORCH_SDPA for compiled runs") + "or TORCH_SDPA for compiled runs; this model stays eager") return None +def _enable_regional_attention_compile(model: nn.Module) -> int: + """Opt in distributed-attention instances owned by ``model`` only.""" + from fastvideo.attention.layer import DistributedAttention + + enabled_count = 0 + for submodule in model.modules(): + if isinstance(submodule, DistributedAttention): + submodule._set_compile_forward_enabled(True) + enabled_count += 1 + return enabled_count + + def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) -> int: """Compile repeated mathematical regions of a loaded model. @@ -321,9 +335,9 @@ def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) -> if "mode" in compile_kwargs: # torch.compile forbids passing both `mode` and `options`, and # regional compile always injects options (emulate_precision_casts) - # for bf16 numerics parity. Fail here with an actionable message - # instead of letting torch raise a mode/options conflict about an - # `options` key the user never wrote. + # to match the training-side regional-compile configuration. Fail here + # with an actionable message instead of letting torch raise a + # mode/options conflict about an `options` key the user never wrote. raise ValueError("Regional compile sets inductor options " "(emulate_precision_casts) and cannot be combined " "with torch_compile_kwargs['mode']. Remove 'mode' or " diff --git a/fastvideo/tests/inference/test_inference_regional_compile.py b/fastvideo/tests/inference/test_inference_regional_compile.py index 080c2b137..1f72e0dfd 100644 --- a/fastvideo/tests/inference/test_inference_regional_compile.py +++ b/fastvideo/tests/inference/test_inference_regional_compile.py @@ -9,23 +9,31 @@ must not drift from the #1718 training-port semantics: - ``_regional_compile_unsupported_reason``: VSA backends (and the attention eager escape hatch) degrade to eager with a reason instead of hard-failing fullgraph capture at the first denoising forward. +- attention forward dispatch: ordinary instances retain the historical + compiler-disabled boundary; regional compile opts in only the selected + model's instances. - ``_compile_model_regions``: exactly the ``_compile_conditions`` blocks are compiled, fullgraph=True plus inductor ``emulate_precision_casts`` are injected, and ``mode`` kwargs are rejected (torch.compile forbids mode+options). -CPU-safe: torch.compile is monkeypatched, no CUDA needed. +CPU-safe: the real capture check uses torch's eager recording backend; no CUDA +or inductor kernel build is needed. """ -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace import pytest import torch from torch import nn +from fastvideo.attention import layer as attention_layer +from fastvideo.attention.layer import DistributedAttention from fastvideo.models.loader import fsdp_load from fastvideo.models.loader.fsdp_load import ( _compile_model_regions, + _enable_regional_attention_compile, _regional_compile_unsupported_reason, ) @@ -57,6 +65,94 @@ def test_attention_compile_escape_hatch_degrades_to_eager(monkeypatch) -> None: assert "FASTVIDEO_DISABLE_ATTENTION_COMPILE" in reason +@pytest.mark.parametrize(("fa_version", "allowed"), [("2", True), ("3", False), ("4", True)]) +def test_flash_attention_3_degrades_to_eager(fa_version, allowed, monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + fake_module = ModuleType("fastvideo.attention.utils.flash_attn_default") + fake_module.fa_version = fa_version + monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module) + + reason = _regional_compile_unsupported_reason(_init_params_for("FLASH_ATTN")) + assert (reason is None) is allowed + if not allowed: + assert "flash-attn 3" in reason + assert "eager" in reason + + +def test_default_attention_dispatch_stays_compiler_disabled(monkeypatch) -> None: + monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False) + disabled_calls: list[str] = [] + + def _fake_disable(fn): + + def _disabled(self): + disabled_calls.append("disabled") + return fn(self) + + return _disabled + + monkeypatch.setattr(attention_layer.torch.compiler, "disable", _fake_disable) + + class _ToyAttention: + + def __init__(self) -> None: + self._compile_forward_enabled = not attention_layer._attention_compile_disabled() + + @attention_layer._maybe_compiler_disable + def forward(self) -> str: + return "forward" + + ordinary = _ToyAttention() + assert ordinary.forward() == "forward" + assert disabled_calls == ["disabled"] + + # Preserve the existing process-wide escape hatch for callers that opt in + # before constructing their attention modules. + monkeypatch.setenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", "0") + explicit = _ToyAttention() + assert explicit.forward() == "forward" + assert disabled_calls == ["disabled"] + + +class _BareDistributedAttention(DistributedAttention): + + def __init__(self) -> None: + nn.Module.__init__(self) + self._compile_forward_enabled = False + + +class _CompileProbe(nn.Module): + + def __init__(self, enabled: bool) -> None: + super().__init__() + self._compile_forward_enabled = enabled + + @attention_layer._maybe_compiler_disable + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + +def test_attention_instance_flag_controls_real_fullgraph_capture() -> None: + try: + ordinary = torch.compile(_CompileProbe(False), backend="eager", fullgraph=True) + with pytest.raises(torch._dynamo.exc.Unsupported, match="Skip calling"): + ordinary(torch.ones(2)) + + regional = torch.compile(_CompileProbe(True), backend="eager", fullgraph=True) + torch.testing.assert_close(regional(torch.ones(2)), torch.full((2,), 2.0)) + finally: + torch._dynamo.reset() + + +def test_regional_compile_enables_only_loaded_model_attention() -> None: + selected_model = nn.Sequential(_BareDistributedAttention(), _BareDistributedAttention()) + unrelated_model = nn.Sequential(_BareDistributedAttention()) + + assert _enable_regional_attention_compile(selected_model) == 2 + assert all(module._compile_forward_enabled for module in selected_model) + assert not unrelated_model[0]._compile_forward_enabled + + class _Block(nn.Module): def __init__(self) -> None: