From 792ef0d19454d5ad7aac2bdd2c4479469c8f0b2d Mon Sep 17 00:00:00 2001 From: root Date: Wed, 19 Aug 2026 07:01:01 +0000 Subject: [PATCH] [bugfix]: fix compiled Wan modulation gradients --- docs/training/train_infra.md | 13 +- fastvideo/models/dits/wanvideo.py | 143 +++++++++++++++++- fastvideo/models/loader/fsdp_load.py | 31 ++-- .../models/test_wan_modulation_custom_op.py | 43 ++++++ .../tests/train/utils/test_torch_compile.py | 36 ++++- 5 files changed, 228 insertions(+), 38 deletions(-) create mode 100644 fastvideo/tests/models/test_wan_modulation_custom_op.py diff --git a/docs/training/train_infra.md b/docs/training/train_infra.md index 7aa03325e..9a863fb42 100644 --- a/docs/training/train_infra.md +++ b/docs/training/train_infra.md @@ -189,11 +189,14 @@ block is then compiled inside its eager checkpoint wrapper. The checkpoint control logic, FSDP hooks, and the model-level block loop all remain outside the compiled region; checkpoint state-dict names are preserved. -Only no-grad block calls use the compiled path. Gradient-bearing student and -critic calls remain eager because the current Inductor backward is not -numerically reliable for the FA4 + FSDP + activation-checkpointed Wan stack. -This still accelerates DMD2 simulated rollouts and teacher evaluation without -invoking the affected compiled backward path. +Both forward and backward block calls use the compiled path. Wan modulation is +kept behind a small opaque operator boundary: this ensures all six modulation +gradient slices are materialized before they are returned to FSDP instead of +being fused into an incomplete Inductor output buffer. Eager execution keeps +the original native PyTorch path. Regional training compile also enables +Inductor's `emulate_precision_casts` option by default so fused BF16 operators +retain eager's intermediate rounding points. It can be overridden explicitly +under `torch_compile_kwargs.options`. Additional supported `torch.compile` options can be provided under `training.model.torch_compile_kwargs`; `fullgraph: false` is rejected because diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index 4c971ca9e..91d2ccec6 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -30,6 +30,129 @@ from fastvideo.distributed.parallel_state import get_sp_world_size logger = init_logger(__name__) +@torch.library.custom_op( + "fastvideo::_wan_modulation_forward", + mutates_args=(), + device_types="cuda", +) +def _wan_modulation_forward( + scale_shift_table: torch.Tensor, + temb: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Keep Wan modulation gradients outside Inductor's fused reductions. + + The six cloned outputs are intentional: custom operators may not return + aliased views, and keeping the outputs distinct prevents AOTAutograd from + reassembling their gradients inside the compiled block backward. + """ + if temb.dim() == 4: + modulation = scale_shift_table.unsqueeze(0).float() + temb.float() + chunk_dim = 2 + else: + modulation = scale_shift_table.float() + temb.float() + chunk_dim = 1 + return tuple(chunk.clone() for chunk in modulation.chunk(6, dim=chunk_dim)) # type: ignore[return-value] + + +@torch.library.register_fake("fastvideo::_wan_modulation_forward") +def _wan_modulation_forward_fake( + scale_shift_table: torch.Tensor, + temb: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if temb.dim() == 4: + shape = (*temb.shape[:2], 1, temb.shape[-1]) + else: + shape = (temb.shape[0], 1, temb.shape[-1]) + return tuple(temb.new_empty(shape, dtype=torch.float32) for _ in range(6)) # type: ignore[return-value] + + +@torch.library.custom_op( + "fastvideo::_wan_modulation_backward", + mutates_args=(), + device_types="cuda", +) +def _wan_modulation_backward_op( + grad_shift_msa: torch.Tensor, + grad_scale_msa: torch.Tensor, + grad_gate_msa: torch.Tensor, + grad_c_shift_msa: torch.Tensor, + grad_c_scale_msa: torch.Tensor, + grad_c_gate_msa: torch.Tensor, + scale_shift_table: torch.Tensor, + temb: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + chunk_dim = 2 if temb.dim() == 4 else 1 + chunk_shape = ((*temb.shape[:2], 1, temb.shape[-1]) + if temb.dim() == 4 else (temb.shape[0], 1, temb.shape[-1])) + grad_modulation = torch.cat( + tuple(grad.reshape(chunk_shape) for grad in ( + grad_shift_msa, + grad_scale_msa, + grad_gate_msa, + grad_c_shift_msa, + grad_c_scale_msa, + grad_c_gate_msa, + )), + dim=chunk_dim, + ) + grad_scale_shift = grad_modulation.sum_to_size(scale_shift_table.shape) + return ( + grad_scale_shift.to(scale_shift_table.dtype).clone(), + grad_modulation.to(temb.dtype).clone(), + ) + + +@torch.library.register_fake("fastvideo::_wan_modulation_backward") +def _wan_modulation_backward_fake( + grad_shift_msa: torch.Tensor, + grad_scale_msa: torch.Tensor, + grad_gate_msa: torch.Tensor, + grad_c_shift_msa: torch.Tensor, + grad_c_scale_msa: torch.Tensor, + grad_c_gate_msa: torch.Tensor, + scale_shift_table: torch.Tensor, + temb: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + del grad_shift_msa, grad_scale_msa, grad_gate_msa + del grad_c_shift_msa, grad_c_scale_msa, grad_c_gate_msa + return torch.empty_like(scale_shift_table), torch.empty_like(temb) + + +def _wan_modulation_setup_context(ctx: torch.autograd.function.FunctionCtx, inputs, output) -> None: + del output + scale_shift_table, temb = inputs + ctx.save_for_backward(scale_shift_table, temb) + + +def _wan_modulation_backward( + ctx: torch.autograd.function.FunctionCtx, + grad_shift_msa: torch.Tensor, + grad_scale_msa: torch.Tensor, + grad_gate_msa: torch.Tensor, + grad_c_shift_msa: torch.Tensor, + grad_c_scale_msa: torch.Tensor, + grad_c_gate_msa: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + scale_shift_table, temb = ctx.saved_tensors + return torch.ops.fastvideo._wan_modulation_backward( + grad_shift_msa, + grad_scale_msa, + grad_gate_msa, + grad_c_shift_msa, + grad_c_scale_msa, + grad_c_gate_msa, + scale_shift_table, + temb, + ) + + +torch.library.register_autograd( + "fastvideo::_wan_modulation_forward", + _wan_modulation_backward, + setup_context=_wan_modulation_setup_context, +) + + class WanImageEmbedding(torch.nn.Module): def __init__(self, in_features: int, out_features: int): @@ -372,7 +495,17 @@ class WanTransformerBlock(nn.Module): orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 - if temb.dim() == 4: + if hidden_states.is_cuda and torch.compiler.is_compiling(): + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + torch.ops.fastvideo._wan_modulation_forward(self.scale_shift_table, temb)) + if temb.dim() == 4: + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + elif temb.dim() == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) @@ -530,8 +663,12 @@ class WanTransformerBlock_VSA(nn.Module): bs, seq_length, _ = hidden_states.shape orig_dtype = hidden_states.dtype # assert orig_dtype != torch.float32 - e = self.scale_shift_table + temb.float() - shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) + if hidden_states.is_cuda and torch.compiler.is_compiling(): + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + torch.ops.fastvideo._wan_modulation_forward(self.scale_shift_table, temb)) + else: + e = self.scale_shift_table + temb.float() + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=1) assert shift_msa.dtype == torch.float32 # 1. Self-attention diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 1c097bf46..f39da6dbd 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -9,7 +9,6 @@ import os import contextlib import re from collections.abc import Callable, Generator -from functools import wraps from itertools import chain from typing import Any @@ -382,15 +381,11 @@ def _enable_regional_attention_compile(model: nn.Module) -> int: def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) -> int: - """Compile no-grad calls to repeated regions after FSDP setup. + """Compile repeated mathematical regions after FSDP setup. Only the selected module ``forward`` is replaced. This keeps activation checkpoint wrappers structurally transparent while FSDP pre/post hooks - execute outside the compiled region. Gradient-bearing calls deliberately - stay eager: Inductor's compiled backward is not numerically reliable for - the current FA4 + FSDP + checkpointed Wan block stack, while DMD2's - simulated rollouts and teacher calls still provide substantial no-grad - regions to optimize. + execute outside the compiled region. """ compile_conditions = getattr(model, "_compile_conditions", None) if not compile_conditions: @@ -399,33 +394,25 @@ def _compile_model_regions(model: nn.Module, compile_kwargs: dict[str, Any]) -> if compile_kwargs.get("fullgraph", True) is not True: raise ValueError("Regional training compile requires fullgraph=True") kwargs = {**compile_kwargs, "fullgraph": True} + options = {"emulate_precision_casts": True} + options.update(kwargs.get("options") or {}) + kwargs["options"] = options compiled_count = 0 for name, submodule in list(model.named_modules()): if not name: continue if any(condition(name, submodule) for condition in compile_conditions): # Activation checkpoint wrappers are control-flow boundaries, not - # mathematical regions. Compiling the wrapper asks AOTAutograd to - # trace checkpoint's saved-tensor/recompute machinery and has - # produced invalid training gradients. Keep that machinery eager - # and compile only no-grad calls to the repeated block it owns. + # mathematical regions. Keep their saved-tensor/recompute logic + # eager and compile only the repeated block they own. compile_target = getattr(submodule, "_checkpoint_wrapped_module", submodule) - eager_forward = compile_target.forward - compiled_forward = torch.compile(eager_forward, **kwargs) - - @wraps(eager_forward) - def forward(*args, _eager=eager_forward, _compiled=compiled_forward, **forward_kwargs): - if torch.is_grad_enabled(): - return _eager(*args, **forward_kwargs) - return _compiled(*args, **forward_kwargs) - - compile_target.forward = forward + compile_target.forward = torch.compile(compile_target.forward, **kwargs) compiled_count += 1 if compiled_count == 0: raise ValueError(f"No submodules in {type(model).__name__} matched _compile_conditions") logger.info( - "Enabled regional no-grad torch.compile for %d submodules in %s after FSDP setup with kwargs=%s", + "Enabled regional torch.compile for %d submodules in %s after FSDP setup with kwargs=%s", compiled_count, type(model).__name__, kwargs, diff --git a/fastvideo/tests/models/test_wan_modulation_custom_op.py b/fastvideo/tests/models/test_wan_modulation_custom_op.py new file mode 100644 index 000000000..77ac68bb4 --- /dev/null +++ b/fastvideo/tests/models/test_wan_modulation_custom_op.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compiled Wan modulation forward/backward regression tests.""" + +from __future__ import annotations + +import pytest +import torch + +# Import registers the model-local custom operators. +import fastvideo.models.dits.wanvideo # noqa: F401 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("temb_shape", [(2, 6, 32), (2, 3, 6, 32)]) +def test_wan_modulation_fullgraph_backward_matches_eager(temb_shape: tuple[int, ...]) -> None: + torch.manual_seed(0) + scale_base = torch.randn(1, 6, 32, device="cuda", dtype=torch.bfloat16) + temb_base = torch.randn(temb_shape, device="cuda", dtype=torch.bfloat16) + chunk_shape = (*temb_shape[:2], 1, temb_shape[-1]) if len(temb_shape) == 4 else (temb_shape[0], 1, + temb_shape[-1]) + weights = [torch.randn(chunk_shape, device="cuda") for _ in range(6)] + + scale_eager = scale_base.clone().requires_grad_(True) + temb_eager = temb_base.clone().requires_grad_(True) + chunk_dim = 2 if len(temb_shape) == 4 else 1 + modulation = (scale_eager.unsqueeze(0).float() if len(temb_shape) == 4 else + scale_eager.float()) + temb_eager.float() + eager_outputs = modulation.chunk(6, dim=chunk_dim) + eager_loss = sum((output * weight).sum() for output, weight in zip(eager_outputs, weights)) + eager_grads = torch.autograd.grad(eager_loss, (scale_eager, temb_eager)) + + scale_compiled = scale_base.clone().requires_grad_(True) + temb_compiled = temb_base.clone().requires_grad_(True) + + def loss(scale_shift_table: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + outputs = torch.ops.fastvideo._wan_modulation_forward(scale_shift_table, temb) + return sum((output * weight).sum() for output, weight in zip(outputs, weights)) + + compiled_loss = torch.compile(loss, fullgraph=True)(scale_compiled, temb_compiled) + compiled_grads = torch.autograd.grad(compiled_loss, (scale_compiled, temb_compiled)) + + for compiled_grad, eager_grad in zip(compiled_grads, eager_grads): + torch.testing.assert_close(compiled_grad, eager_grad, rtol=0, atol=0) diff --git a/fastvideo/tests/train/utils/test_torch_compile.py b/fastvideo/tests/train/utils/test_torch_compile.py index aad815176..9e458405f 100644 --- a/fastvideo/tests/train/utils/test_torch_compile.py +++ b/fastvideo/tests/train/utils/test_torch_compile.py @@ -44,13 +44,13 @@ def test_regional_compile_preserves_checkpoint_state_dict(monkeypatch) -> None: block._checkpoint_wrapped_module for block in model.blocks ] assert [kwargs for _, kwargs in calls] == [ - {"fullgraph": True}, - {"fullgraph": True}, + {"fullgraph": True, "options": {"emulate_precision_casts": True}}, + {"fullgraph": True, "options": {"emulate_precision_casts": True}}, ] assert list(model.state_dict()) == state_dict_keys -def test_regional_compile_dispatches_only_no_grad_calls(monkeypatch) -> None: +def test_regional_compile_dispatches_grad_and_no_grad_calls(monkeypatch) -> None: model = _RepeatedModel() compiled_calls = 0 @@ -69,11 +69,11 @@ def test_regional_compile_dispatches_only_no_grad_calls(monkeypatch) -> None: value = torch.randn(2, 4) model(value) - assert compiled_calls == 0 + assert compiled_calls == 2 with torch.no_grad(): model(value) - assert compiled_calls == 2 + assert compiled_calls == 4 def test_regional_compile_forwards_supported_kwargs(monkeypatch) -> None: @@ -86,11 +86,31 @@ def test_regional_compile_forwards_supported_kwargs(monkeypatch) -> None: monkeypatch.setattr(torch, "compile", _fake_compile) - _compile_model_regions(model, {"dynamic": False}) + _compile_model_regions(model, { + "dynamic": False, + "options": { + "emulate_precision_casts": False, + "max_autotune": True, + }, + }) assert calls == [ - {"fullgraph": True, "dynamic": False}, - {"fullgraph": True, "dynamic": False}, + { + "fullgraph": True, + "dynamic": False, + "options": { + "emulate_precision_casts": False, + "max_autotune": True, + }, + }, + { + "fullgraph": True, + "dynamic": False, + "options": { + "emulate_precision_casts": False, + "max_autotune": True, + }, + }, ]