[draft] gmm v2 quantization

[shuningjin]

Description

Integrate quantization of Tokamax GMM V2. This is joint effort with @CaptainO5

What this does

1. Fork newer changes from tokamax gmm v2: in particular tgmm has rhs_scale support. gmm, tgmm
2. Integrate quantization in gmm v2

• flags use_gmm_v2_fwd, use_gmm_v2_dlhs, use_gmm_v2_drhs
• src/maxtext/kernels/megablox/ops.py

3. Add unit tests
4. 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.

Tests

Unit test: smoke train

pytest -v -m integration_test tests/integration/tokamax_test.py

Unit test: gradient equivalence

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

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.