From b902f53054773527e4ff697fcb86542f89608bca Mon Sep 17 00:00:00 2001
From: nan <zjn32202153@gmail.com>
Date: Tue, 1 Sep 2026 18:24:31 +0800
Subject: [PATCH] Apply learned RMSNorm scales in action attention

---
 Matrix-Game-3/tests/test_action_module.py  | 23 ++++++++++++++++++++++
 Matrix-Game-3/wan/modules/action_module.py |  7 +++++--
 2 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 Matrix-Game-3/tests/test_action_module.py

diff --git a/Matrix-Game-3/tests/test_action_module.py b/Matrix-Game-3/tests/test_action_module.py
new file mode 100644
index 0000000..f541984
--- /dev/null
+++ b/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()
diff --git a/Matrix-Game-3/wan/modules/action_module.py b/Matrix-Game-3/wan/modules/action_module.py
index 694d936..6f2a216 100644
--- a/Matrix-Game-3/wan/modules/action_module.py
+++ b/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
\ No newline at end of file
+        return hidden_states