[CI] Initialize fused gated RMSNorm weights (#50377)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 30, 2026 at 14:34 UTC
bdc98bf5060db48f5844fa35f9ffe1608710ba51
1 file changed
+7
tests/kernels/core/test_fused_rms_norm_gated.py
+7
@@ -49,6 +49,11 @@ def test_compiled_vs_eager(
49
device=device,
50
dtype=dtype,
51
)
52
+ # Model parameters use torch.empty because checkpoint loading overwrites
53
+ # them. Initialize the standalone test module so allocator contents cannot
54
+ # introduce NaNs and make this comparison flaky.
55
+ if module.weight is not None:
56
+ module.weight.uniform_(-1, 1)
57
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
58
g = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
59
@@ -94,6 +99,8 @@ def test_compiled_vs_eager_multidim(
99
device=device,
100
dtype=dtype,
101
)
102
+ if module.weight is not None:
103
+ module.weight.uniform_(-1, 1)
104
x = torch.randn(*shape, dtype=dtype, device=device)
105
g = torch.randn(*shape, dtype=dtype, device=device)
106