perf(fastwam): add switchable TE RMSNorm for DiT and VAE

Author: XueSongTap

configs/models/embodied/fastwam.yaml

@@ -22,6 +22,8 @@ model:
   action_dit_pretrained_path: checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt
   redirect_common_files: true
   dtype: bfloat16
+  # RMSNorm kernel for DiT q/k and VAE: te for torch < 2.9, wan for torch >= 2.9
+  rmsnorm_impl: te
 
 data:
   num_video_frames: 9

loongforge/embodied/model/fastwam/action/dit.py

@@ -24,6 +24,7 @@
 import torch.nn as nn
 
 from loongforge.embodied.model.fastwam.utils.gradient import gradient_checkpoint_forward
+from loongforge.embodied.model.fastwam.utils.state_dict import EXTRA_STATE_SUFFIX
 from loongforge.embodied.model.fastwam.wan.dit import (
     DiTBlock,
     precompute_freqs_cis,
@@ -79,6 +80,7 @@ def __init__(
         attn_head_dim: int,
         num_layers: int,
         use_gradient_checkpointing: bool = False,
+        rmsnorm_impl: str = "wan",
     ):
         """Initialize the action DiT backbone and embedding layers."""
         super().__init__()
@@ -121,6 +123,7 @@ def __init__(
                     num_heads=num_heads,
                     ffn_dim=ffn_dim,
                     eps=eps,
+                    rmsnorm_impl=rmsnorm_impl,
                 )
                 for _ in range(num_layers)
             ]
@@ -135,11 +138,17 @@ def __init__(
 
     @classmethod
     def backbone_key_set(cls, keys) -> set[str]:
-        """Return pretrained keys that belong to the shared action backbone."""
+        """Return pretrained keys that belong to the shared action backbone.
+
+        ``_extra_state`` is excluded so the expected key set is the same under both
+        RMSNorm implementations; the TE bookkeeping a module built for itself is kept
+        by seeding the load from its own ``state_dict()``.
+        """
         return {
             key
             for key in keys
             if not any(key.startswith(prefix) for prefix in cls.ACTION_BACKBONE_SKIP_PREFIXES)
+            and not key.endswith(EXTRA_STATE_SUFFIX)
         }
 
     @classmethod

loongforge/embodied/model/fastwam/modeling_configuration_fastwam.py

@@ -70,6 +70,12 @@ class FastWAMModelConfig:
     )
     redirect_common_files: bool = True
     dtype: str = "bfloat16"
+    # q/k RMSNorm implementation for both DiT experts: "wan" (upstream module built
+    # on F.rms_norm) or "te" (TransformerEngine).
+    # torch < 2.9.0 has no fused rms_norm CUDA kernel, so "te" is far faster there;
+    # from torch 2.9.0 on the native kernel wins at the DiT's hidden size (3072),
+    # which TE 2.9 has no tuned kernel for. See `wan.dit.make_rmsnorm`.
+    rmsnorm_impl: str = "wan"  # {"wan", "te"}
 
     # ── Nested architecture configs (fixed for Wan2.2-5B, not in YAML) ────────
     video_dit_config: dict[str, Any] = field(default_factory=lambda: {
@@ -114,6 +120,16 @@ def __post_init__(self) -> None:
                 f"got {self.variant!r}"
             )
 
+        # ── Validate rmsnorm_impl ─────────────────────────────────────────────
+        # Mirrors the names accepted by `wan.dit.make_rmsnorm`; kept as a literal
+        # here so importing this config does not pull in torch / TransformerEngine.
+        valid_rmsnorm_impls = {"wan", "te"}
+        if self.rmsnorm_impl not in valid_rmsnorm_impls:
+            raise ValueError(
+                f"FastWAMModelConfig.rmsnorm_impl must be one of {sorted(valid_rmsnorm_impls)}, "
+                f"got {self.rmsnorm_impl!r}"
+            )
+
         # ── Validate num_video_frames constraint via data config ───────────────
         # (num_video_frames lives in DataConfig; validation happens there)
 
@@ -134,3 +150,7 @@ def __post_init__(self) -> None:
         # ── Sync action_dim → nested dit configs ──────────────────────────────
         self.video_dit_config["action_dim"] = self.action_dim
         self.action_dit_config["action_dim"] = self.action_dim
+
+        # ── Sync rmsnorm_impl → nested dit configs ────────────────────────────
+        self.video_dit_config["rmsnorm_impl"] = self.rmsnorm_impl
+        self.action_dit_config["rmsnorm_impl"] = self.rmsnorm_impl

loongforge/embodied/model/fastwam/mot/fastwam.py

@@ -25,6 +25,7 @@
 import logging
 
 from loongforge.embodied.model.fastwam.action.dit import ActionDiT
+from loongforge.embodied.model.fastwam.utils.state_dict import drop_extra_state
 from loongforge.embodied.model.fastwam.wan.loader import load_wan22_ti2v_5b_components
 from loongforge.embodied.model.fastwam.mot.model import MoT
 from loongforge.embodied.model.fastwam.action.schedulers import WanContinuousFlowMatchScheduler
@@ -1157,7 +1158,9 @@ def infer(
     def save_checkpoint(self, path, optimizer=None, step=None):
         """Save MoT, optional proprio encoder, and optimizer checkpoint state."""
         payload = {
-            "mot": self.mot.state_dict(),
+            # `_extra_state` is dropped so the file does not depend on the RMSNorm
+            # implementation it was trained with (see `utils.state_dict`).
+            "mot": drop_extra_state(self.mot.state_dict()),
             "step": step,
             "torch_dtype": str(self.torch_dtype),
         }
@@ -1171,10 +1174,10 @@ def load_checkpoint(self, path, optimizer=None):
         """Load MoT, optional proprio encoder, and optimizer checkpoint state."""
         payload = torch.load(path, map_location="cpu")
         if "mot" in payload:
-            self.mot.load_state_dict(payload["mot"], strict=False)
+            self.mot.load_state_dict(drop_extra_state(payload["mot"]), strict=False)
         elif "dit" in payload:
             logger.warning("Loading legacy `dit` checkpoint into video expert only.")
-            self.video_expert.load_state_dict(payload["dit"], strict=False)
+            self.video_expert.load_state_dict(drop_extra_state(payload["dit"]), strict=False)
         else:
             raise ValueError(f"Checkpoint missing both `mot` and `dit` keys: {path}")
         if self.proprio_encoder is not None:

loongforge/embodied/model/fastwam/utils/state_dict.py

@@ -15,6 +15,43 @@
 # limitations under the License.
 """State-dict conversion helpers for FastWAM Wan components."""
 
+EXTRA_STATE_SUFFIX = "._extra_state"
+
+
+def drop_extra_state(state_dict):
+    """Return ``state_dict`` without Transformer Engine ``_extra_state`` entries.
+
+    Transformer Engine modules add a ``<prefix>._extra_state`` key holding FP8 recipe
+    bookkeeping (an empty ``uint8`` tensor under bf16, so dropping it is lossless).
+    Their learnable parameter is still named ``weight``, so this key is the only
+    schema difference between the TE and Wan RMSNorm implementations.
+
+    Apply this on both sides of a checkpoint round trip: on save so the file is
+    implementation-agnostic, and on load so an existing TE-saved file carries no
+    keys the Wan implementation cannot place. A freshly built TE module already
+    holds valid bookkeeping, so leaving it untouched is the correct outcome.
+
+    Filtering the key set up front is preferred over a
+    ``register_load_state_dict_post_hook`` that edits ``incompatible_keys``: the
+    hook would also mask a genuinely absent key at every call site it is installed
+    on, whereas here ``strict`` keeps its full meaning for real weights.
+
+    This works because every load site that can see a cross-implementation file
+    passes ``strict=False`` (``mot.fastwam.load_checkpoint``, ``wan.core``,
+    ``wan.loader``). Under ``strict=True`` a TE module would still report its own
+    ``_extra_state`` as missing, so seed such a load from the module's own
+    ``state_dict()`` — this is what ``ActionDiT.from_pretrained`` does, together
+    with ``ActionDiT.backbone_key_set`` filtering the same suffix out of the keys
+    it expects the payload to provide.
+
+    Args:
+        state_dict: Mapping of parameter name to tensor.
+
+    Returns:
+        A new dict with every ``._extra_state`` key removed.
+    """
+    return {key: value for key, value in state_dict.items() if not key.endswith(EXTRA_STATE_SUFFIX)}
+
 
 def wan_video_vae_state_dict_converter(state_dict):
     """Convert Wan VAE checkpoint keys to the local module prefix layout."""

loongforge/embodied/model/fastwam/wan/core.py

@@ -24,6 +24,7 @@
 from PIL import Image
 
 from loongforge.embodied.model.fastwam.action.schedulers import WanContinuousFlowMatchScheduler
+from loongforge.embodied.model.fastwam.utils.state_dict import drop_extra_state
 from loongforge.embodied.model.fastwam.wan.dit import WanVideoDiT
 from loongforge.embodied.model.fastwam.wan.loader import load_wan22_ti2v_5b_components
 
@@ -437,7 +438,9 @@ def infer(
     def save_checkpoint(self, path, optimizer=None, step=None):
         """Save DiT weights and optional optimizer state to a checkpoint path."""
         payload = {
-            "dit": self.dit.state_dict(),
+            # `_extra_state` is dropped so the file does not depend on the RMSNorm
+            # implementation it was trained with (see `utils.state_dict`).
+            "dit": drop_extra_state(self.dit.state_dict()),
             "step": step,
             "torch_dtype": str(self.torch_dtype),
         }
@@ -448,7 +451,7 @@ def save_checkpoint(self, path, optimizer=None, step=None):
     def load_checkpoint(self, path, optimizer=None):
         """Load DiT weights and optional optimizer state from a checkpoint path."""
         payload = torch.load(path, map_location="cpu")
-        self.dit.load_state_dict(payload["dit"], strict=False)
+        self.dit.load_state_dict(drop_extra_state(payload["dit"]), strict=False)
         if optimizer is not None and "optimizer" in payload:
             optimizer.load_state_dict(payload["optimizer"])
         return payload

loongforge/embodied/model/fastwam/wan/dit.py

@@ -22,6 +22,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import transformer_engine.pytorch as te
 from einops import rearrange
 
 from loongforge.embodied.model.fastwam.utils.gradient import gradient_checkpoint_forward
@@ -177,28 +178,75 @@ def create_group_causal_attn_mask(
     return attn_mask
 
 
-class RMSNorm(nn.Module):
-    """Root-mean-square normalization for Wan DiT projections."""
+class WanRMSNorm(nn.Module):
+    """Root-mean-square normalization for Wan DiT projections via ``F.rms_norm``.
+
+    This is upstream Wan's own implementation, hence the name.
+    """
 
     def __init__(self, dim, eps=1e-5):
         """Initialize RMSNorm scale and epsilon."""
         super().__init__()
         self.eps = eps
         self.weight = nn.Parameter(torch.ones(dim))
 
-    def norm(self, x):
-        """Normalize input by root mean square magnitude."""
-        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
-
     def forward(self, x):
-        """Apply RMS normalization and learned scaling via the fused ATen kernel."""
+        """Apply RMS normalization and learned scaling via the ATen kernel."""
         # F.rms_norm only fuses when weight and x share a dtype; otherwise scale
         # separately to keep the original dtype promotion (fp32 weight -> fp32 out).
         if self.weight.dtype == x.dtype:
             return F.rms_norm(x, self.weight.shape, weight=self.weight, eps=self.eps)
         return F.rms_norm(x, self.weight.shape, eps=self.eps) * self.weight
 
 
+class TERMSNorm(te.RMSNorm):
+    """Transformer Engine RMSNorm for Wan DiT attention projections."""
+
+    def __init__(self, dim, eps=1e-5):
+        """Initialize a standard (non-zero-centered) RMSNorm."""
+        # FastWAM constructs experts on CPU before moving them to the target GPU.
+        super().__init__(dim, eps=eps, device="cpu", zero_centered_gamma=False)
+
+
+def make_rmsnorm(dim: int, eps: float, impl: str) -> nn.Module:
+    """Build the q/k RMSNorm module selected by ``impl``.
+
+    Both implementations expose a single ``weight`` parameter of shape ``(dim,)``,
+    so checkpoints are interchangeable except for the ``_extra_state`` key that TE
+    modules add (see ``utils.state_dict.drop_extra_state``).
+
+    Which one is faster depends on the torch build, so the choice is configuration
+    rather than autodetection:
+
+    * **torch < 2.9.0** has no fused ``rms_norm`` CUDA kernel — ``F.rms_norm``
+      decomposes into seven fp32 elementwise/reduction kernels, and TE is ~9.5x
+      faster. Use ``te``.
+    * **torch >= 2.9.0** ships ``vectorized_layer_norm_kernel<..., rms_norm=true>``
+      (22 registers, 100% occupancy). TE 2.9 only precompiles tuned kernels for
+      hidden sizes {512, 768, 1024, 2048, 4096, 8192}, so the DiT's 3072 falls back
+      to ``rmsnorm_fwd_general_kernel`` (182 registers, 12.5% occupancy) and the
+      native kernel is ~2.0x faster. Use ``wan``.
+
+    Args:
+        dim: Normalized (last) dimension size.
+        eps: Epsilon added to the mean square before the reciprocal square root.
+        impl: ``"wan"`` (upstream ``F.rms_norm`` module) or ``"te"``.
+
+    Returns:
+        The constructed RMSNorm module.
+
+    Raises:
+        ValueError: If ``impl`` is not a recognized implementation name.
+    """
+    if impl == "wan":
+        return WanRMSNorm(dim, eps=eps)
+    if impl == "te":
+        return TERMSNorm(dim, eps=eps)
+    raise ValueError(
+        f"Unknown RMSNorm implementation {impl!r}; expected one of ['wan', 'te']."
+    )
+
+
 class AttentionModule(nn.Module):
     """Small wrapper around FastWAM flash attention."""
 
@@ -216,7 +264,14 @@ def forward(self, q, k, v, ctx_mask=None):
 class SelfAttention(nn.Module):
     """Wan DiT self-attention block with RoPE."""
 
-    def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, eps: float = 1e-6):
+    def __init__(
+        self,
+        hidden_dim: int,
+        attn_head_dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        rmsnorm_impl: str = "wan",
+    ):
         """Initialize self-attention projections and RMS norms."""
         super().__init__()
         self.hidden_dim = hidden_dim
@@ -228,8 +283,8 @@ def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, eps: flo
         self.k = nn.Linear(hidden_dim, self.attn_hidden_dim)
         self.v = nn.Linear(hidden_dim, self.attn_hidden_dim)
         self.o = nn.Linear(self.attn_hidden_dim, hidden_dim)
-        self.norm_q = RMSNorm(self.attn_hidden_dim, eps=eps)
-        self.norm_k = RMSNorm(self.attn_hidden_dim, eps=eps)
+        self.norm_q = make_rmsnorm(self.attn_hidden_dim, eps=eps, impl=rmsnorm_impl)
+        self.norm_k = make_rmsnorm(self.attn_hidden_dim, eps=eps, impl=rmsnorm_impl)
 
         # self.attn = AttentionModule(self.num_heads)
 
@@ -247,7 +302,14 @@ def forward(self, x, freqs, self_attn_mask: Optional[torch.Tensor] = None):
 class CrossAttention(nn.Module):
     """Wan DiT cross-attention block for text context."""
 
-    def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, eps: float = 1e-6):
+    def __init__(
+        self,
+        hidden_dim: int,
+        attn_head_dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        rmsnorm_impl: str = "wan",
+    ):
         """Initialize cross-attention projections and RMS norms."""
         super().__init__()
         self.hidden_dim = hidden_dim
@@ -259,8 +321,8 @@ def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, eps: flo
         self.k = nn.Linear(hidden_dim, self.attn_hidden_dim)
         self.v = nn.Linear(hidden_dim, self.attn_hidden_dim)
         self.o = nn.Linear(self.attn_hidden_dim, hidden_dim)
-        self.norm_q = RMSNorm(self.attn_hidden_dim, eps=eps)
-        self.norm_k = RMSNorm(self.attn_hidden_dim, eps=eps)
+        self.norm_q = make_rmsnorm(self.attn_hidden_dim, eps=eps, impl=rmsnorm_impl)
+        self.norm_k = make_rmsnorm(self.attn_hidden_dim, eps=eps, impl=rmsnorm_impl)
 
         # self.attn = AttentionModule(self.num_heads)
 
@@ -288,16 +350,24 @@ def forward(self, x, gate, residual):
 class DiTBlock(nn.Module):
     """Wan DiT transformer block with self-attention, cross-attention, and MLP."""
 
-    def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6):
+    def __init__(
+        self,
+        hidden_dim: int,
+        attn_head_dim: int,
+        num_heads: int,
+        ffn_dim: int,
+        eps: float = 1e-6,
+        rmsnorm_impl: str = "wan",
+    ):
         """Initialize one DiT block."""
         super().__init__()
         self.hidden_dim = hidden_dim
         self.attn_head_dim = attn_head_dim
         self.num_heads = num_heads
         self.ffn_dim = ffn_dim
 
-        self.self_attn = SelfAttention(hidden_dim, attn_head_dim, num_heads, eps)
-        self.cross_attn = CrossAttention(hidden_dim, attn_head_dim, num_heads, eps)
+        self.self_attn = SelfAttention(hidden_dim, attn_head_dim, num_heads, eps, rmsnorm_impl)
+        self.cross_attn = CrossAttention(hidden_dim, attn_head_dim, num_heads, eps, rmsnorm_impl)
         self.norm1 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False)
         self.norm2 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False)
         self.norm3 = nn.LayerNorm(hidden_dim, eps=eps)
@@ -416,6 +486,7 @@ def __init__(
         action_group_causal_mask_mode: str = "causal",
         video_attention_mask_mode: str = "bidirectional",
         use_gradient_checkpointing: bool = False,
+        rmsnorm_impl: str = "wan",
     ):
         """Initialize patch, text, time, transformer, and output modules."""
         super().__init__()
@@ -459,7 +530,10 @@ def __init__(
         )
         self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6))
         self.blocks = nn.ModuleList(
-            [DiTBlock(hidden_dim, attn_head_dim, num_heads, ffn_dim, eps) for _ in range(num_layers)]
+            [
+                DiTBlock(hidden_dim, attn_head_dim, num_heads, ffn_dim, eps, rmsnorm_impl)
+                for _ in range(num_layers)
+            ]
         )
         self.head = Head(hidden_dim, out_dim, patch_size, eps)
         self.freqs = precompute_freqs_cis_3d(attn_head_dim)