[qat] feat: MXFP4 weight-only QAT for MLite MoE models

[ISEEKYAN]

MXFP4 weight-only QAT for Megatron Lite MoE models

Scope

This PR aligns MXFP4 weight-only QAT in Megatron Lite with the quantization
contract used by rollout and composes the path into Qwen3 MoE, Qwen3.5,
DeepSeek V4, GLM-5, and Kimi K2.

• BF16 master weights remain optimizer-owned while training forwards apply
  differentiable MXFP4 fake quantization.
• Checkpoint loading maps logical ...weight names onto the surviving
  ...parametrizations.weight.original master instead of silently dropping
  tensors.
• Qwen3 MoE export uses the FP4-aware resynchronization path consumed by
  rollout.
• All five model protocols apply QAT before optimizer construction and use a
  shared checkpoint-key canonicalization contract.
• Router replay attaches only to decoder routers. MTP-only routers are
  excluded, and GLM-5 route and replay-mask packing share its contiguous CP layout.
• Replay fails loudly if router auxiliary loss is nonzero because native
  routing-map load statistics do not describe replay dispatch.
• A router-free replay request fails loudly instead of silently doing nothing.

The quantization primitive remains model-agnostic and contains no model-name
allowlist. The verified delivery surface described here is MXFP4.

MXFP4 numerical alignment

Before this change, training fake quantization, checkpoint serialization, and
rollout used different scale and tie-rounding rules. The old training rule
selected a scale 2x too small whenever a block maximum's mantissa exceeded
1.5, clipping the largest value in each affected block.

Measured over 99,090,432 real Qwen3-30B-A3B weight elements
(3,096,576 blocks), the old rule produced a different scale for 43.07% of
blocks and a different decoded value for 12.38% of elements.

The implementation now shares one MXFP4 scale-selection and E2M1-rounding
contract between training fake quantization and serialization, with an
independent ModelOpt transcription as the reference. On the same real-weight
sample:

• decoded-value differences: 0
• maximum absolute difference: 0
• scale-exponent mismatches: 0
• clipped elements: 0
• packed-byte match fraction: 1.0000
• scale-byte match fraction: 1.0000

This is weight-level parity. Training and rollout still use different GEMM
implementations, so it is not an end-to-end equality claim.

Five-model QAT and R3 wiring

Model	QAT composition	Checkpoint mapping	Router replay
Qwen3 MoE	before optimizer construction	shared canonical key mapping	decoder layers only; MTP excluded
Qwen3.5	before optimizer construction	shared canonical key mapping	decoder layers only; MTP excluded
DeepSeek V4	before optimizer construction	shared canonical key mapping	decoder layers only; MTP excluded
GLM-5	before optimizer construction	shared canonical key mapping	shared sigmoid router; contiguous route/mask CP packing; MTP excluded
Kimi K2	before optimizer construction	shared canonical key mapping	shared sigmoid router; MTP excluded

The R3 correction prevents training-only MTP routers from consuming rollout
route entries. With MTP disabled, replay still attaches to every decoder layer;
with MTP enabled, the attachment count remains the decoder-layer count rather
than decoder plus MTP layers.

Four-arm DAPO experiment

All arms were planned for 30 steps but were truncated by a four-hour scheduler
limit: BF16 ended at step 21, QAT-off at step 20, QAT-on at step 22, and R3 at
step 24. These are early-window results.

Arm	Training	Rollout	Router replay
baseline	BF16	BF16	off
QAT off	BF16	MXFP4	off
QAT on	MXFP4 fake quant	MXFP4	off
R3	MXFP4 fake quant	MXFP4	on

QAT-off and QAT-on used byte-identical rollout settings; their only QAT
difference was training-side enablement. The rollout path was vLLM
compressed-tensors plus verl.utils.qat.vllm_patch.

Curves

Reward	Entropy

Train/rollout probability deviation	Validation

The images live on a non-merge asset branch; this PR contains no PNG files.

Measured results

Arm	Last-step probs_diff mean	Pearson	Entropy
BF16	0.00598	0.99923	0.4920
QAT on	0.01657	0.99527	0.5527
R3	0.02443	0.98060	0.3917
QAT off	0.02672	0.99204	1.7349

QAT-on reduced the last-step train/rollout probability discrepancy by about
38% versus QAT-off (0.0166 versus 0.0267), but remained about 2.8x
BF16 (0.00598). The result is therefore that MXFP4 QAT reduces the discrepancy,
not that it eliminates it.

MXFP4 rollout cost about 30% less per generated token than BF16 in this
window. QAT-off and QAT-on differed by only 0.4% after token normalization
(0.04072 versus 0.04090 seconds per token).

Validation was measured only at global steps 4, 9, 14, and 19. The observed
within-run noise floor was about ±0.1, so these points do not support ranking
the arms by validation score.

Two observations remain unexplained and no mechanism is claimed:

1. QAT-off entropy was 1.7349, 3.1-4.4x the other three arms, and its
   validation series trended downward.
2. R3 had worse train/rollout consistency than QAT-on, with the synchronized
   probs_diff gap growing from about 11% to 44% over steps 13-22.

Known limitations and follow-ups

• The R3 layer axis covers decoder routers only. MTP routers remain on native
  routing; replaying MTP or speculative routes has not been validated.
• Route rows omitted for final input tokens are zero-filled placeholders.
  Correctness relies on the corresponding r3_replay_mask rows being false;
  zero is not a sentinel expert ID.
• The 99,090,432-element MXFP4 result is offline weight-level evidence reported
  here, not an end-to-end training/rollout parity test or a committed full-model
  test artifact.
• Replay evidence is printed for the first observed microbatch only.
  changed=0 emits a warning rather than failing, and replay counters are not
  exported as training metrics.
• QAT/R3 operational guidance is documented, but a dedicated entry in the
  MLite skill registry is still missing.

Usage

experimental/lite/examples/verl/QAT.md
documents the MXFP4 training contract, apply-before-optimizer ordering,
checkpoint mapping, safe exclusions, packed snapshots, and the exporter
boundary.

experimental/lite/examples/verl/scripts/run_qwen3moe_mxfp4_qat.sh
provides the four experiment modes and keeps the QAT delta separate from
ordinary training parameters.

Author: Yan Bai

Coverage matrix

Model	Training-side QAT	MXFP4 export	R3
Qwen3 MoE	yes	yes	yes
Qwen3.5	yes	yes	yes
DeepSeek V4	yes	yes	yes
GLM-5	yes	yes	yes
Kimi K2	yes	yes	yes

MXFP4 export now covers all five models. For Qwen3.5, GLM-5, and Kimi K2, the exporter validates the resync target, preserves embeddings, lm_head, and router gates, and emits primitive-identical packed weights plus UE8M0 scales.