From 0431d4e8a9cbc845ef8cfa66720777ce2e96797e Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 21 Jul 2026 06:07:35 -0700 Subject: [PATCH] [perf]: enable regional compile for training --- examples/train/configs/example.yaml | 1 + fastvideo/attention/utils/flash_attn_cute.py | 173 +++++++++++++++--- fastvideo/models/dits/ltx2.py | 1 + fastvideo/models/loader/fsdp_load.py | 38 +++- .../test_flash_attn_cute_custom_op.py | 78 ++++++++ .../contract/test_ltx2_compile_conditions.py | 6 + fastvideo/tests/train/utils/test_config.py | 13 ++ .../train/utils/test_moduleloader_compile.py | 121 ++++++++++++ fastvideo/train/utils/config.py | 6 + fastvideo/train/utils/moduleloader.py | 2 +- fastvideo/train/utils/training_config.py | 1 + 11 files changed, 409 insertions(+), 31 deletions(-) create mode 100644 fastvideo/tests/train/utils/test_moduleloader_compile.py diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index 692677fa0..225529fb6 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -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) diff --git a/fastvideo/attention/utils/flash_attn_cute.py b/fastvideo/attention/utils/flash_attn_cute.py index ca3853992..b25eb0b1b 100644 --- a/fastvideo/attention/utils/flash_attn_cute.py +++ b/fastvideo/attention/utils/flash_attn_cute.py @@ -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, ) diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index 13797b131..184c791c1 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/fastvideo/models/dits/ltx2.py @@ -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. diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index e7045ff3d..7e9e7c3b3 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -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 diff --git a/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py b/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py index 2fc540937..8cb31c55a 100644 --- a/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py +++ b/fastvideo/tests/attention/test_flash_attn_cute_custom_op.py @@ -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) diff --git a/fastvideo/tests/contract/test_ltx2_compile_conditions.py b/fastvideo/tests/contract/test_ltx2_compile_conditions.py index 89d68cbd6..fe8dca0b5 100644 --- a/fastvideo/tests/contract/test_ltx2_compile_conditions.py +++ b/fastvideo/tests/contract/test_ltx2_compile_conditions.py @@ -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 ( diff --git a/fastvideo/tests/train/utils/test_config.py b/fastvideo/tests/train/utils/test_config.py index 77eb7100f..0bd178e3d 100644 --- a/fastvideo/tests/train/utils/test_config.py +++ b/fastvideo/tests/train/utils/test_config.py @@ -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() diff --git a/fastvideo/tests/train/utils/test_moduleloader_compile.py b/fastvideo/tests/train/utils/test_moduleloader_compile.py new file mode 100644 index 000000000..d76c192ad --- /dev/null +++ b/fastvideo/tests/train/utils/test_moduleloader_compile.py @@ -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"] diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index 7bf8f646f..c40896155 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -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, diff --git a/fastvideo/train/utils/moduleloader.py b/fastvideo/train/utils/moduleloader.py index b1aaf83ac..2be564398 100644 --- a/fastvideo/train/utils/moduleloader.py +++ b/fastvideo/train/utils/moduleloader.py @@ -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, ) diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index d47e6dc52..43f1c5b89 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -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)