From 5db61d927fd8c8255b5e147274dfe572c823df33 Mon Sep 17 00:00:00 2001 From: Guoren Li Date: Wed, 17 Jun 2026 09:38:36 -0700 Subject: [PATCH] [QwixODML] Generalize 3D einsum flattening PiperOrigin-RevId: 933774591 --- qwix/_src/providers/odml.py | 157 +++++++++- qwix/_src/providers/odml_ops.py | 22 +- tests/_src/providers/odml_test.py | 476 ++++++++++++++++++++++++++++++ 3 files changed, 649 insertions(+), 6 deletions(-) diff --git a/qwix/_src/providers/odml.py b/qwix/_src/providers/odml.py index 3494ad1..b9ab817 100644 --- a/qwix/_src/providers/odml.py +++ b/qwix/_src/providers/odml.py @@ -27,6 +27,7 @@ from qwix._src import averaging from qwix._src import interception from qwix._src import qconfig +from qwix._src.core import einsum_info from qwix._src.core import qarray from qwix._src.providers import odml_ops from qwix._src.utils import flax_util @@ -380,6 +381,10 @@ def get_intercept_map(self): self._flatten_dot_general, _dot_general=intercept_map['jax.lax.dot_general'], ) + intercept_map['jax.numpy.einsum'] = functools.partial( + self._flatten_einsum, + _einsum=intercept_map['jax.numpy.einsum'], + ) return intercept_map def _flatten_dot_general(self, *args, _dot_general, **kwargs): @@ -393,11 +398,146 @@ def _flatten_dot_general(self, *args, _dot_general, **kwargs): ): args = list(args) dout = args[1].shape[1:] + original_weight = args[1] args[1] = jax.lax.reshape(args[1], (args[1].shape[0], np.prod(dout))) + odml_ops.forward_metadata(original_weight, args[1]) out = _dot_general(*args, **kwargs) - return jax.lax.reshape(out, out.shape[:-1] + dout) + res = jax.lax.reshape(out, out.shape[:-1] + dout) + odml_ops.forward_metadata(out, res) + return res return _dot_general(*args, **kwargs) + def _flatten_einsum(self, *args, _einsum, **kwargs): + """Flatten 3D RHS weights to 2-D when contracting in the middle. + + This handles a limited einsum family: + ...D,NDH->...NH + + Examples include 'TD,NDH->TNH', 'BTD,NDH->BTNH', 'BSTD,NDH->BSTNH', etc., + where LHS may have any number of non-contracting dimensions. + + Args: + *args: Positional arguments to the original einsum operation. + _einsum: The original einsum function being intercepted. + **kwargs: Keyword arguments to the original einsum operation. + + Returns: + The result of the einsum operation. + """ + + if len(args) == 3: + einsum_str, lhs, rhs = args + weight_name = aux_data.get(rhs, odml_ops.AuxDataKey.WEIGHT_NAME, None) + if ( + isinstance(einsum_str, str) + and weight_name is not None + and rhs.ndim == 3 + ): + try: + info = einsum_info.EinsumInfo.parse( + einsum_str, ndims=(lhs.ndim, rhs.ndim) + ) + except (NotImplementedError, ValueError): + return _einsum(*args, **kwargs) + if ( + info.lhs + and len(info.rhs) == 3 + and len(info.contract_chars) == 1 + and info.contract_chars[0] == info.lhs[-1] + and not info.batch_chars + ): + contract_char = info.contract_chars[0] + if contract_char in info.rhs: + contract_idx = info.rhs.index(contract_char) + non_contract_axes = [i for i in (0, 1, 2) if i != contract_idx] + expected_out_chars = ( + info.lhs[:-1] + + info.rhs[non_contract_axes[0]] + + info.rhs[non_contract_axes[1]] + ) + # We allow the output to be a permutation of the expected output. + # E.g., ...F, NHF -> ...HN (instead of ...NH). + if set(info.out) == set(expected_out_chars): + # Matches 3D RHS weight. Rewrite as matmul: + # lhs (..., D) -> (prod(...), D) + # rhs (N, D, H) -> (D, N*H) (or transposed equivalent) + # and then reshape the output back to (..., N, H). + transpose_perm = ( + contract_idx, + non_contract_axes[0], + non_contract_axes[1], + ) + args = list(args) + args[1] = jax.lax.reshape( + lhs, (int(np.prod(lhs.shape[:-1])), lhs.shape[-1]) + ) + odml_ops.forward_metadata(lhs, args[1]) + + # Forward Propagation: If lhs already has a quantized + # tracer cached (FQ_ARRAY) from a sibling branch, reshape + # it to match args[1] shape and cache it on args[1]. + fq_lhs = aux_data.get(lhs, odml_ops.AuxDataKey.FQ_ARRAY, None) + if fq_lhs is not None: + if isinstance(fq_lhs, str) and fq_lhs == 'self': + aux_data.set(args[1], odml_ops.AuxDataKey.FQ_ARRAY, 'self') + else: + fq_lhs_2d = jax.lax.reshape( + fq_lhs, (int(np.prod(lhs.shape[:-1])), lhs.shape[-1]) + ) + odml_ops.forward_metadata(fq_lhs, fq_lhs_2d) + aux_data.set(args[1], odml_ops.AuxDataKey.FQ_ARRAY, fq_lhs_2d) + + args[2] = jax.lax.reshape( + jax.lax.transpose(rhs, transpose_perm), + ( + rhs.shape[contract_idx], + rhs.shape[non_contract_axes[0]] + * rhs.shape[non_contract_axes[1]], + ), + ) + odml_ops.forward_metadata(rhs, args[2]) + aux_data.set( + args[2], + odml_ops.AuxDataKey.FLATTENED_EINSUM_PERM, + transpose_perm, + ) + out = _einsum('ab,bc->ac', args[1], args[2], **kwargs) + + # Backward Propagation: If lhs was not already quantized + # (meaning we are the first sibling to run), the _einsum + # call fake-quantized args[1] and cached it as FQ_ARRAY. + # We reshape it back to lhs shape and cache it on lhs so + # sibling branches can reuse it. + if fq_lhs is None: + fq_lhs_2d = aux_data.get( + args[1], odml_ops.AuxDataKey.FQ_ARRAY, None + ) + if fq_lhs_2d is not None: + if isinstance(fq_lhs_2d, str) and fq_lhs_2d == 'self': + aux_data.set(lhs, odml_ops.AuxDataKey.FQ_ARRAY, 'self') + else: + fq_lhs_3d = jax.lax.reshape(fq_lhs_2d, lhs.shape) + odml_ops.forward_metadata(fq_lhs_2d, fq_lhs_3d) + aux_data.set(lhs, odml_ops.AuxDataKey.FQ_ARRAY, fq_lhs_3d) + + res = jax.lax.reshape( + out, + lhs.shape[:-1] + + ( + rhs.shape[non_contract_axes[0]], + rhs.shape[non_contract_axes[1]], + ), + ) + # If the output layout in info.out is a permutation of the + # expected layout, transpose the result to match info.out. + if info.out != expected_out_chars: + perm = tuple(expected_out_chars.index(c) for c in info.out) + res = jax.lax.transpose(res, perm) + odml_ops.forward_metadata(out, res) + return res + + return _einsum(*args, **kwargs) + def _fake_quant( self, array: jax.Array, @@ -414,12 +554,25 @@ def _fake_quant( assert quant_stat_name is None mdl_path = flax_util.get_current_module_path() weight = self._flatten_params[mdl_path + (weight_name,)] - if weight.shape != array.shape: # when _flatten_dot_general is used. + flattened_perm = aux_data.get( + array, odml_ops.AuxDataKey.FLATTENED_EINSUM_PERM, None + ) + if flattened_perm is not None: + # Apply the same layout as _flatten_einsum to the static weight. + weight = jax.lax.reshape( + jax.lax.transpose(weight, flattened_perm), array.shape + ) + elif weight.shape != array.shape: # when _flatten_dot_general is used. weight = weight.reshape(array.shape) calibration = qarray.calibrate(weight, how) scale, zp = qarray.compute_scale_zero_point(calibration, how.qtype) elif quant_stat_name is not None: # Static-range activations. scale, zp = self._compute_static_scale_zero_point(how, quant_stat_name) + # Match scale/zp rank to the activation flattened by _flatten_einsum. + if scale.ndim != array.ndim and scale.size == 1: + scale = scale.reshape((1,) * array.ndim) + if zp is not None and zp.ndim != array.ndim and zp.size == 1: + zp = zp.reshape((1,) * array.ndim) else: # Dynamic-range activations. scale, zp = None, None diff --git a/qwix/_src/providers/odml_ops.py b/qwix/_src/providers/odml_ops.py index ab4c974..d9c89d1 100644 --- a/qwix/_src/providers/odml_ops.py +++ b/qwix/_src/providers/odml_ops.py @@ -148,6 +148,9 @@ class AuxDataKey(str, enum.Enum): # softmax. FIXED_RANGE = 'fixed_range' # tuple[float, float] + # Permutation used to flatten the einsum weight. + FLATTENED_EINSUM_PERM = 'flattened_einsum_perm' # tuple[int, ...] + # Metadata keys that depend on the value being preserved. # If the value changes (e.g. add, mul), these keys become invalid. @@ -156,6 +159,7 @@ class AuxDataKey(str, enum.Enum): AuxDataKey.FQ_RULE, AuxDataKey.FIXED_RANGE, AuxDataKey.ALLOW_FUSION, + AuxDataKey.FLATTENED_EINSUM_PERM, ) # These ops only change the tensor view or layout, not the values. @@ -214,6 +218,15 @@ def _copy_for_isolation(original_array: jax.Array) -> jax.Array: aux_data.set(array_copy, AuxDataKey.FIXED_RANGE, fixed_range) if aux_data.get(original_array, AuxDataKey.ALLOW_FUSION, False): aux_data.set(array_copy, AuxDataKey.ALLOW_FUSION, True) + flattened_einsum_perm = aux_data.get( + original_array, AuxDataKey.FLATTENED_EINSUM_PERM, None + ) + if flattened_einsum_perm is not None: + aux_data.set( + array_copy, + AuxDataKey.FLATTENED_EINSUM_PERM, + flattened_einsum_perm, + ) return array_copy @@ -346,6 +359,7 @@ def _maybe_fake_quant( Returns: The fake quantized array. """ + # Only quantize float arrays. if array.dtype not in (jnp.float32, jnp.bfloat16): return array @@ -396,7 +410,7 @@ def _maybe_fake_quant( # See if we can reuse cached FQ_ARRAY from a sibling branch. fq_array = aux_data.get(array, AuxDataKey.FQ_ARRAY, None) if fq_array is not None: - if fq_array == 'self': + if isinstance(fq_array, str) and fq_array == 'self': # If the current tensor is already physically fake quantized, return. return array elif aux_data.get(fq_array, AuxDataKey.FQ_RULE, None) == effective_rule: @@ -531,7 +545,7 @@ def __call__(self, x: Any) -> Any: return self._maybe_fake_quant(x, previous_rule, op_id) -def _forward_metadata( +def forward_metadata( inputs: Any, outputs: Any, primitive_name: str | None = None, @@ -659,7 +673,7 @@ class Dropout(QuantizedOp): def __call__(self, *args, **kwargs): out = self._call_original_op(*args, **kwargs) - _forward_metadata(args[self.input_idx[0]], out) + forward_metadata(args[self.input_idx[0]], out) return out @@ -676,7 +690,7 @@ def __init__(self, **kwargs): def __call__(self, primitive, *args, **params): out = self._call_original_op(primitive, *args, **params) - _forward_metadata(args, out, primitive_name=primitive.name) + forward_metadata(args, out, primitive_name=primitive.name) return out diff --git a/tests/_src/providers/odml_test.py b/tests/_src/providers/odml_test.py index 2679146..c9d2870 100644 --- a/tests/_src/providers/odml_test.py +++ b/tests/_src/providers/odml_test.py @@ -21,10 +21,12 @@ from flax import nnx import jax from jax import numpy as jnp +from qwix._src import aux_data from qwix._src import interception from qwix._src import model as qwix_model from qwix._src import qconfig from qwix._src.providers import odml +from qwix._src.providers import odml_ops from qwix._src.utils import flax_util @@ -42,6 +44,65 @@ def __call__(self, x): class OdmlTest(parameterized.TestCase): + def _run_linen_einsum_conversion(self, rules): + class EinsumModel(nn.Module): + + @nn.compact + def __call__(self, x): + return nn.Einsum( + shape=(4, x.shape[-1], 8), + einsum_str='BTD,NDH->BTNH', + use_bias=False, + )(x) + + model = EinsumModel() + qat_provider = odml.OdmlQatProvider(rules) + qat_model = qwix_model.quantize_model(model, qat_provider) + model_input = jnp.arange(2 * 3 * 16, dtype=jnp.float32).reshape(2, 3, 16) + model_input = model_input / jnp.max(model_input) + qat_vars = qat_model.init(jax.random.key(0), model_input) + qat_res, new_vars = qat_model.apply(qat_vars, model_input, mutable=True) + qat_vars.update(new_vars) + + conversion_provider = odml.OdmlConversionProvider( + rules, qat_vars['params'], qat_vars['quant_stats'] + ) + conversion_model = qwix_model.quantize_model(model, conversion_provider) + conversion_res = conversion_model.apply(qat_vars, model_input) + return qat_vars, qat_res, conversion_res + + def _run_nnx_einsum_conversion(self, rules): + class NnxEinsumModel(nnx.Module): + + def __init__(self, rngs: nnx.Rngs): + self.einsum = nnx.Einsum( + 'BTD,NDH->BTNH', + (4, 16, 8), + rngs=rngs, + ) + + def __call__(self, x): + return self.einsum(x) + + model = NnxEinsumModel(nnx.Rngs(0)) + qat_provider = odml.OdmlQatProvider(rules) + model_input = jnp.arange(2 * 3 * 16, dtype=jnp.float32).reshape(2, 3, 16) + model_input = model_input / jnp.max(model_input) + qat_model = qwix_model.quantize_model(model, qat_provider, model_input) + qat_res = qat_model(model_input) + + quant_stats = nnx.to_pure_dict(nnx.state(qat_model, flax_util.QuantStat)) + params = nnx.to_pure_dict(nnx.state(qat_model, nnx.Param)) + + conversion_provider = odml.OdmlConversionProvider( + rules, params, quant_stats + ) + conversion_model = qwix_model.quantize_model( + model, conversion_provider, model_input + ) + conversion_res = conversion_model(model_input) + return quant_stats, qat_res, conversion_res + def test_linen(self): class LinenModel(nn.Module): @@ -154,6 +215,68 @@ def __call__(self, x): }, ) + def test_linen_einsum_multi_axis_weight_conversion(self): + """Tests 3D einsum weights with a middle contracting dimension.""" + rules = [ + qconfig.QuantizationRule( + module_path='.*', + weight_qtype=jnp.int8, + act_qtype=jnp.int8, + ), + ] + + qat_vars, qat_res, conversion_res = self._run_linen_einsum_conversion(rules) + self.assertIn('einsum0_lhs', qat_vars['quant_stats']['Einsum_0']) + self.assertEqual(conversion_res.shape, (2, 3, 4, 8)) + self.assertTrue(jnp.allclose(qat_res, conversion_res)) + + def test_linen_einsum_flattened_activation_static_scale_conversion(self): + """Tests static activation scales survive ODML's einsum flattening.""" + rules = [ + qconfig.QuantizationRule( + module_path='.*', + act_qtype=jnp.int8, + ), + ] + qat_vars, qat_res, conversion_res = self._run_linen_einsum_conversion(rules) + + # QAT stores static activation stats in the original BTD rank. Conversion + # flattens the activation to 2-D before fake quantizing it. + self.assertEqual( + qat_vars['quant_stats']['Einsum_0']['einsum0_lhs']['sum_of_max'].shape, + (1, 1, 1), + ) + self.assertEqual(conversion_res.shape, (2, 3, 4, 8)) + self.assertTrue(jnp.allclose(qat_res, conversion_res)) + + def test_linen_einsum_repeated_labels_fall_back(self): + """Tests repeated-label einsums use the default einsum path.""" + + class RepeatedLabelEinsumModel(nn.Module): + + @nn.compact + def __call__(self, x): + return nn.Einsum( + shape=(4, x.shape[-1], 8), + einsum_str='DD,NDH->DNH', + use_bias=False, + )(x) + + model = RepeatedLabelEinsumModel() + model_input = jnp.arange(16 * 16, dtype=jnp.float32).reshape(16, 16) + model_input = model_input / jnp.max(model_input) + variables = model.init(jax.random.key(0), model_input) + fp_res = model.apply(variables, model_input) + + conversion_provider = odml.OdmlConversionProvider( + [], variables['params'], {} + ) + conversion_model = qwix_model.quantize_model(model, conversion_provider) + conversion_res = conversion_model.apply(variables, model_input) + + self.assertEqual(conversion_res.shape, (16, 4, 8)) + self.assertTrue(jnp.allclose(fp_res, conversion_res)) + def test_nnx(self): class NnxModel(nnx.Module): @@ -204,6 +327,44 @@ def __call__(self, x): conversion_res = conversion_model(model_input) self.assertTrue(jnp.allclose(qat_res, conversion_res)) + def test_nnx_einsum_multi_axis_weight_conversion(self): + """Tests 3D einsum weights with a middle contracting dimension.""" + rules = [ + qconfig.QuantizationRule( + module_path='.*', + weight_qtype=jnp.int8, + act_qtype=jnp.int8, + ), + ] + + quant_stats, qat_res, conversion_res = self._run_nnx_einsum_conversion( + rules + ) + self.assertIn('einsum0_lhs', quant_stats['einsum']) + self.assertEqual(conversion_res.shape, (2, 3, 4, 8)) + self.assertTrue(jnp.allclose(qat_res, conversion_res)) + + def test_nnx_einsum_flattened_activation_static_scale_conversion(self): + """Tests static activation scales survive ODML's einsum flattening for NNX.""" + rules = [ + qconfig.QuantizationRule( + module_path='.*', + act_qtype=jnp.int8, + ), + ] + quant_stats, qat_res, conversion_res = self._run_nnx_einsum_conversion( + rules + ) + + # QAT stores static activation stats in the original BTD rank. Conversion + # flattens the activation to 2-D before fake quantizing it. + self.assertEqual( + quant_stats['einsum']['einsum0_lhs']['sum_of_max'].shape, + (1, 1, 1), + ) + self.assertEqual(conversion_res.shape, (2, 3, 4, 8)) + self.assertTrue(jnp.allclose(qat_res, conversion_res)) + def test_odml_interception_stack(self): """Verifies that ODML providers return interceptors in the correct order.""" rules = [qconfig.QuantizationRule(module_path='.*')] @@ -474,6 +635,65 @@ def __call__(self, x): conversion_res = conversion_model.apply(qat_vars, model_input) self.assertIsNotNone(conversion_res) + def test_sibling_einsum_sharing_keyerror(self): + """Test that sibling branches (one being a flattened 3D einsum, the other not flattened) convert successfully.""" + + class SiblingEinsumModel(nn.Module): + + @nn.compact + def __call__(self, x): + # Preceding quantized layer to set FQ_RULE + x = nn.Dense(features=16, name='pre_dense')(x) + # sibling1: 3D RHS weight middle contracting einsum (flattened in + # conversion) + x1 = nn.Einsum( + shape=(4, x.shape[-1], 8), + einsum_str='BTD,NDH->BTNH', + use_bias=False, + name='sibling1', + )(x) + # sibling2: Standard dense (not flattened) + x2 = nn.Dense(features=8, name='sibling2')(x) + return x1, x2 + + model = SiblingEinsumModel() + rules = [ + qconfig.QuantizationRule( + module_path='.*', + weight_qtype=jnp.int8, + act_qtype=jnp.int8, + ), + ] + + qat_provider = odml.OdmlQatProvider(rules) + qat_model = qwix_model.quantize_model(model, qat_provider) + model_input = jnp.arange(2 * 3 * 16, dtype=jnp.float32).reshape(2, 3, 16) + + model_input = model_input / jnp.max(model_input) + qat_vars = qat_model.init(jax.random.key(0), model_input) + + # Run calibration to accumulate stats + _, new_vars = qat_model.apply(qat_vars, model_input, mutable=True) + qat_vars.update(new_vars) + + # sibling1 runs first, so it should register stats. sibling2 reuses + # tracer and should not register. + flat_stats = flax.traverse_util.flatten_dict(qat_vars['quant_stats']) + stat_keys = {'/'.join(k[:-1]) for k in flat_stats} + + self.assertIn('sibling1/einsum0_lhs', stat_keys) + self.assertNotIn('sibling2/dot_general0_lhs', stat_keys) + + # Conversion should succeed without any KeyError + conversion_provider = odml.OdmlConversionProvider( + rules, + qat_vars['params'], + qat_vars['quant_stats'], + ) + conversion_model = qwix_model.quantize_model(model, conversion_provider) + conversion_res = conversion_model.apply(qat_vars, model_input) + self.assertIsNotNone(conversion_res) + def test_immediate_matched_siblings_sharing_stats(self): """Test that matched sibling branches using immediate quantization share the tracer and stats.""" @@ -641,6 +861,262 @@ def __call__(self, x): ) ) + def test_flatten_einsum_aux_data_propagation(self): + """Verifies aux_data on lhs and rhs propagates in _flatten_einsum.""" + provider = odml.OdmlConversionProvider(rules=[], params={}, quant_stats={}) + lhs = jnp.ones((2, 3, 16)) + rhs = jnp.ones((4, 16, 8)) # N=4, D=16, H=8 + + # 1. Attach various metadata keys to the original lhs input + aux_data.set(lhs, odml_ops.AuxDataKey.IS_ACTIVATION, True) + aux_data.set(lhs, odml_ops.AuxDataKey.FIXED_RANGE, (0.0, 1.0)) + rule = qconfig.QuantizationRule(act_qtype=jnp.int8) + aux_data.set(lhs, odml_ops.AuxDataKey.FQ_RULE, rule) + aux_data.set(lhs, odml_ops.AuxDataKey.ALLOW_FUSION, True) + aux_data.set(lhs, odml_ops.AuxDataKey.FQ_ARRAY, 'self') + + # 2. Tag rhs as a weight and attach metadata to test _flatten_einsum + aux_data.set(rhs, odml_ops.AuxDataKey.WEIGHT_NAME, 'test_weight') + aux_data.set(rhs, odml_ops.AuxDataKey.FIXED_RANGE, (-1.0, 1.0)) + rhs_rule = qconfig.QuantizationRule(weight_qtype=jnp.int8) + aux_data.set(rhs, odml_ops.AuxDataKey.FQ_RULE, rhs_rule) + aux_data.set(rhs, odml_ops.AuxDataKey.ALLOW_FUSION, True) + + def mock_einsum(einsum_str, flat_lhs, flat_rhs, **_kwargs): + del einsum_str # Unused in mock + self.assertEqual(flat_lhs.ndim, 2) + self.assertTrue( + aux_data.get(flat_lhs, odml_ops.AuxDataKey.IS_ACTIVATION, False) + ) + self.assertEqual( + aux_data.get(flat_lhs, odml_ops.AuxDataKey.FIXED_RANGE, None), + (0.0, 1.0), + ) + self.assertEqual( + aux_data.get(flat_lhs, odml_ops.AuxDataKey.FQ_RULE, None), rule + ) + self.assertTrue( + aux_data.get(flat_lhs, odml_ops.AuxDataKey.ALLOW_FUSION, False) + ) + self.assertEqual( + aux_data.get(flat_lhs, odml_ops.AuxDataKey.FQ_ARRAY, None), 'self' + ) + + # Check rhs metadata propagation + self.assertEqual(flat_rhs.ndim, 2) + self.assertEqual( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.WEIGHT_NAME, None), + 'test_weight', + ) + self.assertEqual( + aux_data.get( + flat_rhs, odml_ops.AuxDataKey.FLATTENED_EINSUM_PERM, None + ), + (1, 0, 2), + ) + self.assertEqual( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.FIXED_RANGE, None), + (-1.0, 1.0), + ) + self.assertEqual( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.FQ_RULE, None), rhs_rule + ) + self.assertTrue( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.ALLOW_FUSION, False) + ) + out = jnp.ones((flat_lhs.shape[0], flat_rhs.shape[1])) + aux_data.set(out, odml_ops.AuxDataKey.IS_ACTIVATION, True) + aux_data.set(out, odml_ops.AuxDataKey.ALLOW_FUSION, True) + aux_data.set(out, odml_ops.AuxDataKey.FQ_RULE, rule) + return out + + # 4. Trigger the _flatten_einsum interception logic directly + res = provider._flatten_einsum( + 'BTD,NDH->BTNH', lhs, rhs, _einsum=mock_einsum + ) + self.assertEqual(res.shape, (2, 3, 4, 8)) + self.assertTrue(aux_data.get(res, odml_ops.AuxDataKey.IS_ACTIVATION, False)) + self.assertTrue(aux_data.get(res, odml_ops.AuxDataKey.ALLOW_FUSION, False)) + self.assertEqual(aux_data.get(res, odml_ops.AuxDataKey.FQ_RULE, None), rule) + + def test_flatten_einsum_no_non_contracting_dim(self): + provider = odml.OdmlConversionProvider(rules=[], params={}, quant_stats={}) + lhs = jnp.ones((16,)) + rhs = jnp.ones((4, 16, 8)) + aux_data.set(rhs, odml_ops.AuxDataKey.WEIGHT_NAME, 'test_weight') + + def mock_einsum( + unused_einsum_str, unused_flat_lhs, unused_flat_rhs, **_kwargs + ): + return jnp.ones((4, 8)) + + res = provider._flatten_einsum('D,NDH->NH', lhs, rhs, _einsum=mock_einsum) + self.assertEqual(res.shape, (4, 8)) + + def test_flatten_dot_general_aux_data_propagation(self): + """Verifies aux_data on rhs weight and output propagates.""" + provider = odml.OdmlConversionProvider(rules=[], params={}, quant_stats={}) + lhs = jnp.ones((2, 8, 16)) + rhs = jnp.ones((16, 4, 4)) + dimension_numbers = (((2,), (0,)), ((), ())) + + # 1. Attach metadata to rhs weight to verify propagation after flattening + aux_data.set(rhs, odml_ops.AuxDataKey.WEIGHT_NAME, 'test_weight') + aux_data.set(rhs, odml_ops.AuxDataKey.FIXED_RANGE, (-1.0, 1.0)) + rhs_rule = qconfig.QuantizationRule(weight_qtype=jnp.int8) + aux_data.set(rhs, odml_ops.AuxDataKey.FQ_RULE, rhs_rule) + aux_data.set(rhs, odml_ops.AuxDataKey.ALLOW_FUSION, True) + + rule = qconfig.QuantizationRule(act_qtype=jnp.int8) + + # 2. Define a mock inner dot_general to check weight and output metadata + def mock_dot_general(unused_lhs, flat_rhs, unused_dim_nums, **_kwargs): + self.assertEqual(flat_rhs.ndim, 2) + self.assertEqual( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.WEIGHT_NAME, None), + 'test_weight', + ) + self.assertEqual( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.FIXED_RANGE, None), + (-1.0, 1.0), + ) + self.assertEqual( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.FQ_RULE, None), rhs_rule + ) + self.assertTrue( + aux_data.get(flat_rhs, odml_ops.AuxDataKey.ALLOW_FUSION, False) + ) + + # Mock output carrying metadata + out = jnp.ones((2, 8, 16)) + aux_data.set(out, odml_ops.AuxDataKey.IS_ACTIVATION, True) + aux_data.set(out, odml_ops.AuxDataKey.ALLOW_FUSION, True) + aux_data.set(out, odml_ops.AuxDataKey.FQ_RULE, rule) + return out + + # 3. Trigger _flatten_dot_general and verify output metadata propagation + res = provider._flatten_dot_general( + lhs, rhs, dimension_numbers, _dot_general=mock_dot_general + ) + self.assertEqual(res.shape, (2, 8, 4, 4)) + self.assertTrue(aux_data.get(res, odml_ops.AuxDataKey.IS_ACTIVATION, False)) + self.assertTrue(aux_data.get(res, odml_ops.AuxDataKey.ALLOW_FUSION, False)) + self.assertEqual(aux_data.get(res, odml_ops.AuxDataKey.FQ_RULE, None), rule) + + def test_gemma3_kv_einsum_failure(self): + class Gemma3Attention(nn.Module): + + @nn.compact + def __call__(self, x): + kv_einsum = nn.Einsum( + shape=(2, 1, x.shape[-1], 8), + einsum_str='BSD,CKDH->CBSKH', + use_bias=False, + name='kv_einsum', + ) + return kv_einsum(x) + + model = Gemma3Attention() + rules = [ + qconfig.QuantizationRule( + module_path='.*', + weight_qtype=jnp.int8, + act_qtype=jnp.int8, + act_static_scale=True, + ), + ] + + qat_provider = odml.OdmlQatProvider(rules) + qat_model = qwix_model.quantize_model(model, qat_provider) + model_input = jnp.ones((2, 3, 16), dtype=jnp.float32) + qat_vars = qat_model.init(jax.random.key(0), model_input) + _, new_vars = qat_model.apply(qat_vars, model_input, mutable=True) + qat_vars.update(new_vars) + + conversion_provider = odml.OdmlConversionProvider( + rules, qat_vars['params'], qat_vars['quant_stats'] + ) + conversion_model = qwix_model.quantize_model(model, conversion_provider) + + with self.assertRaisesRegex(ValueError, 'Cannot flatten scale with shape'): + _ = conversion_model.apply(qat_vars, model_input) + + def test_gemma3_gating_einsum_transposed(self): + class Gemma3GatingTransposed(nn.Module): + + @nn.compact + def __call__(self, x): + gating = nn.Einsum( + shape=(2, 8, x.shape[-1]), + einsum_str='...F,NHF->...HN', + use_bias=False, + name='gating_einsum', + ) + return gating(x) + + model = Gemma3GatingTransposed() + rules = [ + qconfig.QuantizationRule( + module_path='.*', + weight_qtype=jnp.int8, + act_qtype=jnp.int8, + act_static_scale=True, + ), + ] + + qat_provider = odml.OdmlQatProvider(rules) + qat_model = qwix_model.quantize_model(model, qat_provider) + model_input = jnp.ones((2, 3, 16), dtype=jnp.float32) + qat_vars = qat_model.init(jax.random.key(0), model_input) + _, new_vars = qat_model.apply(qat_vars, model_input, mutable=True) + qat_vars.update(new_vars) + + conversion_provider = odml.OdmlConversionProvider( + rules, qat_vars['params'], qat_vars['quant_stats'] + ) + conversion_model = qwix_model.quantize_model(model, conversion_provider) + + res = conversion_model.apply(qat_vars, model_input) + self.assertEqual(res.shape, (2, 3, 8, 2)) + + def test_gemma3_gating_einsum(self): + class Gemma3Gating(nn.Module): + + @nn.compact + def __call__(self, x): + gating = nn.Einsum( + shape=(2, 8, x.shape[-1]), + einsum_str='...F,NHF->...NH', + use_bias=False, + name='gating_einsum', + ) + return gating(x) + + model = Gemma3Gating() + rules = [ + qconfig.QuantizationRule( + module_path='.*', + weight_qtype=jnp.int8, + act_qtype=jnp.int8, + act_static_scale=True, + ), + ] + + qat_provider = odml.OdmlQatProvider(rules) + qat_model = qwix_model.quantize_model(model, qat_provider) + model_input = jnp.ones((2, 3, 16), dtype=jnp.float32) + qat_vars = qat_model.init(jax.random.key(0), model_input) + _, new_vars = qat_model.apply(qat_vars, model_input, mutable=True) + qat_vars.update(new_vars) + + conversion_provider = odml.OdmlConversionProvider( + rules, qat_vars['params'], qat_vars['quant_stats'] + ) + conversion_model = qwix_model.quantize_model(model, conversion_provider) + + res = conversion_model.apply(qat_vars, model_input) + self.assertEqual(res.shape, (2, 3, 2, 8)) + if __name__ == '__main__': absltest.main()