Apply learned RMSNorm scales in Matrix-Game-3 action attention

Matrix-Game-3/tests/test_action_module.py

@@ -0,0 +1,23 @@
+import torch
+
+from wan.modules.action_module import WanRMSNorm
+
+
+def test_wan_rms_norm_applies_learned_weight():
+    norm = WanRMSNorm(4, eps=1e-6)
+    inputs = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]])
+
+    unit_weight_output = norm(inputs).detach()
+    with torch.no_grad():
+        norm.weight.copy_(torch.tensor([1.0, 1.25, 1.5, 2.0]))
+
+    output = norm(inputs)
+    expected = torch.nn.functional.rms_norm(inputs.float(), (norm.dim,), eps=norm.eps)
+    expected = expected.to(inputs.dtype) * norm.weight
+
+    torch.testing.assert_close(output, expected)
+    assert not torch.allclose(output, unit_weight_output)
+
+    output.sum().backward()
+    assert norm.weight.grad is not None
+    assert torch.count_nonzero(norm.weight.grad) == norm.weight.numel()

Matrix-Game-3/wan/modules/action_module.py

@@ -25,7 +25,10 @@ def forward(self, x):
         Args:
             x(Tensor): Shape [B, L, C]
         """
-        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) # fast_rms_norm(x, self.weight, self.eps)
+        return self._norm(x.float()).to(x.dtype) * self.weight
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
 
 def sinusoidal_embedding_1d(dim, position):
     assert dim % 2 == 0
@@ -330,4 +333,4 @@ def forward(self, x, tt, th, tw, mouse_condition=None, keyboard_condition=None,
             attn = rearrange(attn, '(B S) T H D -> B (T S) (H D)', S=S)
             attn = self.proj_keyboard(attn)
             hidden_states = hidden_states + attn
-        return hidden_states
+        return hidden_states