quantized ragged dot maxtext integration

PiperOrigin-RevId: 836301101

Author: Google-ML-Automation

src/MaxText/layers/moe.py

@@ -18,6 +18,7 @@
 import enum
 import functools
 import math
+import random
 from typing import Iterable, Optional, Tuple, Union
 
 from aqt.jax.v2 import aqt_tensor as aqt
@@ -860,7 +861,14 @@ def gmm(inputs, kernel, tiling, group_sizes, expert_assignments):
             if kernel.bias or kernel.sparsity_mask or len(kernel.scale) > 1:
               raise ValueError("Unsupported usecase for ragged_dot with quantized kernel.")
             rhs_inputs = kernel.qvalue
-          with set_xla_metadata(ragged_dot_tiling=",".join([str(t) for t in tiling])):
+          if self.config.use_qwix_quantization:
+            # Use full contraction for QWIX quantization to allow quantization
+            # fusion (max reduce over contracting dimension).
+            tiling = (tiling[0], k, tiling[2])
+          with set_xla_metadata(
+              ragged_dot_tiling=",".join([str(t) for t in tiling]),
+              mosaic_fusion_group=f"{random.randint(0, 1000000000)}",
+          ):
             output = jax.lax.ragged_dot(
                 lhs=inputs,
                 rhs=rhs_inputs,

src/MaxText/layers/quantizations.py

@@ -672,7 +672,7 @@ def get_quantization_rule(config: Config):
           weight_calibration_method=config.weight_quantization_calibration_method,
           act_calibration_method=config.act_quantization_calibration_method,
           bwd_calibration_method=config.bwd_quantization_calibration_method,
-          op_names=("dot_general", "gmm"),
+          op_names=("dot_general", "gmm", "ragged_dot"),
       )
     case "fp8_gpu":
       return qwix.QtRule(

tests/train_using_ragged_dot_smoke_test.py

@@ -15,54 +15,61 @@
 """Smoke test for MoE using ragged_dot."""
 
 import os
-import unittest
-from tempfile import gettempdir
+import tempfile
 
 from absl.testing import absltest
+from absl.testing import parameterized
+from MaxText import globals as maxtext_globals
+from MaxText import train
 
-from MaxText.globals import MAXTEXT_PKG_DIR
-from MaxText.train import main as train_main
+train_main = train.main
+MAXTEXT_PKG_DIR = maxtext_globals.MAXTEXT_PKG_DIR
+gettempdir = tempfile.gettempdir
 
 
-class Train(unittest.TestCase):
+class Train(parameterized.TestCase):
   """Smoke test for MoE using ragged_dot in G3 only."""
 
-  def test_tiny_config(self):
+  @parameterized.named_parameters(
+      {"testcase_name": "not_quantized", "quantization": ""},
+      {"testcase_name": "fp8_full", "quantization": "fp8_full"},
+  )
+  def test_tiny_config(self, quantization: str):
     test_tmpdir = os.environ.get("TEST_TMPDIR", gettempdir())
     outputs_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", test_tmpdir)
-    train_main(
-        [
-            None,
-            os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
-            f"base_output_directory={test_tmpdir}",
-            "run_name=ragged_dot_smoke_test",
-            "base_emb_dim=128",
-            "base_num_query_heads=4",
-            "base_num_kv_heads=4",
-            "base_mlp_dim=128",
-            "base_moe_mlp_dim=128",
-            "base_num_decoder_layers=8",
-            "head_dim=128",
-            # TODO(b/441100085): When changing the decoder_block we might
-            # need to adjust the tiling.
-            "decoder_block=deepseek",
-            "attention_type=mla",
-            "num_experts=2",
-            # Enable sparse_matmul.
-            "sparse_matmul=True",
-            # Enable ragged_dot.
-            "megablox=False",
-            "per_device_batch_size=2",
-            "max_target_length=1024",
-            "dataset_type=synthetic",
-            "steps=10",
-            "enable_checkpointing=False",
-            "enable_goodput_recording=False",
-            "enable_checkpoint_cloud_logger=False",
-            "monitor_goodput=False",
-            f"metrics_file={os.path.join(outputs_dir, 'metrics.json')}",
-        ]
-    )
+    train_main([
+        None,
+        os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
+        f"base_output_directory={test_tmpdir}",
+        "run_name=ragged_dot_smoke_test",
+        "base_emb_dim=128",
+        "base_num_query_heads=4",
+        "base_num_kv_heads=4",
+        "base_mlp_dim=128",
+        "base_moe_mlp_dim=128",
+        "base_num_decoder_layers=8",
+        "head_dim=128",
+        # TODO(b/441100085): When changing the decoder_block we might
+        # need to adjust the tiling.
+        "decoder_block=deepseek",
+        "attention_type=mla",
+        "num_experts=2",
+        # Enable sparse_matmul.
+        "sparse_matmul=True",
+        # Enable ragged_dot.
+        "megablox=False",
+        f'quantization="{quantization}"',
+        "use_qwix_quantization=True",
+        "per_device_batch_size=2",
+        "max_target_length=1024",
+        "dataset_type=synthetic",
+        "steps=10",
+        "enable_checkpointing=False",
+        "enable_goodput_recording=False",
+        "enable_checkpoint_cloud_logger=False",
+        "monitor_goodput=False",
+        f"metrics_file={os.path.join(outputs_dir, 'metrics.json')}",
+    ])
 
 
 if __name__ == "__main__":