nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix]: guard automatic FA3 regional compile

Browse Source
This commit is contained in:
H1yori233
2026-09-01 19:40:22 -07:00
parent 8775bff28a
commit 5c37eaa60f
2 changed files with 72 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/loader/fsdp_load.py
+31 -8
View File
@@ -312,6 +312,7 @@ def maybe_load_fsdp_model(
else:
unsupported = _regional_compile_unsupported_reason(
init_params,
model=model,
training=True,
vsa_tile_size=inference_vsa_tile_size,
)
@@ -331,6 +332,7 @@ def maybe_load_fsdp_model(
# user kwargs needed (fullgraph + emulate_precision_casts injected).
unsupported = _regional_compile_unsupported_reason(
init_params,
model=model,
vsa_tile_size=inference_vsa_tile_size,
)
if unsupported is None:
@@ -361,6 +363,7 @@ def _strip_checkpoint_wrapper_prefix(name: str) -> str:
def _regional_compile_unsupported_reason(
init_params: dict[str, Any],
*,
model: nn.Module | None = None,
training: bool = False,
vsa_tile_size: int | None = None,
) -> str | None:
@@ -390,21 +393,41 @@ def _regional_compile_unsupported_reason(
"forwards out of compiled graphs via torch.compiler."
"disable, which fullgraph regional compile cannot trace; "
"this model stays eager")
config = init_params.get("config")
resolved = getattr(config, "_resolved_attention_backend", None)
resolved_name = getattr(resolved, "name", "")
if training and resolved_name == "FLASH_ATTN":
backend_names: set[str] = set()
if model is not None:
from fastvideo.attention.layer import DistributedAttention, LocalAttention
for submodule in model.modules():
if not isinstance(submodule, (DistributedAttention, LocalAttention)):
continue
backend = getattr(submodule, "backend", None)
backend_name = getattr(backend, "name", "")
if backend_name:
backend_names.add(backend_name)
# The component config stores the requested backend. ``None`` means
# platform automatic selection, so once a constructed model is available
# its attention instances are the source of truth. Keep the config fallback
# for callers that validate policy before or without model construction.
if not backend_names:
config = init_params.get("config")
resolved = getattr(config, "_resolved_attention_backend", None)
resolved_name = getattr(resolved, "name", "")
if resolved_name:
backend_names.add(resolved_name)
if training and "FLASH_ATTN" in backend_names:
try:
from fastvideo.attention.utils.flash_attn_default import fa_version
except Exception: # pragma: no cover - flash-attn stack not importable
pass
else:
if fa_version == "3":
return ("attention backend resolved to FLASH_ATTN with flash-attn 3, "
return ("model uses 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 training")
if resolved_name == "VIDEO_SPARSE_ATTN_H3":
if "VIDEO_SPARSE_ATTN_H3" in backend_names:
if training:
return ("VIDEO_SPARSE_ATTN_H3 regional compile is supported only for inference; "
"compiled training stays eager")
@@ -417,8 +440,8 @@ def _regional_compile_unsupported_reason(
if vsa_tile_size != 64:
return ("VIDEO_SPARSE_ATTN_H3 regional compile requires VSA_tile_size=64; "
f"got {vsa_tile_size!r}, so tile-256/CuTe VSA stays eager")
if resolved_name == "VIDEO_SPARSE_ATTN":
return (f"attention backend resolved to {resolved_name}, whose Triton "
if "VIDEO_SPARSE_ATTN" in backend_names:
return ("attention backend resolved to VIDEO_SPARSE_ATTN, whose Triton "
"kernels, sequence-parallel collectives, and sync metadata "
"guard graph-break (incompatible with fullgraph regional "
"compile); this model stays eager")
fastvideo/tests/inference/test_inference_regional_compile.py
+41
View File
@@ -38,6 +38,7 @@ from fastvideo.models.loader.fsdp_load import (
_prepare_model_for_compile,
_regional_compile_unsupported_reason,
)
from fastvideo.platforms import AttentionBackendEnum
def _init_params_for(backend_name: str | None) -> dict:
@@ -45,6 +46,15 @@ def _init_params_for(backend_name: str | None) -> dict:
return {"config": SimpleNamespace(_resolved_attention_backend=resolved)}
def _model_with_actual_backend(backend: AttentionBackendEnum) -> nn.Module:
attention = DistributedAttention.__new__(DistributedAttention)
nn.Module.__init__(attention)
attention.backend = backend
model = nn.Module()
model.add_module("attention", attention)
return model
def test_legacy_vsa_backend_degrades_to_eager(monkeypatch) -> None:
monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
backend_name = "VIDEO_SPARSE_ATTN"
@@ -268,6 +278,37 @@ def test_dense_flash_attention_inference_allows_compile(fa_version, monkeypatch)
assert _regional_compile_unsupported_reason(_init_params_for("FLASH_ATTN")) is None
def test_automatic_flash_attention_fa3_training_degrades_to_eager(monkeypatch) -> None:
monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
fake_module = ModuleType("fastvideo.attention.utils.flash_attn_default")
fake_module.fa_version = "3"
monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module)
model = _model_with_actual_backend(AttentionBackendEnum.FLASH_ATTN)
reason = _regional_compile_unsupported_reason(
_init_params_for(None),
model=model,
training=True,
)
assert reason is not None
assert "flash-attn 3" in reason
def test_actual_sdpa_backend_overrides_flash_request_for_fa3_guard(monkeypatch) -> None:
monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
fake_module = ModuleType("fastvideo.attention.utils.flash_attn_default")
fake_module.fa_version = "3"
monkeypatch.setitem(sys.modules, fake_module.__name__, fake_module)
model = _model_with_actual_backend(AttentionBackendEnum.TORCH_SDPA)
assert _regional_compile_unsupported_reason(
_init_params_for("FLASH_ATTN"),
model=model,
training=True,
) is None
def test_default_attention_dispatch_stays_compiler_disabled(monkeypatch) -> None:
monkeypatch.delenv("FASTVIDEO_DISABLE_ATTENTION_COMPILE", raising=False)
disabled_calls: list[str] = []