Compute fp16 batch norm in fp32 instead of downcasting the params (#54)

BatchNorm keeps its params in fp32 even when the activation is fp16. ATen promotes the input to fp32, uses the params and eps at full precision, and narrows only the result (native_batch_norm_helper). The composite body did the opposite: it cast the params down to fp16, losing mantissa bits and turning any running_var above the fp16 max of 65504 into inf, which silently zeroes the output.

Fix: use prepare_compute_type_for_norm, as layer_norm and group_norm already do. fp32 inputs are unchanged apart from dropping redundant f32 -> f32 casts.

Tests: numerical case with large/tiny running_var (fails on main with max abs error 4.945 and 0.5), plus an IR case pinning the promote/compute/narrow body.

Author: gokulkrishna98

coreai_torch/_aten_to_core.py

@@ -706,16 +706,27 @@ def replace_batch_norm(
     def batch_norm(
         input: Value, gamma: Value, beta: Value, mean: Value, variance: Value
     ) -> Value:
+        # Match ATen: compute in fp32 and narrow only the result. Downcasting
+        # the fp32 params to fp16 loses precision and overflows to inf.
+        input_compute, compute_type, needs_downcast = prepare_compute_type_for_norm(
+            input, ele_type, loc
+        )
+
         eps_casted = coreai.constant(node.args[6])
 
-        eps_casted = coreai.cast(eps_casted, ele_type)
+        def to_compute_type(v: Value) -> Value:
+            if v.type.element_type == compute_type:
+                return v
+            return coreai.cast(v, compute_type)
+
+        eps_casted = to_compute_type(eps_casted)
 
         expand_dims = [0] + list(
             range(2, x_type.rank)
         )  # expand [C] to [1, C, 1, ...] for broadcasting
 
         def expand_and_cast(v: Value) -> Value:
-            return coreai.cast(coreai.expand_dims(v, expand_dims), ele_type)
+            return to_compute_type(coreai.expand_dims(v, expand_dims))
 
         weight, bias, running_mean, running_var = (
             expand_and_cast(v) for v in [gamma, beta, mean, variance]
@@ -724,12 +735,15 @@ def expand_and_cast(v: Value) -> Value:
         # ((x - mean) / sqrt(var + eps)) * weight + bias
         std = coreai.sqrt(coreai.broadcasting_add(running_var, eps_casted))
         normalized = coreai.broadcasting_divide(
-            coreai.broadcasting_sub(input, running_mean), std
+            coreai.broadcasting_sub(input_compute, running_mean), std
         )
 
-        return coreai.broadcasting_add(
+        result = coreai.broadcasting_add(
             coreai.broadcasting_mul(normalized, weight), bias
         )
+        if needs_downcast:
+            result = coreai.cast(result, ele_type)
+        return result
 
     # In inference mode, outputs[1] and [2] are empty placeholder tensors (shape [0]).
     np_dtype = _get_coreai_to_numpy_dtype()[ele_type]

tests/ops/test_ops.py

@@ -692,17 +692,9 @@ async def test_nan_propagation(self) -> None:
         await validate_numerical_output(model=model, y=y, x=x)
 
 
-@pytest.mark.parametrize(
-    "x",
-    [
-        torch.rand(2, 3, 8, 8),
-        torch.rand(2, 3, 8, 8, dtype=torch.float16),  # fp16
-    ],
-)
-@pytest.mark.parametrize(
-    "dynamic_dims", [tuple(), (0,), (2,), (3,), (0, 2), (0, 3), (0, 2, 3)]
-)
-async def test_batchnorm(x: Tensor, dynamic_dims: tuple[int]) -> None:
+class TestBatchNorm:
+    """Tests for native_batch_norm → coreai batch_norm composite."""
+
     class BatchNormModel(nn.Module):
         def __init__(self) -> None:
             super().__init__()
@@ -711,10 +703,47 @@ def __init__(self) -> None:
         def forward(self, x: Tensor) -> Tensor:
             return self.bn(x)
 
-    model = BatchNormModel().eval()
-    dim_names = {0: "batch", 1: "channels", 2: "height", 3: "width"}
-    dynamic_shapes = make_dynamic_shapes(x={d: dim_names[d] for d in dynamic_dims})
-    await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
+    @pytest.mark.parametrize(
+        "x",
+        [
+            torch.rand(2, 3, 8, 8),
+            torch.rand(2, 3, 8, 8, dtype=torch.float16),  # fp16
+        ],
+    )
+    @pytest.mark.parametrize(
+        "dynamic_dims", [tuple(), (0,), (2,), (3,), (0, 2), (0, 3), (0, 2, 3)]
+    )
+    async def test_basic(self, x: Tensor, dynamic_dims: tuple[int]) -> None:
+        model = self.BatchNormModel().eval()
+        dim_names = {0: "batch", 1: "channels", 2: "height", 3: "width"}
+        dynamic_shapes = make_dynamic_shapes(x={d: dim_names[d] for d in dynamic_dims})
+        await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
+
+    @pytest.mark.parametrize(
+        "running_var, running_mean, weight",
+        [
+            # running_var past the fp16 max (65504) becomes inf when the params are
+            # downcast before the sqrt, which silently zeroes the whole output.
+            ([1e6, 2e5, 9e4], [0.0, 1.0, -2.0], [2000.0, 1000.0, 500.0]),
+            # Tiny variances make eps and the param mantissas precision-critical.
+            ([1e-6, 4e-5, 2.0], [0.5, -0.25, 100.0], [1.0, 2.0, 0.5]),
+        ],
+    )
+    async def test_fp16_input_fp32_params(
+        self,
+        running_var: list[float],
+        running_mean: list[float],
+        weight: list[float],
+    ) -> None:
+        """An fp16 activation with fp32 params must be computed in fp32."""
+        model = self.BatchNormModel().eval()
+        model.bn.running_var.data = torch.tensor(running_var)
+        model.bn.running_mean.data = torch.tensor(running_mean)
+        model.bn.weight.data = torch.tensor(weight)
+        model.bn.bias.data = torch.tensor([0.0, 1.0, -1.0])
+
+        x = (torch.rand(2, 3, 4, 4) * 2 - 1).half()
+        await validate_numerical_output(model=model, x=x)
 
 
 @pytest.mark.parametrize("dynamic", [False, True])

tests/ops/test_ops_ir.py

@@ -577,16 +577,16 @@ def forward(self, x: Tensor) -> Tensor:
 
 
 class TestBatchNormIR:
-    def test_static_composite(self) -> None:
-        class BatchNormModel(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.bn = nn.BatchNorm2d(3)
+    class BatchNormModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.bn = nn.BatchNorm2d(3)
 
-            def forward(self, x: Tensor) -> Tensor:
-                return self.bn(x)
+        def forward(self, x: Tensor) -> Tensor:
+            return self.bn(x)
 
-        ir = get_ir(BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4))
+    def test_static_composite(self) -> None:
+        ir = get_ir(self.BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4))
         filecheck_pattern(
             ir,
             check_file="""
@@ -616,6 +616,43 @@ def forward(self, x: Tensor) -> Tensor:
             """,
         )
 
+    def test_fp16_input_computes_in_fp32(self) -> None:
+        """fp16 input is promoted, the fp32 params are used as-is, and only the
+        result is narrowed back to fp16."""
+        ir = get_ir(
+            self.BatchNormModel().eval(), x=torch.rand(1, 3, 4, 4, dtype=torch.float16)
+        )
+        filecheck_pattern(
+            ir,
+            check_file="""
+                // CHECK-LABEL: module {
+                // CHECK-NEXT:   coreai.graph private noinline @batch_norm_{{.*}}(%[[ARG0:.*]]: tensor<1x3x4x4xf16> {coreai.name = "input"}, %[[ARG1:.*]]: tensor<3xf32> {coreai.name = "gamma"}, %[[ARG2:.*]]: tensor<3xf32> {coreai.name = "beta"}, %[[ARG3:.*]]: tensor<3xf32> {coreai.name = "mean"}, %[[ARG4:.*]]: tensor<3xf32> {coreai.name = "variance"}) -> tensor<1x3x4x4xf16> attributes {__coreai_pure__, composite_decl = #coreai.composite_declaration<"batch_norm" = {input_names = ["input", "gamma", "beta", "mean", "variance"], op_attrs = {eps = 9.99999974E-6 : f32, version = 1 : si64}, output_names = ["output"]}>, template_op = "batch_norm"} {
+                // CHECK-NEXT:     %[[V0:.*]] = coreai.constant dense<[1, 3, 1, 1]> : tensor<4xui32>
+                // CHECK-NEXT:     %[[V1:.*]] = coreai.constant dense<9.99999974E-6> : tensor<f32>
+                // CHECK-NEXT:     %[[V2:.*]] = coreai.cast %[[ARG0]] : tensor<1x3x4x4xf16> to tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V3:.*]] = coreai.reshape %[[ARG1]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V4:.*]] = coreai.reshape %[[ARG2]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V5:.*]] = coreai.reshape %[[ARG3]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V6:.*]] = coreai.decomposable.broadcasting_add %[[ARG4]], %[[V1]] : (tensor<3xf32>, tensor<f32>) -> tensor<3xf32>
+                // CHECK-NEXT:     %[[V7:.*]] = coreai.reshape %[[V6]], %[[V0]] : (tensor<3xf32>, tensor<4xui32>) -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V8:.*]] = coreai.sqrt %[[V7]] : tensor<1x3x1x1xf32> -> tensor<1x3x1x1xf32>
+                // CHECK-NEXT:     %[[V9:.*]] = coreai.decomposable.broadcasting_sub %[[V2]], %[[V5]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V10:.*]] = coreai.decomposable.broadcasting_divide %[[V9]], %[[V8]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V11:.*]] = coreai.decomposable.broadcasting_mul %[[V10]], %[[V3]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V12:.*]] = coreai.decomposable.broadcasting_add %[[V11]], %[[V4]] : (tensor<1x3x4x4xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x4x4xf32>
+                // CHECK-NEXT:     %[[V13:.*]] = coreai.cast %[[V12]] : tensor<1x3x4x4xf32> to tensor<1x3x4x4xf16>
+                // CHECK-NEXT:     coreai.output %[[V13]] : tensor<1x3x4x4xf16>
+                // CHECK-NEXT:   }
+                // CHECK-NEXT:   coreai.graph @main(%[[ARG0]]: tensor<1x3x4x4xf16> {coreai.name = "x"}) -> (tensor<1x3x4x4xf16> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
+                // CHECK-NEXT:     %[[V0]] = coreai.constant dense<1.000000e+00> : tensor<3xf32>
+                // CHECK-NEXT:     %[[V1]] = coreai.constant dense<0.000000e+00> : tensor<3xf32>
+                // CHECK-NEXT:     %[[V2]] = coreai.invoke @batch_norm_{{.*}}(%[[ARG0]], %[[V0]], %[[V1]], %[[V1]], %[[V0]])  : (tensor<1x3x4x4xf16>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>, tensor<3xf32>) -> tensor<1x3x4x4xf16>
+                // CHECK-NEXT:     coreai.output %[[V2]] : tensor<1x3x4x4xf16>
+                // CHECK-NEXT:   }
+                // CHECK-NEXT: }
+            """,
+        )
+
 
 class TestToCopyIR:
     def test_cast_f32_to_f16(self) -> None: