Fix quantized MoE on the dense_matmul path

src/maxtext/layers/moe.py

@@ -63,6 +63,9 @@
 
 DISPATCH = "dispatch"
 COMBINE = "combine"
+WI_0 = "wi_0"
+WI_1 = "wi_1"
+WO = "wo"
 
 
 @struct.dataclass
@@ -486,6 +489,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,
@@ -2703,15 +2716,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
@@ -2915,7 +2935,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:
@@ -2929,7 +2949,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:
@@ -2943,7 +2963,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,
@@ -2993,7 +3013,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:
@@ -3002,7 +3022,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:
@@ -3013,7 +3033,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,

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."""
 

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

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
@@ -1688,6 +1688,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,