[sophiayyya]

What does this PR do?

MXFP4 QAT precision and data flow

This change adds routed-expert weight-only MXFP4 fake quantization on top of Transformer Engine (TE) FP8 training for DeepSeek-V4-Flash-0731. Checkpoint storage, trainable parameter dtype, and GEMM precision are separate:

Module / tensors	Official 0731 checkpoint storage	Parameters after Bridge import	Added ModelOpt QAT
Routed experts: w1, w2, w3	MXFP4 E2M1, two values per I8 byte; one E8M0 scale per 32 contiguous input-channel elements	BF16	MXFP4 weight fake-quant on mlp.experts.linear_fc1/linear_fc2
Shared experts; MLA wq_a, wq_b, wkv, wo_a, wo_b; indexer wq_b	FP8 E4M3 with E8M0 scales over 128 × 128 weight blocks	BF16	Disabled
Embeddings, prediction head, router, norms; compressor projections and indexer weights_proj	BF16	BF16	Disabled
Native FP32 tensors, including mHC hc_attn_fn/hc_ffn_fn, attn_sink, and compressor ape	FP32	FP32 retained	Disabled

I8 is the byte container for packed FP4 values, not INT8 quantization. The checkpoint precision map establishes tensor storage and scale geometry; it does not establish runtime kernel precision. Bridge's E2M1 unpacking and the QAT configuration provide the implementation-level interpretation of the packed expert weights. Only indexer wq_b is FP8 in this table; the entire indexer is not uniformly FP8.

flowchart TD
    C["0731 checkpoint: MXFP4 experts / FP8 projections / BF16 / FP32"]
    I["Bridge import: unpack and dequantize quantized weights"]
    P["Trainable model parameters: BF16 plus native FP32 tensors"]
    Q["Routed expert weights: dynamic MXFP4 Q/DQ; BF16 temporary output"]
    T["TE FP8 quantization and GEMM: hybrid / blockwise recipe"]
    O["Adam update: FP32 main parameters and optimizer state; HDO CPU offload"]
    E["Bridge export: requantize updated parameters to MXFP4 / FP8 weight-scale pairs"]
    V["vLLM rollout / validation: quantization=fp8; KV cache=fp8"]
    C --> I --> P
    P -->|routed expert weights| Q
    Q --> T
    P -->|other TE FP8 linear paths| T
    T --> O
    O -->|update model parameters| P
    P -->|weight synchronization| E
    E --> V

Loading

QAT scope and kernel. qat.mode=mxfp4_experts disables all ModelOpt quantizers first, then enables only *mlp.experts*weight_quantizer, with num_bits=(2, 1) and block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)}. Scales are computed on each call; algorithm=None requires no calibration pass. Shared experts, attention, router, and all ModelOpt input/output quantizers remain disabled. This change does not implement FP4 indexer QK computation or additional KV/output-projection activation fake-quantization.

In the validated ModelOpt 0.44 implementation, MXFP4 fake-quant uses modelopt_cuda_ext_mx.fused_amax_convert → CUDA amax_and_convert_kernel. It returns BF16 values on the MXFP4 grid without permanently packing the trainable parameters. TE then quantizes weights and activations for its FP8 GEMMs. On the validated H100 setup, this simulates MXFP4 weight error while computing the supported GEMMs in FP8; it is not native FP4 GEMM. bypass_te_fp8_assert=True allows this explicit composition. Direct FP8 parameter loading and fp8_param are not enabled.

Scale and rollout handling. Checkpoint E8M0 scales are distinct from runtime compute scales: the recipe sets NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 for TE and VLLM_USE_DEEP_GEMM_E8M0=0 for the vLLM DeepGEMM path. Bridge exports the updated model parameters and independently recreates packed MXFP4 expert weights and FP8 projection weights with their scales. It does not export a previous forward pass's temporary fake-quant result. The generic NVFP4 QATWeightExporter is skipped for this mode, and rollout.qat.enable=False still retains Bridge's quantized weight synchronization. The diagram describes export formats and rollout configuration, not the compute precision of every vLLM kernel.

Checkpoint compatibility. Dynamic MXFP4 quantizers have no persistent calibration tensors. preserve_mxfp4_checkpoint_methods therefore preserves native TE/Megatron checkpoint methods while retaining the QAT forward path, with a guard rejecting persistent quantizer state. HDO resume restores optimizer main shards to the GPU before loading optimizer state. Neither fix changes the quantization recipe. Native FP32 parameters with HDO additionally require the detach fix in NVIDIA/Megatron-LM#6982, as documented in the example script.

Implementation: quantization configuration, checkpoint compatibility, and 40-step QAT example.

Quantization experiment evidence

The completed H100 experiment ran 40 steps with a 4096-token response limit, validation at steps 10/20/30/40, checkpoints at steps 20/40, and optimizer-state resume between the two 20-step segments. Step-40 AIME validation accuracy was 64.21% with MXFP4 QAT, versus 64.59% for the TE FP8 control without ModelOpt QAT. These are single-run measurements, not evidence of a statistically significant difference. See the QAT run and control run.

This experiment predates the final rebase. It used the implementation before publication cleanup and Megatron 1ff25ca7e with the equivalent local FP32-shard detach patch; the runtime patch helper is not included in this PR. These results validate that experiment, not a fresh 40-step run of the rebased head.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward, fully_async, one_step_off
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)
• If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}