Support quantization for tokamax gmm v2

docs/guides/optimization/benchmark_and_performance.md

@@ -191,7 +191,7 @@ Different quantization recipes are available, including` "int8", "fp8", "fp8_ful
 
 For v6e and earlier generation TPUs, use the "int8" recipe. For v7x and later generation TPUs, use "fp8_full". GPUs should use “fp8_gpu” for NVIDIA and "nanoo_fp8" for AMD.
 
-See [](quantization-doc).
+See [Quantization guide](../../reference/core_concepts/quantization.md).
 
 ### Choose sharding strategy
 

docs/reference/models/supported_models_and_architectures.md

@@ -12,7 +12,7 @@ MaxText is an open-source, high-performance LLM framework written in Python/JAX.
 
 - **Supported Precisions**: FP32, BF16, INT8, and FP8.
 - **Ahead-of-Time Compilation (AOT)**: For faster model development/prototyping and earlier OOM detection.
-- **Quantization**: Via **Qwix** (recommended) and AQT (deprecated). See the [Quantization guide](quantization-doc).
+- **Quantization**: Via **Qwix** (recommended) and AQT (deprecated). See the [Quantization guide](../core_concepts/quantization.md).
 - **Diagnostics**: Simple logging via `max_logging`, profiling in **XProf**, and visualization in **TensorBoard**.
 - **Multi-Token Prediction (MTP)**: Enables token efficient training with multi-token prediction.
 - **Elastic Training**: Fault-tolerant and dynamic scale-up/scale-down on Cloud TPUs with Pathways.

src/maxtext/common/checkpointing.py

@@ -1065,6 +1065,9 @@ def _should_save_checkpoint_at_step(checkpoint_manager, step, config, force):
   """Returns whether MaxText should build and dispatch checkpoint args."""
   if force:
     return True
+  if step == 0 and not config.save_checkpoint_on_start:
+    # if step = 0, `step % config.checkpoint_period == 0` is always true, force skip
+    return False
   if config.enable_continuous_checkpointing:
     base_checkpoint_due = bool(checkpoint_manager.should_save(step))
   else:

src/maxtext/configs/base.yml

@@ -53,6 +53,7 @@ load_full_state_path: ""
 # async_checkpointing is true, else a synchronous one is used. If you have
 # problems with the checkpointer we recommend trying the synchronous one.
 enable_checkpointing: true
+save_checkpoint_on_start: true
 save_checkpoint_on_completion: true
 async_checkpointing: true
 checkpoint_period: 10_000
@@ -268,6 +269,16 @@ wo_tile_drhs_mlp_dim: 1024
 
 merge_gating_gmm: false
 
+# Use tokamax library for gmm kernel implementation
+use_tokamax_gmm: false
+# Whether to use GMM v2 for MoE. (Requires use_tokamax_gmm: true)
+# fwd: tokamax_v2_gmm if true, else tokamax_v1_gmm
+use_gmm_v2_fwd: false
+# bwd_dlhs: tokamax_v2_gmm if true, else tokamax_v1_gmm
+use_gmm_v2_dlhs: false
+# bwd_drhs: tokamax_v2_tgmm if true, else tokamax_v1_tgmm
+use_gmm_v2_drhs: false
+
 norm_topk_prob: false # boolean to enable the top-k probability normalization. qwen3-specific normalization of router weights.
 
 # when moe weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls
@@ -1260,11 +1271,6 @@ use_qk_norm_in_gdn: true
 # The ratio of dimension to apply ROPE on
 partial_rotary_factor: 1.0
 
-# Use tokamax library for gmm kernel implementation
-use_tokamax_gmm: false
-# Whether to use GMM v2 for MoE.
-# Requires use_tokamax_gmm: true. Currently incompatible with quantization.
-use_gmm_v2: false
 use_tokamax_splash: false
 # Setting this flag will use a non-pallas implementation.
 use_jax_splash: false

src/maxtext/configs/pyconfig.py

@@ -39,7 +39,8 @@
 from maxtext.utils import max_logging
 
 logger = logging.getLogger(__name__)
-logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))
+# TODO: revert before merge
+# logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))
 
 _BASE_CONFIG_ATTR = "base_config"
 _MAX_PREFIX = "M_"

src/maxtext/configs/types.py

@@ -352,6 +352,7 @@ class Checkpointing(BaseModel):
   source_checkpoint_layout: Literal["orbax", "safetensors", "safetensors_dynamic"] = Field(
       "orbax", description="The layout of the source checkpoint to load."
   )
+  save_checkpoint_on_start: bool = Field(True, description="If True, saves an initial checkpoint upon training start.")
   save_checkpoint_on_completion: bool = Field(
       True, description="If True, saves a final checkpoint upon training completion."
   )
@@ -444,7 +445,7 @@ class Quantization(BaseModel):
   use_qwix_quantization: bool = Field(False, description="Whether to use qwix for quantization.")
   use_manual_quantization: bool = Field(
       False,
-      description="Whether to use manual quantization for batch split. Only used if use_batch_split_schedule is True.",
+      description="Whether to use manual quantization for batch split. Only used if `use_batch_split_schedule=True`.",
   )
   weight_quantization_calibration_method: str = Field(
       "absmax",
@@ -616,20 +617,6 @@ class Attention(BaseModel):
   use_post_attn_norm: bool = Field(False, description="Apply LayerNorm after the attention block.")
   use_post_ffw_norm: bool = Field(False, description="Apply LayerNorm after the feed-forward block.")
   use_ragged_attention: bool = Field(False, description="Whether to use ragged attention kernels.")
-  use_tokamax_gmm: bool = Field(
-      False,
-      description="Whether to use the Tokamax library for GMM kernel implementation.",
-  )
-  use_gmm_v2: bool = Field(
-      False,
-      description=(
-          "Whether to use GMM v2 (with bf16 activations and weights) for MoE."
-          " Requires use_tokamax_gmm: true. Currently incompatible with quantization."
-      ),
-  )
-  num_moe_emb_chunks: int = Field(
-      0, description="Number of chunks for overlapping token all-gather and GMM computation along embedding dimension."
-  )
   ragged_block_size: int = Field(256, description="Block size for ragged attention.")
   enable_padding_causal_mask: bool = Field(True, description="Temporary flag for Transformer Engine padding.")
   use_tokamax_splash: bool = Field(False, description="Whether to use tokamax splash attention.")
@@ -969,6 +956,27 @@ class MoEKernels(BaseModel):
 
   merge_gating_gmm: bool = Field(False, description="whether to merge the two gating gmm kernels into one.")
 
+  # tokamax gmm
+  use_tokamax_gmm: bool = Field(
+      False,
+      description="Whether to use the Tokamax library for GMM kernel implementation.",
+  )
+  use_gmm_v2_fwd: bool = Field(
+      False,
+      description="Whether to use GMM v2 for MoE forward pass. (Requires use_tokamax_gmm: true)",
+  )
+  use_gmm_v2_dlhs: bool = Field(
+      False,
+      description="Whether to use GMM v2 for MoE backward pass (dlhs). (Requires use_tokamax_gmm: true)",
+  )
+  use_gmm_v2_drhs: bool = Field(
+      False,
+      description="Whether to use GMM v2 for MoE backward pass (drhs). (Requires use_tokamax_gmm: true)",
+  )
+  num_moe_emb_chunks: int = Field(
+      0, description="Number of chunks for overlapping token all-gather and GMM computation along embedding dimension."
+  )
+
 
 class DeepSeekMoE(BaseModel):
   """Configuration specific to DeepSeek-style MoE layers."""
@@ -2691,6 +2699,7 @@ def validate_num_moe_emb_chunks(self):
     Validates that num_moe_emb_chunks is used with supported settings.
     """
     if self.num_moe_emb_chunks > 0:
+      # TODO: update self.use_gmm_v2
       if not self.use_gmm_v2 or not self.use_ring_of_experts:
         raise ValueError(
             f"num_moe_emb_chunks > 0 requires use_gmm_v2=True and use_ring_of_experts=True. "
@@ -2946,6 +2955,10 @@ def get_num_target_devices():
             "Please migrate to Qwix by setting use_qwix_quantization=True."
         )
 
+    # Check quant config is non-empty for Qwix quantization
+    if self.use_qwix_quantization and not self.quantization:
+      raise ValueError("Qwix quantization is enabled but quantization is not set.")
+
     # Default quantization sharding count to number of local devices if not set.
     if self.quantization_local_shard_count == -1:
       try:
@@ -3573,11 +3586,36 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
     if self.share_kv_projections and self.attention_type == "mla":
       raise ValueError("`share_kv_projections` is not compatible with `attention_type='mla'`.")
 
-    if self.use_gmm_v2 and (self.quantization or self.use_qwix_quantization):
-      raise ValueError("Quantization with GMM v2 is not supported yet.")
-    if self.use_gmm_v2 and not self.use_tokamax_gmm:
+    if self.use_manual_quantization and not self.use_batch_split_schedule:
+      raise ValueError("manual quantization is only used when `use_batch_split_schedule=True`.")
+
+    if (self.use_gmm_v2_fwd or self.use_gmm_v2_dlhs or self.use_gmm_v2_drhs) and not self.use_tokamax_gmm:
       raise ValueError("GMM v2 requires `use_tokamax_gmm=true`.")
 
+    # Ensure GMM v2 flags (fwd, dlhs, drhs) are set to a valid combination:
+    # - v1+v1+v1 (False, False, False)
+    # - v2+v2+v2 (True, True, True)
+    # - v2+v1+v2 (True, False, True)
+    gmm_v2_combo = (
+        self.use_gmm_v2_fwd,
+        self.use_gmm_v2_dlhs,
+        self.use_gmm_v2_drhs,
+    )
+    valid_combos = {
+        (False, False, False),  # 111 (v1+v1+v1)
+        (True, True, True),  # 222 (v2+v2+v2)
+        (True, False, True),  # 212 (v2+v1+v2)
+    }
+    if gmm_v2_combo not in valid_combos:
+      raise ValueError(
+          "Invalid GMM v2 configuration combination (fwd, dlhs, drhs). "
+          "Allowed combinations are:\n"
+          "  - v1+v1+v1 (False, False, False)\n"
+          "  - v2+v2+v2 (True, True, True)\n"
+          "  - v2+v1+v2 (True, False, True)\n"
+          f"But got: {gmm_v2_combo}"
+      )
+
     for val in self.compress_ratios:
       if val != 0 and val < 4:
         raise ValueError(f"compress_ratio must be 0 (disabled) or >= 4, got {val}")