Support quantization for tokamax gmm v2 by shuningjin · Pull Request #4451 · AI-Hypercomputer/maxtext

shuningjin · 2026-07-13T21:54:27Z

Description

Integrate quantization of Tokamax GMM V2. This is joint effort with @CaptainO5, built on top of PR#4197.

Fix: b/534817371

What this does

Integrate quantization in gmm v2

base.yml: introduce flags use_gmm_v2_fwd, use_gmm_v2_dlhs, use_gmm_v2_drhs:
- pure mode: fwd_gmm2 + dlhs_gmm2 + drhs_tgmm2
- hybrid mode: fwd_gmm2 + dlhs_gmm1 + drhs_tgmm2, which may have better perf for fp8
src/maxtext/kernels/megablox/ops.py: quantization and scaling integration of gmm2 / tgmm2

Add unit tests for quantized gmm

tests.integration.tokamax_test: smoke train
tests.unit.moe_test: value and gradient equivalence

Other changes

Improve robustness of quantization. For example, gmm checks agaisnt empty rule to ensure quantization is applied;use_qwix_quantization only when quantization is enabled.
Add save_checkpoint_on_start to allow skipping for testing

Tests

Unit test

Unit test: smoke train

python3 -m pytest -v --pyargs tests.integration.tokamax_test -rP -s

https://paste.googleplex.com/6354648482054144

Unit test: value and gradient equivalence

python3 -m pytest -v --pyargs tests.unit.moe_test -rP -s -k "test_gmm_grad_equivalence"

https://paste.googleplex.com/4962768028565504

E2E test

Detail: b/534817371#comment3

fsdp convergence - v212+fp8 / v222+fp8 compare with v111+fp8: loss is close (b/518725224#comment36, run5, run6)
fsdp convergence - v212+fp8 compare with v1111+bf16: reach target eval loss in similar step,
ep+fsdp convergence: v212+fp8 compare with v222+bf16: b/532148952#comment3

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

CaptainO5 · 2026-07-15T23:27:13Z

-      False,
-      description="Whether to use the Tokamax library for GMM kernel implementation.",
-  )
-  use_gmm_v2: bool = Field(


Can we still have this config that sets all fwd, dlhs and drhs to true in types.py when true?

github-actions · 2026-07-16T02:12:33Z

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions

## 📋 Review Summary

This pull request introduces quantization support for Tokamax GMM V2, enabling memory-efficient Mixture of Experts (MoE) operations with robust performance profiles. The implementation is highly structured and includes comprehensive value/gradient equivalence unit tests and integration smoke training. While the overall quality of the code is high, a critical logic issue in the backward-pass selection needs to be resolved prior to merging.

🔍 General Feedback

Comprehensive Testing: Excellent inclusion of both value/gradient equivalence checks and TPU integration smoke tests. This is a very robust testing strategy.
Robust Parameter Validation: The added configuration checks in types.py are thorough and prevent invalid parameter combinations cleanly.
Temporary Code Cleanup: There are a couple of left-over development configurations and TODOs that should be reverted before finalizing the merge.

github-actions · 2026-07-16T02:17:17Z

+# TODO: revert before merge
+from absl import logging
+logging.set_verbosity(logging.INFO)


🟢 This temporary absl logging override and its accompanying TODO should be removed before merge.

Suggested change

# TODO: revert before merge

from absl import logging

logging.set_verbosity(logging.INFO)

github-actions · 2026-07-16T02:17:17Z

 import qwix.pallas as qpl
 import tokamax
+from tokamax._src.ops.ragged_dot import pallas_mosaic_tpu_kernel as tokamax_backend



🟡 Importing from the internal module `tokamax._src.ops.ragged_dot` is a maintainability risk since private APIs (anything under `_src`) do not guarantee backward compatibility and can change in any release. If there is no public equivalent, please document this as known technical debt.

github-actions · 2026-07-16T02:17:17Z

+# TODO: revert before merge
+# logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))


🟡 The logger level restriction was commented out during development with a "revert before merge" TODO. Leaving this in the codebase can lead to verbosity issues or inconsistent logger behavior in production. Please restore the logger level setup and remove the TODO.

Suggested change

# TODO: revert before merge

# logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))

logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))

github-actions · 2026-07-16T02:17:17Z

          implementation="mosaic",
          **dlhs_kwargs,
      )
+    if not use_gmm_v2_dlhs and use_gmm_v2_fwd:


🔴 There is a critical logic bug in the `dlhs` path selection. Because the second condition is checked using a separate `if` statement rather than an `elif`, when both `use_gmm_v2_dlhs` and `use_gmm_v2_fwd` are `False` (which is the default v1 configuration), Python will enter the first block, but then also execute the `else` block of the second `if`. This will overwrite the computed `dlhs` using GMM v2 instead of leaving it as GMM v1, breaking backward compatibility and potentially leading to compilation failures, hardware crashes, or silent gradient corruption.
Please change the second if to elif to make the conditional branches mutually exclusive.

Suggested change

if not use_gmm_v2_dlhs and use_gmm_v2_fwd:

elif not use_gmm_v2_dlhs and use_gmm_v2_fwd:

Thanks for catching this! I have updated. It previously affected v1 path. My experiments with current head is around v2+v1+v2 or v2+v2+v2. Therefore, they may not be affected.

github-actions · 2026-07-16T02:17:17Z

+    raise ValueError("Reference and actual pytrees must have the same structure.")
+
+  log_lines = ["CALCULATE DIFF"]
+  for (path, a), (_, b) in zip(leaves_a, leaves_b):


🟡 If any leaf in the pytree is `None` (which often happens for non-differentiable parameters or state variables where the gradient is `None`), calling `a.dtype` will raise an `AttributeError`. Please consider handling `None` leaves gracefully by asserting that both are `None` and continuing.

Suggested change

for (path, a), (_, b) in zip(leaves_a, leaves_b):

if a is None or b is None:

assert a is None and b is None, f"Structure mismatch: one leaf is None at {path}"

continue

assert a.dtype in (jnp.float32, jnp.bfloat16, jnp.float16)

RissyRan

Thanks for the change!

RissyRan · 2026-07-16T04:23:18Z


+# Use tokamax library for gmm kernel implementation
+use_tokamax_gmm: false
+# Whether to use GMM v2 for MoE. (Requires use_tokamax_gmm: true)


Wondering if we should simplify back to use_gmm_v2 config for users? What's the gain for using different combination?

Also, could you help update this file when applies?

maxtext/docs/reference/core_concepts/moe_configuration.md

Line 4 in fe5032a

Licensed under the Apache License, Version 2.0 (the "License");

RissyRan · 2026-07-16T04:24:54Z

+      False,
+      description="Whether to use GMM v2 for MoE backward pass (drhs). (Requires use_tokamax_gmm: true)",
+  )
+  num_moe_emb_chunks: int = Field(


RissyRan · 2026-07-16T04:34:24Z

+
+  if scale.ndim == 2:  # Per-Channel quantization
+    rhs_scale = jnp.expand_dims(scale, axis=(1, 2))
+  elif scale.ndim == 3:  # Block-wise quantization


Does this mean we could support block-wise?

RissyRan · 2026-07-16T04:36:29Z

 # Forked from:
-# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/
-# tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py
+# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py


When will we deprecate and move?

CaptainO5 · 2026-07-16T08:44:01Z

+          transpose_rhs=not transpose_rhs,
+          interpret=interpret,
+          input_buffer_count=input_buffer_count[1],
+      )


Suggested change

)

**dlhs_kwargs,

)

CaptainO5 · 2026-07-16T08:53:27Z

-    if self.use_gmm_v2 and not self.use_tokamax_gmm:
+    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`.")



Shall we also add a check to disable use_gmm_v2* and use_manual_quantization?

Support quantization for tokamax gmm v2

f814413

shuningjin force-pushed the gmm2_quantize branch from 2f59d6c to f814413 Compare July 15, 2026 07:48

shuningjin added the pull ready label Jul 15, 2026

shuningjin added 2 commits July 15, 2026 08:18

format

2cec661

clean test

b135e8e

shuningjin changed the title ~~[draft] gmm v2 quantization~~ Support quantization for tokamax gmm v2 Jul 15, 2026

shuningjin marked this pull request as ready for review July 15, 2026 22:05

shuningjin removed the pull ready label Jul 15, 2026

shuningjin assigned RissyRan, NuojCheng, CaptainO5 and gobbleturk Jul 15, 2026

CaptainO5 requested changes Jul 15, 2026

View reviewed changes

RissyRan added the gemini-review label Jul 16, 2026

github-actions Bot reviewed Jul 16, 2026

View reviewed changes

fix

f9f37b8

shuningjin unassigned RissyRan, CaptainO5, gobbleturk and NuojCheng Jul 16, 2026

RissyRan reviewed Jul 16, 2026

View reviewed changes

CaptainO5 reviewed Jul 16, 2026

View reviewed changes

	# TODO: revert before merge
	from absl import logging
	logging.set_verbosity(logging.INFO)

		# TODO: revert before merge
		# logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))

	# TODO: revert before merge
	# logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))
	logger.setLevel(os.environ.get("LOGLEVEL", "INFO"))

	if not use_gmm_v2_dlhs and use_gmm_v2_fwd:
	elif not use_gmm_v2_dlhs and use_gmm_v2_fwd:

-  for (path, a), (_, b) in zip(leaves_a, leaves_b):
+    if a is None or b is None:
+      assert a is None and b is None, f"Structure mismatch: one leaf is None at {path}"
+      continue
+    assert a.dtype in (jnp.float32, jnp.bfloat16, jnp.float16)

Uh oh!

Conversation

shuningjin commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Unit test

E2E test

Checklist

Uh oh!

CaptainO5 Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented Jul 16, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

github-actions Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

shuningjin commented Jul 13, 2026 •

edited

Loading

CaptainO5 Jul 15, 2026 •

edited

Loading