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

[perf]: enable regional compile for training

Browse Source
This commit is contained in:
SolitaryThinker
2026-07-21 06:07:35 -07:00
parent cfc110cb6b
commit 0431d4e8a9
11 changed files with 409 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/train/configs/example.yaml
+1
View File
@@ -150,6 +150,7 @@ training:
precondition_outputs: false # default: false
moba_config: {} # default: {}
enable_gradient_checkpointing_type: full # default: null ("full" or null)
enable_torch_compile: false # default: false; regionally compile matched DiT blocks before FSDP
# --- training top-level [TYPED] ---
dit_precision: fp32 # default: "fp32" (master weight precision)
fastvideo/attention/utils/flash_attn_cute.py
+149 -24
View File
@@ -9,6 +9,12 @@ from fastvideo.platforms import current_platform
logger = init_logger(__name__)
def _empty_like_fa4_backward_input(x: torch.Tensor) -> torch.Tensor:
"""Mirror FA4's last-dimension contiguity normalization in fake kernels."""
return torch.empty_like(x.contiguous() if x.stride(-1) != 1 else x)
if torch.cuda.is_available():
try:
from flash_attn.cute.interface import _flash_attn_bwd, _flash_attn_fwd
@@ -108,6 +114,7 @@ def _flash_attn_cute_forward(
softcap=0.0,
num_splits=1,
pack_gqa=None,
return_lse=True,
)[:2]
return out, lse
@@ -132,38 +139,92 @@ def _flash_attn_cute_setup_context(ctx: torch.autograd.function.FunctionCtx, inp
q, k, v, softmax_scale, causal, deterministic = inputs
out, lse = output
ctx.save_for_backward(q, k, v, out, lse)
ctx.mark_non_differentiable(lse)
ctx.softmax_scale = softmax_scale
ctx.causal = causal
ctx.deterministic = deterministic
def _flash_attn_cute_backward(
ctx: torch.autograd.function.FunctionCtx,
# ``register_autograd`` backward callbacks must be AOT-traceable. CuTe's
# Python backward skips kernel launches under fake mode, so expose the real
# dense and varlen backward kernels as opaque custom ops instead of tracing
# their Python wrappers into allocation-only graphs.
@torch.library.custom_op(
"fastvideo::_flash_attn_cute_backward",
mutates_args=(),
device_types="cuda",
)
def _flash_attn_cute_backward_op(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor,
grad_out: torch.Tensor,
grad_lse: torch.Tensor | None,
):
del grad_lse
q, k, v, out, lse = ctx.saved_tensors
dq, dk, dv = _flash_attn_bwd(
lse: torch.Tensor,
softmax_scale: float | None,
causal: bool,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return _flash_attn_bwd(
q,
k,
v,
out,
grad_out,
lse,
softmax_scale=ctx.softmax_scale,
causal=ctx.causal,
softmax_scale=softmax_scale,
causal=causal,
softcap=0.0,
window_size_left=None,
window_size_right=None,
deterministic=ctx.deterministic,
deterministic=deterministic,
)
@torch.library.register_fake("fastvideo::_flash_attn_cute_backward")
def _flash_attn_cute_backward_fake(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor,
grad_out: torch.Tensor,
lse: torch.Tensor,
softmax_scale: float | None,
causal: bool,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
del out, grad_out, lse, softmax_scale, causal, deterministic
return (
_empty_like_fa4_backward_input(q),
_empty_like_fa4_backward_input(k),
_empty_like_fa4_backward_input(v),
)
def _flash_attn_cute_autograd_backward(
ctx: torch.autograd.function.FunctionCtx,
grad_out: torch.Tensor,
grad_lse: torch.Tensor | None,
):
del grad_lse
q, k, v, out, lse = ctx.saved_tensors
dq, dk, dv = torch.ops.fastvideo._flash_attn_cute_backward(
q,
k,
v,
out,
grad_out,
lse,
ctx.softmax_scale,
ctx.causal,
ctx.deterministic,
)
return dq, dk, dv, None, None, None
torch.library.register_autograd(
"fastvideo::_flash_attn_cute_forward",
_flash_attn_cute_backward,
_flash_attn_cute_autograd_backward,
setup_context=_flash_attn_cute_setup_context,
)
@@ -200,6 +261,7 @@ def _flash_attn_cute_varlen_forward(
softcap=0.0,
num_splits=1,
pack_gqa=None,
return_lse=True,
)[:2]
return out, lse
@@ -241,6 +303,7 @@ def _flash_attn_cute_varlen_setup_context(ctx: torch.autograd.function.FunctionC
) = inputs
out, lse = output
ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)
ctx.mark_non_differentiable(lse)
ctx.max_seqlen_q = max_seqlen_q
ctx.max_seqlen_k = max_seqlen_k
ctx.softmax_scale = softmax_scale
@@ -248,37 +311,99 @@ def _flash_attn_cute_varlen_setup_context(ctx: torch.autograd.function.FunctionC
ctx.deterministic = deterministic
def _flash_attn_cute_varlen_backward(
ctx: torch.autograd.function.FunctionCtx,
@torch.library.custom_op(
"fastvideo::_flash_attn_cute_varlen_backward",
mutates_args=(),
device_types="cuda",
)
def _flash_attn_cute_varlen_backward_op(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor,
grad_out: torch.Tensor,
grad_lse: torch.Tensor | None,
):
del grad_lse
q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
dq, dk, dv = _flash_attn_bwd(
lse: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_q: int,
max_seqlen_k: int,
softmax_scale: float | None,
causal: bool,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
return _flash_attn_bwd(
q,
k,
v,
out,
grad_out,
lse,
softmax_scale=ctx.softmax_scale,
causal=ctx.causal,
softmax_scale=softmax_scale,
causal=causal,
softcap=0.0,
window_size_left=None,
window_size_right=None,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=ctx.max_seqlen_q,
max_seqlen_k=ctx.max_seqlen_k,
deterministic=ctx.deterministic,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
deterministic=deterministic,
)
@torch.library.register_fake("fastvideo::_flash_attn_cute_varlen_backward")
def _flash_attn_cute_varlen_backward_fake(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor,
grad_out: torch.Tensor,
lse: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
max_seqlen_q: int,
max_seqlen_k: int,
softmax_scale: float | None,
causal: bool,
deterministic: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
del out, grad_out, lse, cu_seqlens_q, cu_seqlens_k
del max_seqlen_q, max_seqlen_k, softmax_scale, causal, deterministic
return (
_empty_like_fa4_backward_input(q),
_empty_like_fa4_backward_input(k),
_empty_like_fa4_backward_input(v),
)
def _flash_attn_cute_varlen_autograd_backward(
ctx: torch.autograd.function.FunctionCtx,
grad_out: torch.Tensor,
grad_lse: torch.Tensor | None,
):
del grad_lse
q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
dq, dk, dv = torch.ops.fastvideo._flash_attn_cute_varlen_backward(
q,
k,
v,
out,
grad_out,
lse,
cu_seqlens_q,
cu_seqlens_k,
ctx.max_seqlen_q,
ctx.max_seqlen_k,
ctx.softmax_scale,
ctx.causal,
ctx.deterministic,
)
return dq, dk, dv, None, None, None, None, None, None, None
torch.library.register_autograd(
"fastvideo::_flash_attn_cute_varlen_forward",
_flash_attn_cute_varlen_backward,
_flash_attn_cute_varlen_autograd_backward,
setup_context=_flash_attn_cute_varlen_setup_context,
)
fastvideo/models/dits/ltx2.py
+1
View File
@@ -1955,6 +1955,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
self.norm_eps = norm_eps
@torch.compiler.disable
def _register_fsdp_backward_hooks_on_output(self, vx, ax):
"""Register backward hooks on output tensors to trigger FSDP2 unshard.
fastvideo/models/loader/fsdp_load.py
+32 -6
View File
@@ -26,6 +26,34 @@ from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available
logger = init_logger(__name__)
def _compile_matched_submodule_forwards(
model: nn.Module,
torch_compile_kwargs: dict[str, Any] | None = None,
) -> int:
"""Compile only forwards selected by ``model._compile_conditions``."""
compile_conditions = getattr(model, "_compile_conditions", None)
if not compile_conditions:
raise ValueError("Regional torch.compile requires model._compile_conditions; "
"refusing to compile the whole FSDP model.")
compile_kwargs = dict(torch_compile_kwargs or {})
if compile_kwargs.get("fullgraph") is True:
raise ValueError("Regional FSDP training compile requires "
"fullgraph=False so runtime-only hooks can remain graph breaks.")
compile_kwargs["fullgraph"] = False
compiled_count = 0
for name, submodule in model.named_modules():
if name and any(condition(name, submodule) for condition in compile_conditions):
submodule.forward = torch.compile(submodule.forward, **compile_kwargs)
compiled_count += 1
if compiled_count == 0:
raise ValueError("model._compile_conditions matched no submodules; "
"refusing to compile the whole FSDP model.")
return compiled_count
def _maybe_quantize_model(model: nn.Module) -> None:
"""Quantize NVFP4- or FP8-tagged linear layers in-place after weights are loaded.
@@ -151,6 +179,10 @@ def maybe_load_fsdp_model(
use_fsdp = False
logger.info("Disabling FSDP for MPS platform as it's not compatible")
if enable_torch_compile and training_mode:
compiled_count = _compile_matched_submodule_forwards(model, torch_compile_kwargs)
logger.info("Enabled regional torch.compile for %d training submodules before FSDP sharding", compiled_count)
if use_fsdp:
pin_cpu_memory = pin_cpu_memory and is_pin_memory_available()
world_size = hsdp_replicate_dim * hsdp_shard_dim
@@ -209,12 +241,6 @@ def maybe_load_fsdp_model(
# are present (lazy imports inside the helper).
_maybe_quantize_model(model)
compile_in_loader = enable_torch_compile and training_mode
if compile_in_loader:
compile_kwargs = torch_compile_kwargs or {}
logger.info("Enabling torch.compile for FSDP training module with kwargs=%s", compile_kwargs)
model = torch.compile(model, **compile_kwargs)
logger.info("torch.compile enabled for %s", type(model).__name__)
return model
fastvideo/tests/attention/test_flash_attn_cute_custom_op.py
+78
View File
@@ -56,6 +56,14 @@ def _extract_out(x):
return x
def _strided_input_pair(shape: tuple[int, ...], *, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
source = torch.randn(*shape, 2, device="cuda", dtype=dtype)
return (
source[..., 0].detach().requires_grad_(True),
source.clone()[..., 0].detach().requires_grad_(True),
)
def test_capability_gate_is_lazy_and_device_keyed(monkeypatch) -> None:
if not torch.cuda.is_available():
pytest.skip("CUDA is required to import flash_attn.cute")
@@ -185,6 +193,37 @@ def test_flash_attn_func_parity_forward_backward(
_assert_close(dv_test, dv_ref, dtype=dtype, is_grad=True)
def test_flash_attn_func_torch_compile_forward_backward_parity(flash_attn_impls):
custom_flash_attn_func, _, _, _ = flash_attn_impls
dtype = torch.float16
shape = (1, 32, 2, 64)
torch.manual_seed(2)
q_eager, q_compiled = _strided_input_pair(shape, dtype=dtype)
k_eager, k_compiled = _strided_input_pair(shape, dtype=dtype)
v_eager, v_compiled = _strided_input_pair(shape, dtype=dtype)
def attention(q, k, v):
return custom_flash_attn_func(
q,
k,
v,
dropout_p=0.0,
softmax_scale=None,
causal=False,
deterministic=False,
)
out_eager = attention(q_eager, k_eager, v_eager)
out_compiled = torch.compile(attention, fullgraph=True)(q_compiled, k_compiled, v_compiled)
_assert_close(out_compiled, out_eager, dtype=dtype)
dout = torch.randn_like(out_eager)
eager_grads = torch.autograd.grad((out_eager * dout).sum(), (q_eager, k_eager, v_eager))
compiled_grads = torch.autograd.grad((out_compiled * dout).sum(), (q_compiled, k_compiled, v_compiled))
for compiled_grad, eager_grad in zip(compiled_grads, eager_grads):
_assert_close(compiled_grad, eager_grad, dtype=dtype, is_grad=True)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("causal", [False, True])
def test_flash_attn_varlen_func_parity_forward_backward(
@@ -260,3 +299,42 @@ def test_flash_attn_varlen_func_parity_forward_backward(
_assert_close(dq_test, dq_ref, dtype=dtype, is_grad=True)
_assert_close(dk_test, dk_ref, dtype=dtype, is_grad=True)
_assert_close(dv_test, dv_ref, dtype=dtype, is_grad=True)
def test_flash_attn_varlen_func_torch_compile_forward_backward_parity(flash_attn_impls):
_, custom_flash_attn_varlen_func, _, _ = flash_attn_impls
dtype = torch.float16
seqlens = (16, 8)
total = sum(seqlens)
max_seqlen = max(seqlens)
cu_seqlens = torch.tensor((0, seqlens[0], total), device="cuda", dtype=torch.int32)
shape = (total, 2, 64)
torch.manual_seed(3)
q_eager, q_compiled = _strided_input_pair(shape, dtype=dtype)
k_eager, k_compiled = _strided_input_pair(shape, dtype=dtype)
v_eager, v_compiled = _strided_input_pair(shape, dtype=dtype)
def attention(q, k, v, cu):
return custom_flash_attn_varlen_func(
q,
k,
v,
cu,
cu,
max_seqlen,
max_seqlen,
dropout_p=0.0,
softmax_scale=None,
causal=False,
deterministic=False,
)
out_eager = attention(q_eager, k_eager, v_eager, cu_seqlens)
out_compiled = torch.compile(attention, fullgraph=True)(q_compiled, k_compiled, v_compiled, cu_seqlens)
_assert_close(out_compiled, out_eager, dtype=dtype)
dout = torch.randn_like(out_eager)
eager_grads = torch.autograd.grad((out_eager * dout).sum(), (q_eager, k_eager, v_eager))
compiled_grads = torch.autograd.grad((out_compiled * dout).sum(), (q_compiled, k_compiled, v_compiled))
for compiled_grad, eager_grad in zip(compiled_grads, eager_grads):
_assert_close(compiled_grad, eager_grad, dtype=dtype, is_grad=True)
fastvideo/tests/contract/test_ltx2_compile_conditions.py
+6
View File
@@ -36,6 +36,12 @@ def test_is_ltx2_blocks_matches_only_top_level_blocks():
assert not is_ltx2_blocks("audio_transformer_blocks.0", None)
def test_fsdp_output_hook_registration_stays_outside_compiled_graph():
register_hooks = ltx2.BasicAVTransformerBlock._register_fsdp_backward_hooks_on_output
assert getattr(register_hooks, "_torchdynamo_disable", False) is True
def test_freq_grid_cache_is_bypassed_while_compiling():
args = (10000.0, 3, 96)
for generator, cached_generator in (
fastvideo/tests/train/utils/test_config.py
+13
View File
@@ -88,6 +88,7 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None:
assert t.model.weighting_scheme == "uniform"
assert t.model.precondition_outputs is False
assert t.model.moba_config == {}
assert t.model.enable_torch_compile is False
assert t.dit_precision == "fp32"
assert t.vsa_sparsity == 0.0
@@ -148,6 +149,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
"logit_mean": 0.5,
"logit_std": 1.5,
"precondition_outputs": True,
"enable_torch_compile": True,
},
"dit_precision": "bf16",
}
@@ -183,6 +185,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
assert t.vsa_sparsity == pytest.approx(0.5)
assert t.model.weighting_scheme == "logit_normal"
assert t.model.precondition_outputs is True
assert t.model.enable_torch_compile is True
assert t.dit_precision == "bf16"
@@ -301,6 +304,7 @@ def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None:
"--training.optimizer.fused=true",
"--training.distributed.pin_cpu_memory=true",
"--training.distributed.reshard_after_forward=false",
"--training.model.enable_torch_compile=true",
"--training.tracker.project_name=overridden",
]
cfg = load_run_config(path, overrides=overrides)
@@ -310,6 +314,7 @@ def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None:
assert cfg.training.optimizer.fused is True
assert cfg.training.distributed.pin_cpu_memory is True
assert cfg.training.distributed.reshard_after_forward is False
assert cfg.training.model.enable_torch_compile is True
assert cfg.training.tracker.project_name == "overridden"
@@ -345,6 +350,14 @@ def test_reshard_after_forward_rejects_non_bool(tmp_path: Path) -> None:
load_run_config(_write_yaml(tmp_path, data))
def test_model_torch_compile_rejects_non_bool(tmp_path: Path) -> None:
data = _minimal_yaml()
data["training"] = {"model": {"enable_torch_compile": "true"}}
with pytest.raises(ValueError, match="training.model.enable_torch_compile must be a bool"):
load_run_config(_write_yaml(tmp_path, data))
def test_overrides_create_intermediate_keys(tmp_path: Path) -> None:
"""Overrides into a nested key absent from YAML should still apply."""
data = _minimal_yaml()
fastvideo/tests/train/utils/test_moduleloader_compile.py
+121
View File
@@ -0,0 +1,121 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import torch
import pytest
from fastvideo.models.loader import fsdp_load
from fastvideo.train.utils import moduleloader
from fastvideo.train.utils.training_config import (
DistributedConfig,
ModelTrainingConfig,
TrainingConfig,
)
class _RepeatedModel(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.blocks = torch.nn.ModuleList([torch.nn.Linear(1, 1), torch.nn.Linear(1, 1)])
self.output = torch.nn.Linear(1, 1)
self._compile_conditions = [lambda name, _module: name in {"blocks.0", "blocks.1"}]
self._fsdp_shard_conditions = self._compile_conditions
self.param_names_mapping = {}
def test_make_training_args_propagates_regional_compile() -> None:
config = TrainingConfig(
distributed=DistributedConfig(hsdp_shard_dim=1),
model=ModelTrainingConfig(enable_torch_compile=True),
)
args = moduleloader._make_training_args(config, model_path="fake/model")
assert args.enable_torch_compile is True
def test_compile_matched_submodule_forwards_is_regional_and_allows_graph_breaks(monkeypatch) -> None:
model = _RepeatedModel()
output_forward = model.output.forward
calls: list[tuple[object, dict[str, object]]] = []
def fake_compile(forward, **kwargs):
calls.append((forward.__self__, kwargs))
return forward
monkeypatch.setattr(fsdp_load.torch, "compile", fake_compile)
count = fsdp_load._compile_matched_submodule_forwards(model, {"backend": "eager"})
assert count == 2
assert [module for module, _ in calls] == list(model.blocks)
assert [kwargs for _, kwargs in calls] == [
{
"backend": "eager",
"fullgraph": False
},
{
"backend": "eager",
"fullgraph": False
},
]
assert model.output.forward == output_forward
@pytest.mark.parametrize(
"conditions",
[[], [lambda _name, _module: False]],
)
def test_compile_matched_submodule_forwards_rejects_missing_matches(monkeypatch, conditions) -> None:
model = _RepeatedModel()
model._compile_conditions = conditions
monkeypatch.setattr(fsdp_load.torch, "compile", lambda forward, **kwargs: forward)
with pytest.raises(ValueError, match="refusing to compile the whole FSDP model"):
fsdp_load._compile_matched_submodule_forwards(model)
def test_compile_matched_submodule_forwards_rejects_fullgraph() -> None:
with pytest.raises(ValueError, match="requires fullgraph=False"):
fsdp_load._compile_matched_submodule_forwards(_RepeatedModel(), {"fullgraph": True})
def test_fsdp_loader_compiles_regions_before_sharding(monkeypatch) -> None:
events: list[str] = []
def fake_compile(forward, **_kwargs):
events.append("compile")
return forward
def fake_shard_model(model, **_kwargs):
del model
events.append("shard")
def fake_load_state(model, _iterator, device, _dtype, **_kwargs):
model.to_empty(device=device)
monkeypatch.setattr(fsdp_load.torch, "compile", fake_compile)
monkeypatch.setattr(fsdp_load, "init_device_mesh", lambda *_args, **_kwargs: object())
monkeypatch.setattr(fsdp_load, "shard_model", fake_shard_model)
monkeypatch.setattr(fsdp_load, "safetensors_weights_iterator", lambda *_args, **_kwargs: iter(()))
monkeypatch.setattr(fsdp_load, "get_param_names_mapping", lambda _mapping: None)
monkeypatch.setattr(fsdp_load, "load_model_from_full_model_state_dict", fake_load_state)
monkeypatch.setattr(fsdp_load, "_maybe_quantize_model", lambda _model: None)
monkeypatch.setattr("fastvideo.platforms.current_platform.is_mps", lambda: False)
fsdp_load.maybe_load_fsdp_model(
model_cls=_RepeatedModel,
init_params={},
weight_dir_list=[],
device=torch.device("cpu"),
hsdp_replicate_dim=1,
hsdp_shard_dim=1,
default_dtype=torch.float32,
param_dtype=torch.float32,
reduce_dtype=torch.float32,
enable_torch_compile=True,
)
assert events == ["compile", "compile", "shard"]
fastvideo/train/utils/config.py
+6
View File
@@ -425,6 +425,12 @@ def _build_training_config(
precondition_outputs=bool(m.get("precondition_outputs", False)),
moba_config=dict(m.get("moba_config", {}) or {}),
enable_gradient_checkpointing_type=(m.get("enable_gradient_checkpointing_type")),
enable_torch_compile=require_bool(
m,
"enable_torch_compile",
default=False,
where="training.model.enable_torch_compile",
),
),
pipeline_config=pipeline_config,
model_path=model_path,
fastvideo/train/utils/moduleloader.py
+1 -1
View File
@@ -60,7 +60,7 @@ def _make_training_args(
text_encoder_cpu_offload=False,
image_encoder_cpu_offload=False,
use_fsdp_inference=False,
enable_torch_compile=False,
enable_torch_compile=tc.model.enable_torch_compile,
)
fastvideo/train/utils/training_config.py
+1
View File
@@ -79,6 +79,7 @@ class ModelTrainingConfig:
precondition_outputs: bool = False
moba_config: dict = field(default_factory=dict)
enable_gradient_checkpointing_type: str | None = None
enable_torch_compile: bool = False
@dataclass(slots=True)