From d11b168db7fe44bfa57452d3a96465187a7c4f9f Mon Sep 17 00:00:00 2001 From: sukru tikves Date: Tue, 22 Sep 2026 06:55:54 -0700 Subject: [PATCH 1/2] Exclude RMSNormGated from 4bit weight quantization (#277) RMSNormGated does not subclass RMSNorm, so it missed the name-based quantization exclusion and the global INT4 spec (per_block, axis=1) was applied to its rank-1 weight, raising "axis 1 out of bounds for tensor of rank 1". Add its module path to _TORCH_MODULE_EXCLUSIONS alongside the other RMSNorm composite variants. --- python/src/coreai_models/export/presets.py | 1 + .../test_macos/test_rms_norm.py | 21 +++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/python/src/coreai_models/export/presets.py b/python/src/coreai_models/export/presets.py index bfca1312..85fed754 100644 --- a/python/src/coreai_models/export/presets.py +++ b/python/src/coreai_models/export/presets.py @@ -33,6 +33,7 @@ "coreai_models.primitives.macos.rope.RoPE": None, "coreai_models.primitives.macos.rms_norm.RMSNorm": None, "coreai_models.primitives.macos.rms_norm.RMSNormPlusOne": None, + "coreai_models.primitives.macos.rms_norm.RMSNormGated": None, } # Embedding modules excluded from iOS palettization. diff --git a/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py b/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py index 030dca7c..ec7d8c13 100644 --- a/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py +++ b/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py @@ -132,6 +132,27 @@ def test_rms_norm_plus_one(self): torch.testing.assert_close(out, ref_out, rtol=1e-3, atol=1e-3) +@pytest.mark.skipif(not HAS_COREAI, reason="coreai-torch not available") +class TestRMSNormGatedPresetExclusion: + """RMSNormGated must be excluded from 4bit weight quantization (#277).""" + + def test_4bit_preset_excludes_gated_rms_norm(self): + from coreai_models.export.presets import get_preset + + quant_config = get_preset("4bit")["torch_quantization_config"] + gated_path = f"{CoreaiTorchRMSNormGated.__module__}.{CoreaiTorchRMSNormGated.__qualname__}" + + assert quant_config["module_type_configs"].get(gated_path, "missing") is None + + # The exclusion is load-bearing: the global weight spec blocks along axis 1, + # which is out of bounds for RMSNormGated's rank-1 weight. + weight_spec = quant_config["global_config"]["op_state_spec"]["weight"] + global_axis = weight_spec["granularity"]["axis"] + weight = CoreaiTorchRMSNormGated(dim=8).weight + assert weight.ndim == 1 + assert global_axis >= weight.ndim + + # ============================================================================= # Functional-parity tests # ============================================================================= From 489f95ab2a881f26ef60e15f6fe1663b3cf9c666 Mon Sep 17 00:00:00 2001 From: Sukru Date: Tue, 22 Sep 2026 20:15:33 -0700 Subject: [PATCH 2/2] Update python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py Co-authored-by: Prathamesh <46148373+pkmandke@users.noreply.github.com> --- .../test_macos/test_rms_norm.py | 21 ------------------- 1 file changed, 21 deletions(-) diff --git a/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py b/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py index ec7d8c13..030dca7c 100644 --- a/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py +++ b/python/tests/test_model_units/test_primitives/test_macos/test_rms_norm.py @@ -132,27 +132,6 @@ def test_rms_norm_plus_one(self): torch.testing.assert_close(out, ref_out, rtol=1e-3, atol=1e-3) -@pytest.mark.skipif(not HAS_COREAI, reason="coreai-torch not available") -class TestRMSNormGatedPresetExclusion: - """RMSNormGated must be excluded from 4bit weight quantization (#277).""" - - def test_4bit_preset_excludes_gated_rms_norm(self): - from coreai_models.export.presets import get_preset - - quant_config = get_preset("4bit")["torch_quantization_config"] - gated_path = f"{CoreaiTorchRMSNormGated.__module__}.{CoreaiTorchRMSNormGated.__qualname__}" - - assert quant_config["module_type_configs"].get(gated_path, "missing") is None - - # The exclusion is load-bearing: the global weight spec blocks along axis 1, - # which is out of bounds for RMSNormGated's rank-1 weight. - weight_spec = quant_config["global_config"]["op_state_spec"]["weight"] - global_axis = weight_spec["granularity"]["axis"] - weight = CoreaiTorchRMSNormGated(dim=8).weight - assert weight.ndim == 1 - assert global_axis >= weight.ndim - - # ============================================================================= # Functional-parity tests # =============================================================================