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
  • gmm_type={megablox, tokamax gmm1, tokamax gmm2} x quant_scheme={bf16, dynamic, fp8_static_wa}

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

local v5p-8

Unit test: smoke train

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

log

Unit test: value and gradient equivalence

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

log

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 68.42105% with 54 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	0.00%	1 Missing and 1 partial ⚠️
src/maxtext/layers/decoders.py	0.00%	1 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	0.00%	1 Missing ⚠️
src/maxtext/models/deepseek.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[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]

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

[shuningjin]

Current quantization tests do not have correctness check for value or gradient. Quantization would inevitably introduce precision loss.

• For unit test, we can only have a loose check that gaps are not too large. It is hard to determine a proper threshold. Ultimately, we would need E2E to determine the quality.
• In this specific case, since we know what megablox or tokamax gmm1 was previously successful, we can compare the diff of tokamx gmm2 against them.

I expanded the test cases, here

• gmm_type={megablox, tokamax gmm1, tokamax gmm2} x quant_scheme={bf16, dynamic, fp8_static_wa }, where fp8_dynamic=absmax for weight/activation/grad and fp8_static_wa=fixed weight+fixed activation+absmax grad
• diff is against bf16 dense, log from my v5p-8

Observations:

• For the same gmm type: bf16 < fp8_dymanic < fp8_static_wa. This is expected.
• For the same quantization scheme, different gmm has similar diff. This is expected.

1. relative_norm_diff Across All Cases

Test Case	output	state_grad	grad_gate-kernel	grad_wi_0	grad_wi_1	grad_wo
megablox_bf16_ep1	0.004347	0.007798	0.007537	0.007493	0.007523	0.006855
megablox_fp8_dynamic_ep1	0.064479	0.064210	0.022872	0.100576	0.100811	0.101557
megablox_fp8_static_ep1	0.153188	0.049933	0.031198	0.189635	0.190141	0.215808
tokamax_v1_bf16_ep1	0.004347	0.007794	0.007531	0.007492	0.007523	0.006855
tokamax_v1_fp8_dynamic_ep1	0.064479	0.064168	0.022875	0.100574	0.100811	0.101557
tokamax_v1_fp8_static_ep1	0.153188	0.049933	0.031198	0.189635	0.190141	0.215808
tokamax_v2_bf16_ep1	0.004347	0.007798	0.007537	0.007493	0.007523	0.006855
tokamax_v2_fp8_dynamic_ep1	0.064551	0.064635	0.023685	0.100734	0.100966	0.101706
tokamax_v2_fp8_static_ep1	0.153188	0.049926	0.031198	0.189628	0.190157	0.215808
tokamax_v2_bf16_ep4	0.004347	0.007798	0.007537	0.007099	0.007130	0.006371
tokamax_v2_fp8_dynamic_ep4	0.064549	0.065041	0.023742	0.100691	0.100953	0.101660
tokamax_v2_fp8_static_ep4	0.153188	0.049926	0.031198	0.189619	0.190156	0.215799

2. max_abs_diff Across All Cases

Test Case	output	state_grad	grad_gate-kernel	grad_wi_0	grad_wi_1	grad_wo
megablox_bf16_ep1	2.44e-04	1.64e-11	5.95e-09	2.33e-10	1.75e-10	5.89e-11
megablox_fp8_dynamic_ep1	3.05e-03	1.04e-10	1.97e-08	1.11e-09	9.31e-10	3.16e-10
megablox_fp8_static_ep1	7.39e-03	8.37e-11	2.55e-08	2.62e-09	2.44e-09	6.86e-10
tokamax_v1_bf16_ep1	2.44e-04	1.64e-11	5.95e-09	2.33e-10	1.75e-10	5.89e-11
tokamax_v1_fp8_dynamic_ep1	3.05e-03	1.04e-10	1.97e-08	1.11e-09	9.31e-10	3.16e-10
tokamax_v1_fp8_static_ep1	7.39e-03	8.37e-11	2.55e-08	2.62e-09	2.44e-09	6.86e-10
tokamax_v2_bf16_ep1	2.44e-04	1.64e-11	5.95e-09	2.33e-10	1.75e-10	5.89e-11
tokamax_v2_fp8_dynamic_ep1	3.17e-03	1.05e-10	2.18e-08	1.08e-09	1.05e-09	3.26e-10
tokamax_v2_fp8_static_ep1	7.39e-03	8.37e-11	2.55e-08	2.62e-09	2.44e-09	6.86e-10
tokamax_v2_bf16_ep4	2.44e-04	1.64e-11	5.95e-09	1.75e-10	1.16e-10	4.43e-11
tokamax_v2_fp8_dynamic_ep4	3.36e-03	1.05e-10	2.18e-08	1.11e-09	1.06e-09	3.48e-10
tokamax_v2_fp8_static_ep4	7.39e-03	8.37e-11	2.55e-08	2.68e-09	2.44e-09	6.43e-10

[RissyRan]

Thanks for the comprehensive tests!

[tabakg]

LGTM thank you!

One small suggestion for static scales relative error test: inspect the abs-max and insert appropriate scaling values in the test. This way we can avoid excessively loose bounds.

[suhithr]

Thank you! LGTM