[bugfix]: fix compiled Wan modulation gradients

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

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

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,

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)

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,
+            },
+        },
     ]