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

[zjn20030811]

Summary

• apply the learned RMSNorm scale to Matrix-Game-3 action-attention queries and keys
• perform the RMS reduction in float32 and restore the input dtype, matching the main Wan transformer implementation
• add a CPU regression for the forward value and weight gradient

Problem

The action module's WanRMSNorm registers a learned weight, and the released base and distilled checkpoints contain these action-normalization parameters. Its forward path currently ignores the weight and returns only an unscaled RMS-normalized tensor.

Each enabled action block uses four of these norms before mouse and keyboard attention. The released configuration enables 15 action blocks, so 60 learned scale vectors are loaded but have no effect on the attention logits.

The canonical WanRMSNorm in wan/modules/model.py already applies the trained scale and computes the reduction in float32.

Released checkpoint evidence

At official model revision 382aa382a03cec057360755aa72c7735025770bb, both the base and distilled safetensors contain all 60 expected action RMSNorm weights with no missing keys.

For blocks.9.action_model.img_attn_q_norm.weight:

• shape: 64
• minimum: 1.078125
• maximum: 1.484375
• all 64 values differ from 1
• base and distilled values are identical

Loading that released weight into the current and fixed implementations changes all 64 output dimensions. The relative L2 difference is 15.16%, and the current implementation leaves weight.grad unset. The fixed implementation matches weighted RMSNorm and gives all 64 weight elements nonzero gradients.

Fix

Mirror the canonical implementation in the action module:

1. normalize a float32 view
2. cast the normalized value back to the input dtype
3. multiply by the learned weight

Validation

• focused CPU pytest: 1 passed
• output matches torch.nn.functional.rms_norm followed by the learned scale
• changing the scale changes the output and produces nonzero gradients for every scale element
• Black check on the regression test
• Python compilation for both changed files
• git diff --check