Support quantization for tokamax gmm v2

[shuningjin]

Description

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

Bug: b/534817371

What this does

1. Integrate quantization in gmm v2

• base.yml: use_gmm_v2=true should work for quantization
• src/maxtext/kernels/megablox/ops.py: quantization and scaling integration of gmm2 / tgmm2

2. Add unit tests for quantized gmm

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

3. Significant refactor of ops.py: modularize GMM fwd and bwd functions

• Previously, _gmm_fwd and _gmm_bwd were large, monolithic functions containing complex branching and quantization logic. This is difficult to read and maintain.
• Breaks them down into smaller, domain-specific sub-functions, significantly improving code readability and maintainability. Execution logic and runtime behavior remain identical, except for the aforementioned gmm v2 change.

4. Other changes

• Improve robustness of quantization. For example, gmm checks agaisnt empty rule to ensure quantization is applied; refactor condition use_qwix_quantization to be co-occur with quantization; fix broken link for quantization doc.
• Add save_checkpoint_on_start to allow skipping for testing
• update doc for moe configuration: add use_gmm_v2

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

We tested in b/534817371, convergence on par with bf16 for fsdp+ep or fsdp.

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.

[codecov]

Codecov Report

❌ Patch coverage is 69.41176% with 52 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/kernels/megablox/ops.py	72.48%	24 Missing and 17 partials ⚠️
src/maxtext/layers/quantizations.py	28.57%	5 Missing ⚠️
src/maxtext/models/deepseek_batchsplit_fp8.py	0.00%	3 Missing ⚠️
src/maxtext/common/checkpointing.py	50.00%	1 Missing ⚠️
src/maxtext/layers/decoders.py	0.00%	1 Missing ⚠️
src/maxtext/models/deepseek.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[CaptainO5]

lhs needs to be quantized here, since we are skipping it in the fwd pass when gmm_v2_fwd is enabled.

Suggested change

	if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray):
	lhs = qpl.quantize(
	lhs,
	quantization_rule.act_qtype,
	channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0],
	calibration_method=quantization_rule.act_calibration_method,
	)

[shuningjin]

Thanks for the catch! I have corrected the asdtype to qpl.quantize for lhs quantization in bwd.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[RissyRan]

LGTM, just minor comments, related to tolerance.

[RissyRan]

nit: tokamax_v2_bf16_ep, and similar to test name bellow

[RissyRan]

Do you need sharding_tolerance in your end-to-end run? If so, could we create a bug to track which part is unsharded. I think we should debug this with

maxtext/src/maxtext/configs/base.yml

Line 1077 in 17c7172

debug_sharding: false # Prints model weights sharding info

[RissyRan]

The tolerance seems high, and is this aligned with other existing quantization tests?