Fix rms_norm eps inflation from max-|x| rescale; guard all-zero input

coremltools/converters/mil/frontend/torch/ops.py

@@ -3201,21 +3201,26 @@ def rms_norm(context, node):
     # - Prevents FP16 overflow on ANE.
     # - Maintains ANE placement (avoiding CPU/GPU fallback).
     #
-    # Trade-offs:
-    # - May introduce slight numerical differences compared
-    #   to the standard operation due to the division
-    #   and rescaling operations.
-    # - Maximum relative error is typically < 0.1% in practice.
-    #
-    # Note: For applications requiring exact PyTorch parity,
-    # consider using CPU/GPU compute units.
+    # The rescale is algebraically exact because epsilon is rescaled together
+    # with the input (see below):
+    #   sqrt(mean((x/m)^2) + eps/m^2) * m == sqrt(mean(x^2) + eps)
+    # so the only differences vs the standard formulation are floating-point
+    # rounding in the extra division/multiplication.
 
     max_val_tensor = mb.reduce_max(
         x=mb.abs(x=x, name=node.name + "_abs"),
         axes=axes,
         keep_dims=True,
         name=node.name + "_max_val"
         )
+    # Clamp the scale to >= 1: inputs with max|x| <= 1 cannot overflow fp16
+    # when squared (so no rescale is needed), and the clamp prevents a 0/0
+    # (NaN) on all-zero rows, where the reduce_max above is 0.
+    max_val_tensor = mb.maximum(
+        x=max_val_tensor,
+        y=1.0,
+        name=node.name + "_max_val_clamped"
+        )
     x_scaled = mb.real_div(x=x, y=max_val_tensor, name=node.name + "_scale")
     x_scale_squared = mb.square(x=x_scaled, name=node.name + "_square")
     mean_squared = mb.reduce_mean(
@@ -3224,9 +3229,18 @@ def rms_norm(context, node):
         keep_dims=True,
         name=node.name + "_mean_squared"
         )
+    # Rescale epsilon by 1/m^2 to match the rescaled input. Adding the raw
+    # epsilon here would compute sqrt(mean(x^2) + eps * m^2) after the
+    # scale-back, inflating the effective epsilon by m^2 — a large error for
+    # spiky activations (max|x| >> rms(x)) at any compute precision.
+    eps_scaled = mb.real_div(
+        x=eps_val,
+        y=mb.square(x=max_val_tensor, name=node.name + "_max_val_sq"),
+        name=node.name + "_eps_scaled"
+        )
     mean_plus_eps = mb.add(
         x=mean_squared,
-        y=eps_val,
+        y=eps_scaled,
         name=node.name + "_add_eps"
         )
     rms = mb.sqrt(x=mean_plus_eps, name=node.name + "_rms")

coremltools/converters/mil/test/test_rms_norm.py

@@ -116,6 +116,94 @@ def forward(self, x):
         assert not np.isinf(coreml_out).any(), \
             f"Test '{test_name}' produced Inf values"
 
+    @staticmethod
+    def test_spiky_input_eps_not_inflated():
+        """
+        Regression test for Issue #2821: the max-|x| overflow-protection
+        rescale must rescale eps along with the input. Otherwise the
+        effective epsilon becomes eps * max(|x|)^2, which produces large
+        errors for spiky inputs (max|x| >> rms(x)) even at FLOAT32.
+        """
+        class TestModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.norm = torch.nn.RMSNorm(64, eps=1e-5)
+
+            def forward(self, x):
+                return self.norm(x)
+
+        model = TestModel()
+        model.eval()
+
+        torch.manual_seed(0)
+        # Mostly small values with one large component per row:
+        # max(|x|)^2 / mean(x^2) ~ 6e4, so an inflated epsilon dominates.
+        example = 0.01 * torch.randn(1, 8, 64)
+        example[..., 0] = 25.0
+
+        torch_out = model(example).detach().numpy()
+        traced = torch.jit.trace(model, example)
+        mlmodel = ct.convert(
+            traced,
+            inputs=[ct.TensorType(
+                shape=example.shape,
+                dtype=np.float32,
+                name="input"
+            )],
+            outputs=[ct.TensorType(name="output", dtype=np.float32)],
+            convert_to="mlprogram",
+            compute_precision=ct.precision.FLOAT32,
+            compute_units=ct.ComputeUnit.CPU_ONLY,
+        )
+        coreml_out = mlmodel.predict({"input": example.numpy()})["output"]
+
+        rel_l2 = (
+            np.linalg.norm(torch_out - coreml_out) / np.linalg.norm(torch_out)
+        )
+        # fp32 rounding is ~1e-7; the eps-inflation bug produced 3.2e-4 here.
+        assert rel_l2 < 1e-5, (
+            f"rms_norm eps inflation regression: rel_l2={rel_l2:.2e}"
+        )
+
+    @staticmethod
+    def test_zero_input_converted_model():
+        """
+        The converted model (not just the PyTorch reference) must handle an
+        all-zero input: the max-|x| rescale divides by reduce_max(|x|),
+        which is 0 for a zero row unless clamped.
+        """
+        class TestModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.norm = torch.nn.RMSNorm(64)
+
+            def forward(self, x):
+                return self.norm(x)
+
+        model = TestModel()
+        model.eval()
+
+        example = torch.zeros(2, 8, 64)
+        torch_out = model(example).detach().numpy()
+        traced = torch.jit.trace(model, torch.randn(2, 8, 64))
+        mlmodel = ct.convert(
+            traced,
+            inputs=[ct.TensorType(
+                shape=example.shape,
+                dtype=np.float32,
+                name="input"
+            )],
+            outputs=[ct.TensorType(name="output", dtype=np.float32)],
+            convert_to="mlprogram",
+            compute_precision=ct.precision.FLOAT32,
+            compute_units=ct.ComputeUnit.CPU_ONLY,
+        )
+        coreml_out = mlmodel.predict({"input": example.numpy()})["output"]
+
+        assert not np.isnan(coreml_out).any(), \
+            "Converted RMSNorm produced NaN on all-zero input"
+        np.testing.assert_allclose(torch_out, coreml_out, atol=1e-6)
+
     @staticmethod
     def test_edge_cases():
         """