Fix fp8 MoE on the sparse_matmul path

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

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