From 7f372730b4b4ef06702022b70e701f88a4c49cc8 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Tue, 25 Aug 2026 19:16:20 +0000 Subject: [PATCH] Fix quantized MoE on the dense_matmul path `RoutedMoE.get_einsum` builds a Linen einsum and calls it inline, which only worked while the MoE layer was itself a Linen module. Since the move to NNX there is no Linen scope to bind to, so every quantization fails on the dense_matmul path. A tiny Mixtral, one step on CPU: int8 CallCompactUnboundModuleError fp8 CallCompactUnboundModuleError nanoo_fp8 TypeError: Quantization.einsum() got an unexpected keyword argument 'mesh_axes' Bridge the einsums into NNX instead. fp8 keeps its scaling factors and amax histories in Linen variables, so those are created with the parent module rather than on the first call: allocating them later would grow the module graph inside the scanned layer loop, which NNX rejects, and puts the allocation under a trace. Their shape is fixed, so a canonical operand pair materializes them and the einsum still takes operands of any shape. AQT bridges on first use, since its state is shaped after the operands. Each call site passes a stable name, so no two share quantization state; looking up a name that was never registered raises a ValueError listing the ones that were, rather than a bare KeyError. NANOO also lost its `einsum` and its place in the isinstance check in the same migration; both are restored, which gives the orphaned `Fp8Einsum` class its purpose back. The four integration tests carry no hardware marker: this is a binding bug that shows up on every backend, and both fp8 flavors are emulated in XLA, so they run on CPU in seconds. All four fail before this change. The unit tests alongside them pin the binding itself: `create_fp8_einsum` for both fp8 flavors, both branches of `apply_einsum_in_nnx` (the early return for a plain callable, and the AQT bridge, which must reuse the wrapper it built rather than rebuild one per call), and the fp8 and AQT paths through `get_einsum`. Review asked for correctness coverage alongside those. The forward pass of a freshly built layer is exactly the unquantized einsum with both operands cast to the scheme's e4m3, since the scaling factors start at 1 and only move on the backward pass, so the bridged einsum is pinned against that cast rather than a tolerance: e4m3fn for fp8, e4m3fnuz for nanoo_fp8. At layer level the quantized MoE is compared against the same layer run unquantized, which is loose (three mantissa bits across three einsums) but catches a silently unquantized fallback, since the two must not agree exactly either. Still broken and left alone: fp8 on the sparse_matmul path, where `get_quantization_dtypes` reads `self.quant.quant_dg`, which the fp8 classes do not have. --- src/maxtext/layers/moe.py | 44 ++++++++---- src/maxtext/layers/quantizations.py | 57 +++++++++++++++ tests/integration/train_tests.py | 44 ++++++++++++ tests/unit/moe_test.py | 107 +++++++++++++++++++++++++++- tests/unit/quantizations_test.py | 46 ++++++++++++ 5 files changed, 285 insertions(+), 13 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index f051e7a4be..92e52b771f 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -64,6 +64,9 @@ DISPATCH = "dispatch" COMBINE = "combine" +WI_0 = "wi_0" +WI_1 = "wi_1" +WO = "wo" @struct.dataclass @@ -497,6 +500,16 @@ def __init__( else: self._expert_parallelism_name = "expert" + if isinstance(self.quant, (quantizations.Fp8Quantization, quantizations.NANOOFp8Quantization)): + einsum_names = [WI_0, WI_1, WO] + if self.config.capacity_factor > 0: + einsum_names += [DISPATCH, COMBINE] + self.quant_einsums = nnx.Dict( + {name: quantizations.create_fp8_einsum(self.quant, self.dtype, self.rngs) for name in einsum_names} + ) + else: + self.quant_einsums = None + self.gate = GateLogit( in_features_shape=self.moe_expert_input_dim, out_features_shape=self.num_experts, @@ -2717,15 +2730,22 @@ def get_einsum( return jnp.einsum if self.quant: + op_id = einsum_name if einsum_name is not None else "einsum" - def aqt_einsum(*args, **kwargs): # pylint: disable=unused-argument + def quant_einsum(*args, **kwargs): # pylint: disable=unused-argument # simply skip kwargs, since aqt einsum doesn't support any kwargs # like precision - is_aqt = not isinstance(self.quant, quantizations.Fp8Quantization) - kw = {"mesh_axes": rhs_mesh_axes} if is_aqt else {"dtype": self.dtype} - return self.quant.einsum(**kw)(*args) # pytype: disable=attribute-error - - einsum_op = aqt_einsum + if self.quant_einsums is not None: + if op_id not in self.quant_einsums: + raise ValueError( + f"Einsum name '{op_id}' is not registered in quant_einsums. " + f"Available names: {list(self.quant_einsums.keys())}" + ) + return self.quant_einsums[op_id](*args, mutable=["_overwrite_with_gradient"]) + einsum = self.quant.einsum(mesh_axes=rhs_mesh_axes) # pytype: disable=attribute-error + return quantizations.apply_einsum_in_nnx(self, op_id, einsum, ["aqt"], *args) + + einsum_op = quant_einsum else: einsum_op = jnp.einsum return einsum_op @@ -2929,7 +2949,7 @@ def dense_matmul( with jax.named_scope("wi_0"): w0_kernel_axes = ("exp", None, "mlp") w0_kernel = self.maybe_all_gather_kernel_weight_in_expert_parallelism(w0_kernel, w0_kernel_axes) - layer_w0 = self.get_einsum(rhs_mesh_axes=w0_kernel_axes)( + layer_w0 = self.get_einsum(rhs_mesh_axes=w0_kernel_axes, einsum_name=WI_0)( mlp_up_einsum, dispatch, w0_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -2943,7 +2963,7 @@ def dense_matmul( with jax.named_scope("wi_1"): w1_kernel_axes = ("exp", None, "mlp") w1_kernel = self.maybe_all_gather_kernel_weight_in_expert_parallelism(w1_kernel, w1_kernel_axes) - layer_w1 = self.get_einsum(rhs_mesh_axes=w1_kernel_axes)( + layer_w1 = self.get_einsum(rhs_mesh_axes=w1_kernel_axes, einsum_name=WI_1)( mlp_up_einsum, dispatch, w1_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -2957,7 +2977,7 @@ def dense_matmul( with jax.named_scope("wo"): wo_kernel_axes = ("exp", "mlp", None) wo_kernel = self.maybe_all_gather_kernel_weight_in_expert_parallelism(wo_kernel, wo_kernel_axes) - intermediate_layer = self.get_einsum(rhs_mesh_axes=wo_kernel_axes)( + intermediate_layer = self.get_einsum(rhs_mesh_axes=wo_kernel_axes, einsum_name=WO)( mlp_down_einsum, layer_multiply, wo_kernel, @@ -3007,7 +3027,7 @@ def dense_matmul( ), ) with jax.named_scope("wi_0"): - layer_w0 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes)( + layer_w0 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes, einsum_name=WI_0)( "BSM,EMH -> BSEH", inputs, w0_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -3016,7 +3036,7 @@ def dense_matmul( layer_w0 = layer_w0.astype(jnp.float32) layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0") with jax.named_scope("wi_1"): - layer_w1 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes)( + layer_w1 = self.get_einsum(rhs_mesh_axes=self.wi_kernel_axes, einsum_name=WI_1)( "BSM,EMH -> BSEH", inputs, w1_kernel, precision=matmul_precision ) if self.config.mlp_bias: @@ -3027,7 +3047,7 @@ def dense_matmul( layer_multiply = self.apply_ffn_activation(layer_w0, layer_w1) with jax.named_scope("wo"): - intermediate_layer = self.get_einsum(rhs_mesh_axes=self.wo_kernel_axes)( + intermediate_layer = self.get_einsum(rhs_mesh_axes=self.wo_kernel_axes, einsum_name=WO)( "BSEH,EHM -> BSEM", layer_multiply, wo_kernel, diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index a275a0afa8..29599cc54a 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -401,6 +401,10 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns dot_general configured with aqt params.""" return nn.NANOOFp8DotGeneralOp + def einsum(self, dtype: DType = jnp.float32): + """Returns an einsum using the NANOO (fnuz) fp8 formats of AMD MI300/MI325.""" + return Fp8Einsum(dtype=dtype, e4m3_dtype=jnp.float8_e4m3fnuz, e5m2_dtype=jnp.float8_e5m2fnuz) + def _get_int8_quant_config(config): """Get int8 quantization configuration.""" @@ -768,6 +772,59 @@ def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs): return linen_module_cls(name=op_id)(*args, **kwargs) +def create_fp8_einsum(quant: Quantization, dtype: DType, rngs: nnx.Rngs) -> nnx_wrappers.ToNNX: + """Creates an fp8 einsum that an `nnx.Module` can call. + + An fp8 einsum holds its scaling factors and amax histories in Linen variables, which can + only be created while the module is bound to a Linen scope. The bridge into NNX therefore + has to happen while the parent module is being built; creating the state on the first call + instead would grow the module graph inside the scanned layer loop, which NNX rejects. + + The state has a fixed shape, so a canonical pair of operands is enough to materialize it. + The returned einsum still accepts operands of any shape. + + Args: + quant: The fp8 quantization providing the einsum. + dtype: The computation dtype of the einsum. + rngs: The `nnx.Rngs` of the parent module. + """ + wrapper = nnx_wrappers.ToNNX(quant.einsum(dtype=dtype), rngs=rngs) # pytype: disable=attribute-error + dummy_operand = jnp.zeros((1, 1), dtype=dtype) + wrapper.lazy_init("ab,bc->ac", dummy_operand, dummy_operand) + return wrapper + + +def apply_einsum_in_nnx(parent: nnx.Module, op_id: str, einsum, mutable: Sequence[str], *args): + """Applies a quantized Linen einsum from within an NNX parent module. + + A Linen module cannot be called unbound, which is all an `nnx.Module` can offer it, so the + einsum is bridged into NNX on first use and the bridged instance is reused afterwards. + `op_id` must be unique per call site: two call sites sharing an id would also share + quantization state. Bridging here rather than while building the parent needs the operands, + so it only suits einsums whose state is shaped after them, such as AQT. + + Args: + parent: The NNX module hosting the einsum; it must carry an `rngs` attribute. + op_id: Stable identifier for the call site. + einsum: The einsum returned by a `Quantization`, either a Linen module or a callable. + mutable: Variable collections the einsum writes to. + *args: Arguments forwarded to the einsum. + """ + linen_einsum = einsum.func if isinstance(einsum, functools.partial) else einsum + if not isinstance(linen_einsum, nn.Module): + return einsum(*args) + + # The name must not start with an underscore: NNX treats such attributes as static, + # which cannot hold the bridged module's variables. + attr_name = f"quant_einsum_{op_id}" + wrapper = getattr(parent, attr_name, None) + if wrapper is None: + wrapper = nnx_wrappers.ToNNX(linen_einsum, rngs=parent.rngs) + wrapper.lazy_init(*args) + setattr(parent, attr_name, wrapper) + return wrapper(*args, mutable=list(mutable)) + + class NvidaFp8Provider(qwix.QtProvider): """Wraps nn.Fp8DirectDotGeneralOp with Qwix's provider interface.""" diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index 07f8fb446f..48adc26977 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -56,6 +56,17 @@ class TrainTests(unittest.TestCase): "sharding_tolerance=0.1", ] + # Routes the MoE layer through dense_matmul, which is what runs wherever the megablox and + # ragged kernels are unavailable. + _moe_model_overrides = [ + "decoder_block=mixtral", + "num_experts=4", + "num_experts_per_tok=2", + "base_moe_mlp_dim=32", + "sparse_matmul=False", + "megablox=False", + ] + CONFIGS = { "base": [ # short test for train.py with TFDS c4 None, @@ -135,6 +146,19 @@ class TrainTests(unittest.TestCase): rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", ] + _small_model_overrides, + "moe": [ # tests a MoE model, to be combined with a quantization + None, + get_test_config_path(), + f"base_output_directory={_base_output_directory}", + "run_name=runner_test", + "dataset_type=synthetic", # use synthetic dataset_type to decrease training time + "steps=2", + "enable_checkpointing=False", + "enable_goodput_recording=False", + rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", + ] + + _small_model_overrides + + _moe_model_overrides, "te_fp8_delayedscaling": [ # tests base config with te_fp8_delayedscaling None, get_test_config_path(), @@ -288,6 +312,26 @@ def test_gpu_fp8(self): def test_gpu_nanoo_fp8(self): train_main(TrainTests.CONFIGS["nanoo_fp8"] + ["attention=dot_product"]) + # The quantized MoE tests below carry no hardware marker on purpose. They cover how the + # quantized einsums are bound to the MoE layer, which breaks on every backend when it breaks, + # and both fp8 flavors are emulated in XLA rather than needing hardware support. + @pytest.mark.integration_test + def test_moe_int8(self): + train_main(TrainTests.CONFIGS["moe"] + ["quantization=int8"]) + + @pytest.mark.integration_test + def test_moe_fp8(self): + train_main(TrainTests.CONFIGS["moe"] + ["quantization=fp8"]) + + @pytest.mark.integration_test + def test_moe_nanoo_fp8(self): + train_main(TrainTests.CONFIGS["moe"] + ["quantization=nanoo_fp8"]) + + @pytest.mark.integration_test + def test_moe_fp8_token_dropping(self): + # capacity_factor > 0 adds the dispatch and combine einsums to the ones above. + train_main(TrainTests.CONFIGS["moe"] + ["quantization=fp8", "capacity_factor=1.25"]) + @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available") @pytest.mark.integration_test @pytest.mark.gpu_only diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index c183e33d0c..a5813d62a1 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -33,7 +33,7 @@ from maxtext.layers import moe from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import NdInitializer, nd_dense_init, variable_to_logically_partitioned -from maxtext.layers.quantizations import Fp8Quantization +from maxtext.layers.quantizations import configure_quantization, Fp8Quantization from maxtext.utils import max_logging, maxtext_utils from maxtext.utils.sharding import remove_expert_from_partition_spec from tests.utils.test_helpers import get_test_config_path @@ -1735,6 +1735,111 @@ def loss_fn(params, x): max_logging.log("\n" + diff_summary) +class GetEinsumTest(parameterized.TestCase): + """Tests for the quantized einsums RoutedMoE.get_einsum hands to dense_matmul.""" + + def _make_moe(self, quant): + """Builds a small RoutedMoE on the dense_matmul path with the given quantization.""" + cfg = pyconfig.initialize( + [None, get_test_config_path()], + run_name="get_einsum_test", + enable_checkpointing=False, + decoder_block="mixtral", + num_experts=4, + num_experts_per_tok=2, + base_emb_dim=64, + base_mlp_dim=32, + base_moe_mlp_dim=32, + dtype="float32", + weight_dtype="float32", + megablox=False, + sparse_matmul=False, + max_target_length=8, + per_device_batch_size=1, + ) + devices_array = maxtext_utils.create_device_mesh(cfg) + return moe.RoutedMoE( + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=Mesh(devices_array, cfg.mesh_axes), + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + dtype=jnp.float32, + quant=quant, + rngs=nnx.Rngs(0), + ) + + def _quantization(self, quantization): + """Returns the quantization object the config string maps to.""" + return configure_quantization( + pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization=quantization, + ) + ) + + def test_fp8_einsum_is_bound(self): + model = self._make_moe(Fp8Quantization()) + einsum_fn = model.get_einsum(einsum_name=moe.WI_0) + result = einsum_fn("ab,bc->ac", jnp.ones((2, 3)), jnp.ones((3, 4))) + self.assertEqual(result.shape, (2, 4)) + + def test_aqt_einsum_is_bound(self): + model = self._make_moe(self._quantization("int8")) + self.assertIsNone(model.quant_einsums) + einsum_fn = model.get_einsum(einsum_name=moe.WI_0) + result = einsum_fn("ab,bc->ac", jnp.ones((2, 3)), jnp.ones((3, 4))) + self.assertEqual(result.shape, (2, 4)) + + def test_unregistered_quant_einsum_name_raises(self): + model = self._make_moe(Fp8Quantization()) + einsum_fn = model.get_einsum(einsum_name="not_registered") + with self.assertRaises(ValueError) as ctx: + einsum_fn("ab,bc->ac", jnp.ones((2, 3)), jnp.ones((3, 4))) + self.assertIn("not_registered", str(ctx.exception)) + self.assertIn("Available names", str(ctx.exception)) + + @parameterized.named_parameters( + ("fp8", "fp8", jnp.float8_e4m3fn), + ("nanoo_fp8", "nanoo_fp8", jnp.float8_e4m3fnuz), + ) + def test_fp8_einsum_quantizes_both_operands(self, quantization, e4m3_dtype): + """The bridged einsum is the plain one with both operands cast to the scheme's e4m3.""" + model = self._make_moe(self._quantization(quantization)) + lhs = jax.random.normal(jax.random.PRNGKey(0), (4, 16), dtype=jnp.float32) + rhs = jax.random.normal(jax.random.PRNGKey(1), (16, 8), dtype=jnp.float32) + + actual = model.get_einsum(einsum_name=moe.WI_0)("ab,bc->ac", lhs, rhs) + + # The scaling factors start at 1 and are only updated on the backward pass, so a forward + # call on a freshly built layer quantizes by a plain cast. + quantized = jnp.einsum( + "ab,bc->ac", lhs.astype(e4m3_dtype).astype(jnp.float32), rhs.astype(e4m3_dtype).astype(jnp.float32) + ) + np.testing.assert_array_equal(np.asarray(actual), np.asarray(quantized)) + self.assertFalse(np.array_equal(np.asarray(actual), np.asarray(jnp.einsum("ab,bc->ac", lhs, rhs)))) + + @parameterized.named_parameters(("fp8", "fp8"), ("nanoo_fp8", "nanoo_fp8")) + def test_quantized_dense_matmul_tracks_unquantized(self, quantization): + """A quantized MoE layer follows the same layer run unquantized, to within e4m3.""" + reference = self._make_moe(None) + model = self._make_moe(self._quantization(quantization)) + copy_weights(reference, model) + + inputs = jax.random.normal(jax.random.PRNGKey(42), (1, 8, reference.config.base_emb_dim), dtype=jnp.float32) + expected, _, _ = reference(inputs) + actual, _, _ = model(inputs) + + self.assertTrue(np.isfinite(actual).all()) + # e4m3 keeps three mantissa bits, and the layer is quantized at each of wi_0, wi_1 and wo, + # so the agreement is loose; it is the same threshold the qwix MoE test above uses. + relative_error = np.linalg.norm(actual - expected) / np.linalg.norm(expected) + self.assertLess(relative_error, 0.22) + self.assertFalse(np.allclose(actual, expected)) + + def make_moe(cfg, mesh): return moe.RoutedMoE( config=cfg, diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 1c34d128be..d014a66ee0 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -717,6 +717,52 @@ def test_nnx_abstract_state_has_no_intermediates(self): self.assertNotIn("intermediates", state_dict) +class EinsumParent(nnx.Module): + """Minimal NNX parent for apply_einsum_in_nnx tests.""" + + def __init__(self, rngs: nnx.Rngs): + self.rngs = rngs + + +class MoEQuantizedEinsumTest(unittest.TestCase): + """Tests for MoE quantized einsum helpers.""" + + def test_nanoo_fp8_einsum_uses_fnuz_dtypes(self): + quant = quantizations.NANOOFp8Quantization() + einsum_mod = quant.einsum(dtype=jnp.float32) + self.assertEqual(einsum_mod.e4m3_dtype, jnp.float8_e4m3fnuz) + self.assertEqual(einsum_mod.e5m2_dtype, jnp.float8_e5m2fnuz) + + def test_create_fp8_einsum(self): + for quant_str in ("fp8", "nanoo_fp8"): + quant = _configure_quantization(quant_str=quant_str) + wrapper = quantizations.create_fp8_einsum(quant, jnp.float32, nnx.Rngs(0)) + lhs = jnp.ones((4, 8)) + rhs = jnp.ones((8, 16)) + result = wrapper("ab,bc->ac", lhs, rhs, mutable=["_overwrite_with_gradient"]) + self.assertEqual(result.shape, (4, 16)) + + def test_apply_einsum_in_nnx_plain_callable(self): + parent = EinsumParent(nnx.Rngs(0)) + lhs = jnp.ones((2, 3)) + rhs = jnp.ones((3, 4)) + result = quantizations.apply_einsum_in_nnx(parent, "plain", jnp.einsum, [], "ij,jk->ik", lhs, rhs) + expected = jnp.einsum("ij,jk->ik", lhs, rhs) + self.assertTrue(jnp.allclose(result, expected)) + + def test_apply_einsum_in_nnx_aqt_reuses_wrapper(self): + quant = _configure_quantization(quant_str="int8") + parent = EinsumParent(nnx.Rngs(0)) + lhs = jnp.ones((2, 2)) + rhs = jnp.ones((2, 2)) + einsum = quant.einsum(mesh_axes=()) + result1 = quantizations.apply_einsum_in_nnx(parent, "aqt_test", einsum, ["aqt"], "bc,ab->ac", lhs, rhs) + wrapper = getattr(parent, "quant_einsum_aqt_test") + result2 = quantizations.apply_einsum_in_nnx(parent, "aqt_test", einsum, ["aqt"], "bc,ab->ac", lhs, rhs) + self.assertIs(getattr(parent, "quant_einsum_aqt_test"), wrapper) + self.assertEqual(result1.shape, result2.shape) + + class StaticScaleTest(unittest.TestCase): """Tests for static scale extraction."""