Enable quantized multi-token-prediction

[shuningjin]

Description

Enables FP8 quantization support for Multi-Token Prediction (MTP) layers via Qwix interception.

Bug: b/552423947

Changes:

• Config control: (base.yml), Added quantize_mtp flag (default: false). Requires mtp_num_layers > 0 and quantization='fp8_full'
• Qwix rule update: (quantizaiton.py), In get_fp8_full_qwix_rule_w_sparsity, expanded the rule target module path to include (decoder/.*layers.*|mtp_block/.*) when quantize_mtp=True.
• Added unit test for MTP interception (multi_token_prediction_test.py): MTPQwixInterceptionTest

Additional Change:

• Restore gmm rule to previous behavior, Context
• How it relates here: When fp8_full, main model's MoE blocks match Qwix rule. MTP's MoE blocks can have no Qwix rule if quantize_mtp=False.

More Context - model structure: log

python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext model_name=deepseek3-671b scan_layers=true mtp_num_layers=1 enable_nnx=true 

Tests

Unit Tests

Verified Qwix interception tracing for mtp_block dot_general ops under CPU abstract evaluation:

JAX_PLATFORMS=cpu pytest tests/unit/multi_token_prediction_test.py -k "MTPQwixInterceptionTest"

Auxiliary Test

Qwix log: MTP interception correct. script, log

# quantize_mtp=true
[QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/out' op=dot_general0 rule=0
# quantize_mtp=false
[QWIX] module='mtp_block/mtp_layer_1/mtp_1_transformer_layer/self_attention/out' op=dot_general0 rule=None

E2E Validation

• Profile: dot general and gmm is quantized in MTP
• Pre-training convergence of DeepSeek-V3: FP8 baseline + no-MTP / quantized-MTP reach target eval loss in the same steps (See b/552423947#comment5)

Checklist

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

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links/descriptions above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[gemini-code-assist]

Code Review

This pull request introduces support for quantizing Multi-Token Prediction (MTP) layers in MaxText when using FP8 quantization (fp8_full). It adds a new configuration parameter quantize_mtp (defaulting to false), along with validation rules, updates to Qwix quantization rules to target the mtp_block module path, and handling for cases where GMM quantization rules are absent. Additionally, a verification script and unit tests have been added to ensure correct interception behavior. There are no review comments, so I have no feedback to provide.

[codecov]

Codecov Report

❌ Patch coverage is 60.00000% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/kernels/megablox/ops.py	0.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!