[CI][Test] Deflake the rms_norm scaling-property assertions (#54271)
Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
9662ab0835
commit
1a16c2ad2c
@@ -12,7 +12,7 @@ from tests.ir.ir_test_utils import (
|
||||
clone_args,
|
||||
supported_providers,
|
||||
)
|
||||
from tests.kernels.allclose_default import get_default_rtol
|
||||
from tests.utils import set_random_seed
|
||||
from vllm import ir
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
@@ -52,6 +52,7 @@ def test_rms_norm_registration():
|
||||
)
|
||||
class TestRMSNorm:
|
||||
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
|
||||
set_random_seed(0)
|
||||
x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
|
||||
num_tokens=4,
|
||||
hidden_size=8,
|
||||
@@ -66,9 +67,13 @@ class TestRMSNorm:
|
||||
assert out.dtype == x.dtype
|
||||
assert out.device == x.device
|
||||
|
||||
# Check the scaling property of rms norm
|
||||
# Check the scaling property of rms norm. This holds only
|
||||
# approximately: epsilon does not scale with x, so
|
||||
# rms_norm(2x) = x / sqrt(mean(x^2) + epsilon/4), which differs from
|
||||
# rms_norm(x) by the epsilon term. Use the op's declared tolerance
|
||||
# rather than a tighter hard-coded one.
|
||||
out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon)
|
||||
torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3)
|
||||
assert_close(ir.ops.rms_norm, out2, out)
|
||||
|
||||
# Mean square should be approximately 1 (ignoring epsilon and weight scaling)
|
||||
combined_norm = out.float() / weight.float()
|
||||
@@ -266,6 +271,7 @@ def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
|
||||
)
|
||||
class TestFusedAddRMSNorm:
|
||||
def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
|
||||
set_random_seed(0)
|
||||
x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
|
||||
num_tokens=4,
|
||||
hidden_size=8,
|
||||
@@ -304,7 +310,7 @@ class TestFusedAddRMSNorm:
|
||||
out2, _ = fused_add_rms_norm_native(
|
||||
x * 2.0, torch.zeros_like(x), weight, epsilon=epsilon
|
||||
)
|
||||
torch.testing.assert_close(out2, out1, rtol=get_default_rtol(out), atol=1e-3)
|
||||
assert_close(ir.ops.fused_add_rms_norm, out2, out1)
|
||||
|
||||
# Check behavior with and without weight
|
||||
weight1 = torch.ones_like(weight)
|
||||
|
||||
Reference in New Issue
Block a user