From b5cb2dcd1c7ca9f9419651f0c3b92e70d5fb5dad Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Tue, 25 Aug 2026 19:17:13 +0000 Subject: [PATCH] Fix fp8 MoE on the sparse_matmul path `get_quantization_dtypes` reads `self.quant.quant_dg` whenever a quantization is set, but only AQT carries one, so a MoE model with `sparse_matmul=True` dies before it ever reaches the gmm: fp8 AttributeError: 'Fp8Quantization' object has no attribute 'quant_dg' nanoo_fp8 AttributeError: 'NANOOFp8Quantization' object has no attribute 'quant_dg' Since `sparse_matmul=True` is the default, this is what a user hits first when combining fp8 with a MoE model. `quant_dg` is how AQT spells its numerics, not what the gmm needs. What the hand-crafted rule wants is a forward dtype, which the fp8 schemes can simply name, so they do: float8_e4m3fn for fp8 and float8_e4m3fnuz for nanoo_fp8. The read becomes "AQT describes itself through a `quant_dg`, everyone else names a dtype", and the gmm ends up quantized for every scheme rather than only for AQT. That keeps the three intended behaviours intact. With `use_qwix_quantization=false` every scheme reaches the hand-crafted rule, fp8 included, where it used to crash. With qwix and `fp8_full` the gmm takes the qwix rule, untouched. With qwix and any other scheme `configure_quantization` returns None, so there is no dtype to hand over and the gmm is left alone, which is what a scheme declaring no gmm rule should get. `SparseMatmulQuantizationTest` asserts the dtypes the grouped matmul is actually handed across all six cases, including the two qwix ones that must stay None; it intercepts `mblx.gmm`, so it costs no kernel time. int8 is in there because it is the scheme whose dtypes reach the gmm, covering the `quant_dg` branch the fp8 cases never take. The integration tests take the MoE model the dense-path tests already use and put it back on sparse_matmul with megablox, so the rule is built and used for real rather than falling back to ragged_dot -- which is also what confirms qwix accepts the fnuz dtype AMD needs. A downscaled model needs tiles small enough to divide its operands, hence the tile overrides. They carry no hardware marker, for the same reason the dense ones do not: the failure is an attribute read during tracing rather than anything a kernel does, and fp8 is emulated in XLA, so the quantized gmm still runs on CPU through megablox's interpret mode. The fp8 cases fail before this change. --- src/maxtext/layers/moe.py | 14 ++++--- src/maxtext/layers/quantizations.py | 4 ++ tests/integration/train_tests.py | 56 ++++++++++++++++++++++++-- tests/unit/moe_test.py | 62 +++++++++++++++++++++++++++++ 4 files changed, 127 insertions(+), 9 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 851403e943..e7a0551a97 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -1499,12 +1499,14 @@ def get_tokamax_group_sizes(group_sizes, inputs, _kernel): ) def get_quantization_dtypes(): - lhs_quantize_dtype, rhs_quantize_dtype = None, None - if self.quant is not None: - quant_dg = self.quant.quant_dg - lhs_quantize_dtype = quant_dg.fwd.dg_quantizer.lhs.numerics.get_dtype() - rhs_quantize_dtype = quant_dg.fwd.dg_quantizer.rhs.numerics.get_dtype() - return lhs_quantize_dtype, rhs_quantize_dtype + # AQT describes its numerics through a `quant_dg`, while the fp8 schemes just name the + # dtype they quantize to. Under qwix there is no quantization object here at all, and + # the gmm takes its rule from the qwix rule instead. + quant_dg = getattr(self.quant, "quant_dg", None) + if quant_dg is not None: + return quant_dg.fwd.dg_quantizer.lhs.numerics.get_dtype(), quant_dg.fwd.dg_quantizer.rhs.numerics.get_dtype() + quantize_dtype = getattr(self.quant, "quantize_dtype", None) + return quantize_dtype, quantize_dtype def gmm( inputs, diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index be9611ab7a..0ef3aeca55 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -307,6 +307,9 @@ class Fp8Quantization(Quantization): """Configures Fp8 quantization for NVIDIA GPUs""" quant_mode = "train" + # The forward dtype, for callers that quantize an operand themselves rather than through + # `dot_general_cls` or `einsum`, such as the grouped matmul in MoE. + quantize_dtype = jnp.float8_e4m3fn def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns dot_general configured with aqt params.""" @@ -396,6 +399,7 @@ class NANOOFp8Quantization(Quantization): """Configures NANOO Fp8 quantization for AMD MI300/MI325 GPUs""" quant_mode = "train" + quantize_dtype = jnp.float8_e4m3fnuz def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns dot_general configured with aqt params.""" diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index 05b82e9383..55d3dbec46 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -76,17 +76,37 @@ 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 = [ + _moe_expert_overrides = [ "decoder_block=mixtral", "num_experts=4", "num_experts_per_tok=2", "base_moe_mlp_dim=32", + ] + + # Routes the MoE layer through dense_matmul, which is what runs wherever the megablox and + # ragged kernels are unavailable. + _moe_model_overrides = _moe_expert_overrides + [ "sparse_matmul=False", "megablox=False", ] + # The sparse_matmul path, where megablox builds and uses the gmm quantization rule for real + # rather than falling back to ragged_dot. + _moe_sparse_model_overrides = _moe_expert_overrides + [ + "sparse_matmul=True", + "megablox=True", + ] + + # Every operand has to divide into its tile, and the default tiles are far larger than a + # downscaled model. The embedding dim is 28 or 32 depending on the device count, so 4 is + # the largest tile that fits it either way. + _megablox_tile_overrides = [ + f"{matrix}_tile_{direction}_{dim}={size}" + for matrix in ("wi", "wo") + for direction in ("fwd", "dlhs", "drhs") + for dim, size in (("batch_seq", 16), ("embed_dim", 4), ("mlp_dim", 16)) + ] + _qwen3_overrides = [ "override_model_config=True", "base_num_decoder_layers=2", @@ -195,6 +215,20 @@ class TrainTests(unittest.TestCase): ] + _small_model_overrides + _moe_model_overrides, + "moe_sparse": [ # tests a MoE model on the sparse_matmul path, 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_sparse_model_overrides + + _megablox_tile_overrides, "te_fp8_delayedscaling": [ # tests base config with te_fp8_delayedscaling None, get_test_config_path(), @@ -368,6 +402,22 @@ 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"]) + # The sparse_matmul tests below carry no hardware marker for the same reasons. What they cover + # is an attribute read during tracing rather than anything a kernel does, and megablox runs the + # quantized grouped matmul on CPU through its interpret mode. + @pytest.mark.integration_test + def test_moe_fp8_sparse_matmul(self): + train_main(TrainTests.CONFIGS["moe_sparse"] + ["quantization=fp8"]) + + @pytest.mark.integration_test + def test_moe_nanoo_fp8_sparse_matmul(self): + train_main(TrainTests.CONFIGS["moe_sparse"] + ["quantization=nanoo_fp8"]) + + # int8 takes the `quant_dg` branch of the same read, which the fp8 tests never reach. + @pytest.mark.integration_test + def test_moe_int8_sparse_matmul(self): + train_main(TrainTests.CONFIGS["moe_sparse"] + ["quantization=int8"]) + @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 d05ea9a31a..32edc3b0a1 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -1936,6 +1936,68 @@ def test_quantized_dense_matmul_tracks_unquantized(self, quantization): self.assertFalse(np.allclose(actual, expected)) +class SparseMatmulQuantizationTest(parameterized.TestCase): + """The dtypes RoutedMoE hands the grouped matmul on the sparse_matmul path.""" + + @parameterized.named_parameters( + ("fp8", "fp8", False, jnp.float8_e4m3fn), + ("nanoo_fp8", "nanoo_fp8", False, jnp.float8_e4m3fnuz), + ("int8", "int8", False, jnp.int8), + ("unquantized", "", False, None), + # Under qwix a non-fp8_full scheme declares no gmm rule, so the gmm is left alone. + ("fp8_under_qwix", "fp8", True, None), + ("int8_under_qwix", "int8", True, None), + ) + def test_gmm_quantize_dtypes(self, quantization, use_qwix_quantization, expected_dtype): + cfg = pyconfig.initialize( + [None, get_test_config_path()], + run_name="sparse_matmul_quantization_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", + sparse_matmul=True, + megablox=True, + max_target_length=8, + per_device_batch_size=1, + quantization=quantization, + use_qwix_quantization=use_qwix_quantization, + ) + devices_array = maxtext_utils.create_device_mesh(cfg) + model = 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=configure_quantization(cfg), + rngs=nnx.Rngs(0), + ) + + calls = [] + + def record_gmm(**kwargs): + calls.append(kwargs) + return jnp.zeros((kwargs["lhs"].shape[0], kwargs["rhs"].shape[-1]), dtype=kwargs["preferred_element_type"]) + + inputs = jax.random.normal(jax.random.PRNGKey(0), (1, 8, cfg.base_emb_dim), dtype=jnp.float32) + with mock.patch.object(moe.mblx, "gmm", record_gmm): + with nn_partitioning.axis_rules(cfg.logical_axis_rules): + model(inputs) + + self.assertNotEmpty(calls) + for kwargs in calls: + self.assertEqual(kwargs["lhs_quantize_dtype"], expected_dtype) + self.assertEqual(kwargs["rhs_quantize_dtype"], expected_dtype) + + def make_moe(cfg, mesh): return moe.RoutedMoE( config=cfg,