V4 mxfp4#1372
Conversation
Verified end-to-end on B200 (4-layer prune, colocate, tp=ep=2):
weight-update check (--check-weight-update-equal) passes bitwise with
the deepgemm fp8 GEMM backend; train_rollout_logprob_abs_diff ~= 0.06.
- run_deepseek_v4.py: replace fp8_training with an explicit --recipe
{fp8,bf16} selector. On Blackwell, TE emulates the blockwise recipe
with MXFP8 and asserts power-of-two scales, so instead of exporting
NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 (Hopper-only now) override the
unconditional =1 injection from ray/actor_group.py back to 0 via
--train-env-vars.
- deepseek_v4 qat.py: force fp32 scale storage in fp8_simulate. The
act_quant auto dtype picks float8_e8m0fnu on Blackwell+DeepGEMM, but
TileKernels cast_back only accepts int32/fp32 scales.
Stacks on the blockwise-FP8 base: adds --recipe mxfp8 (Blackwell only) covering offline conversion, rollout serving, and online weight update. - convert_hf_to_mxfp8.py: extend SKIP_WEIGHT_SUBSTRINGS with the V4 weights the recipe keeps in BF16: "head." (output head), "wo_a" (o-projection LoRA-A; wo_b stays FP8), "ffn.gate." (router), and "compressor." (attn + indexer compressor wkv/wgate; scoped so the FP8 attn.wkv is not matched). Verified against the full recipe-checkpoint inventory: every FP8 tensor converts, every BF16 tensor is skipped and lands in modules_to_not_convert. - quantizer_mxfp8.py: add the V4 attention weights to the online weight-update whitelist (wq_a, wkv, wo_b, indexer wq_b/wk), mirroring the quantizer_fp8 block; without them a V4 mxfp8 update silently sends these to sglang as BF16 while the engine expects MXFP8 tensors. - run_deepseek_v4.py: add recipe=mxfp8 - prepare-mxfp8 (BF16->MXFP8, full-train sentinel), hf_checkpoint redirect to the -MXFP8 dir, flashinfer_cutlass dense GEMM for serving, --fp8-recipe mxfp8 for the trainer, and Hopper asserts. Verified on B200 (4-layer prune): conversion + serving + weight update all load (52/52 buckets). Known open issue: with the trtllm-routed MoE backend, --check-weight-update-equal reports routed-expert mismatches (snapshot vs update layout interplay in align_mxfp8_moe_weights_for_flashinfer_trtllm); under investigation.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for DeepSeek V4 MXFP8 and MXFP4 QAT (Quantization-Aware Training) recipes, primarily targeting Blackwell hardware. It adds block-wise FP4 quantization kernels, updates the attention compressor and indexer to simulate MXFP4 QAT, and integrates these recipes into the training and rollout scripts. Feedback on these changes highlights a shape and dtype mismatch in the fp4_act_quant helper when inplace=False, a missing Blackwell hardware assertion for the mxfp4_qat recipe in the training script when hardware detection is set to auto, and a missing MXFP4 QAT check in the non-rotated branch of the compressor's forward pass.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| N = x.size(-1) | ||
| assert N % block_size == 0 | ||
| z = x.contiguous() | ||
| y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2) |
There was a problem hiding this comment.
In fp4_act_quant, when inplace is False, y is created with shape (*z.shape[:-1], N // 2) and dtype torch.float4_e2m1fn_x2. However, the TileLang kernel fp4_quant_kernel expects Y to have shape (M, N) and dtype FP4 (which is float4_e2m1fn). This mismatch in both shape (N vs N // 2) and dtype (float4_e2m1fn vs float4_e2m1fn_x2) will cause a runtime shape/dtype validation error when the kernel is launched.
To fix this, y should be initialized with shape *z.shape and dtype torch.float4_e2m1fn when inplace is False.
| y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2) | |
| y = torch.empty_like(z) if inplace else z.new_empty(*z.shape, dtype=torch.float4_e2m1fn) |
| if args.recipe == "mxfp8": | ||
| assert _is_blackwell(args), "MXFP8 requires Blackwell (B200/B300/GB200/GB300)" |
There was a problem hiding this comment.
The mxfp4_qat recipe is Blackwell-only (as asserted in __post_init__ when self.hardware is explicitly set to H100/H200). However, if self.hardware is set to "auto", the __post_init__ check is bypassed. In _train, there is an assertion to ensure Blackwell is used for mxfp8, but mxfp4_qat is missing from this check. This can lead to runtime crashes on Hopper GPUs when using "auto" hardware detection with the mxfp4_qat recipe.
We should update the assertion in _train to cover both mxfp8 and mxfp4_qat.
| if args.recipe == "mxfp8": | |
| assert _is_blackwell(args), "MXFP8 requires Blackwell (B200/B300/GB200/GB300)" | |
| if args.recipe in ("mxfp8", "mxfp4_qat"): | |
| assert _is_blackwell(args), f"{args.recipe.upper()} requires Blackwell (B200/B300/GB200/GB300)" |
| if self.use_mxfp4_qat: | ||
| kv = mxfp4_simulate_qat(kv, 32) | ||
| elif self.use_fp8_qat: | ||
| kv = fp8_simulate_qat(kv, 128) |
There was a problem hiding this comment.
In forward_raw, when self.rotate is False, the else branch (lines 165-168) only checks self.use_fp8_qat and does not handle self.use_mxfp4_qat. If self.use_mxfp4_qat is enabled and self.rotate is False, the MXFP4 QAT simulation will be silently skipped for the non-rotated nope head dimension.
Consider adding the self.use_mxfp4_qat check to the else branch as well to ensure consistency and robustness.
1 participant
Should based on: radixark/Megatron-LM#59