From 2bc82fd741889c6225c4ce1b5a4112fe3a2c8ea2 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Fri, 21 Aug 2026 01:52:07 +0000 Subject: [PATCH 1/2] 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' 55c368de9 already settled what should happen here: schemes that define no gmm quantization rule "execute unquantized GMM". That change handled the qwix side and left this read alone, so fp8 crashes instead of taking the fallback it was given. Read `quant_dg` defensively so it gets there. Expert matmuls running in the compute dtype while the dense layers run fp8 is easy to miss from the config alone, so the layer says so once when it is built. The two new tests carry no hardware marker: the fp8 schemes are emulated in XLA, so a tiny Mixtral trains on CPU in seconds. Both fail before this change. --- src/maxtext/layers/moe.py | 14 ++++++++++++-- tests/integration/train_tests.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index da9e86e320..779995db48 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -486,6 +486,14 @@ def __init__( else: self._expert_parallelism_name = "expert" + if self.config.sparse_matmul and isinstance( + self.quant, (quantizations.Fp8Quantization, quantizations.NANOOFp8Quantization) + ): + max_logging.log( + "fp8 quantization does not reach the MoE expert matmuls on the sparse_matmul path; they run" + f" in {self.dtype}. Set sparse_matmul=False to quantize them." + ) + self.gate = GateLogit( in_features_shape=self.moe_expert_input_dim, out_features_shape=self.num_experts, @@ -1472,8 +1480,10 @@ 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 + # Only AQT describes its numerics through a `quant_dg`. The fp8 schemes define no gmm + # quantization, so their expert matmuls run unquantized, as with the qwix rule below. + quant_dg = getattr(self.quant, "quant_dg", None) + if quant_dg is not None: 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 diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index 07f8fb446f..42cfef7e71 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -135,6 +135,24 @@ class TrainTests(unittest.TestCase): rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", ] + _small_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')}", + "decoder_block=mixtral", + "num_experts=4", + "num_experts_per_tok=2", + "base_moe_mlp_dim=32", + "sparse_matmul=True", + "megablox=False", + ] + + _small_model_overrides, "te_fp8_delayedscaling": [ # tests base config with te_fp8_delayedscaling None, get_test_config_path(), @@ -288,6 +306,16 @@ def test_gpu_fp8(self): def test_gpu_nanoo_fp8(self): train_main(TrainTests.CONFIGS["nanoo_fp8"] + ["attention=dot_product"]) + # No hardware marker: the fp8 schemes do not reach the expert matmuls on this path on any + # backend, and what is being covered is that the layer still builds and trains. + @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"]) + @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available") @pytest.mark.integration_test @pytest.mark.gpu_only From a46a1ad24f47647a49a6059e39c637d310592624 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Mon, 24 Aug 2026 16:02:56 +0000 Subject: [PATCH 2/2] Cover the AQT side of the quant_dg read The fp8 tests only exercise the branch where there is no `quant_dg`, leaving the read that does find one, and the isinstance check that skips the warning, uncovered. int8 is the scheme that reaches the gmm with quantized operands, so running it on the same sparse_matmul config covers both and guards the behaviour this change had to preserve. --- tests/integration/train_tests.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index 42cfef7e71..86c8f76314 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -316,6 +316,11 @@ def test_moe_fp8_sparse_matmul(self): def test_moe_nanoo_fp8_sparse_matmul(self): train_main(TrainTests.CONFIGS["moe_sparse"] + ["quantization=nanoo_fp8"]) + # int8 is the scheme that does reach the gmm, so it guards the other side of the same read. + @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