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..86c8f76314 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,21 @@ 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"]) + + # 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