NVFP4 RL

[zianglih]

@HumansAnd

Blog post: https://humansand.ai/blog/nvfp4-rl

Summary

This PR adds NVFP4 rollout/train support for Blackwell Miles with TE-compatible NVFP4 checkpoint conversion, Megatron-to-HF weight export, Qwen3 launch-script wiring, and GLM5.2 e2e coverage.

The current production quantization path is TE-direct rather than a FlashInfer/reference wrapper: Miles constructs transformer_engine.pytorch.tensor.nvfp4_tensor.NVFP4Quantizer in rowwise-only mode and extracts the rowwise NVFP4 payload fields needed by ModelOpt/SGLang (weight, weight_scale, weight_scale_2). FlashInfer remains the serving backend path via SGLang TRTLLM routed MoE, and Miles forwards relevant FLASHINFER_* and NVTE_* env vars so precision behavior is controlled by the backend libraries rather than new Miles CLI flags.

Implementation

• Adds a shared TE-backed NVFP4 quantization/export path for HF checkpoint conversion and Megatron-to-HF live weight update.
• Keeps the scope focused on MoE expert weights, while preserving BF16 carve-outs for configured first/last layers and extra high-precision layer patterns.
• Handles gated MoE pairs as one quantization unit so gate/up weights share scale semantics, with a storage-aware fast path that avoids an extra full-size copy when possible.
• Adds Qwen3 rollout/train NVFP4 launch wiring, backend precision env forwarding for NVTE_* / FLASHINFER_*, and disabled-by-default GLM5.2 Blackwell e2e coverage.

Validation

• Fast-gpu test after syncing this branch:

  • python3 -m pytest tests/fast-gpu/test_nvfp4_quantizer.py::test_nvfp4_quantize_params_omits_static_input_scale -q
  • Result: 1 passed, 25 warnings in 4.48s

• Qwen3-30B-A3B B200 NVFP4 after syncing this branch:

  • Step 1 quality sanity: train/train_rollout_logprob_abs_diff=0.030812, train/kl_loss=0.007849

• Earlier branch validation before the latest cleanup/rebase series:

  • NVFP4 quantizer bitwise tests were expanded to cover single tensors and paired gate/up tensors across precision-env variants.
  • GLM NVFP4 e2e coverage was brought up on B200 before being moved to the current GLM5.2 test file.

[gemini-code-assist]

Code Review

This pull request introduces comprehensive support for NVFP4 quantization, including utility functions for environment synchronization, integration with FlashInfer and TransformerEngine, and updates to conversion tools, run scripts, and end-to-end tests. The review feedback highlights two important improvements in miles/utils/nvfp4.py: first, adding a check to fall back to the TransformerEngine reference implementation when weights are on the CPU to prevent FlashInfer from crashing; second, moving the environment synchronization call inside the try block of the context manager to ensure that environment variables are properly restored even if synchronization fails.

[zianglih]

Root cause update for the NVFP4 4over6 bitwise mismatch:

The mismatch is not coming from the Miles quantizer plumbing. Miles is comparing FlashInfer NVFP4 quantization against the TE/FourOverSix-style reference. TE and FourOverSix compute the 4over6 candidate error by dequantizing with the original tensor amax expression, effectively:

e2m1_value * e4m3_scale * amax / (E2M1_MAX * E4M3_MAX_FOUROVERSIX)

The current FlashInfer installation computes the same error-domain quantity through the global decode scale path instead. That is algebraically equivalent, but not bitwise identical. On near-tie blocks this changes the strict 4-vs-6 candidate selection, so the packed NVFP4 values and scale bytes can differ from the TE reference.

Because of that, test_nvfp4_quantize_matches_te_reference_bitwise is not expected to pass with current FlashInfer until the FlashInfer-side fix in flashinfer-ai/flashinfer#3448 is available in the installed FlashInfer build. After installing a FlashInfer build with that PR, the Miles bitwise test passes on the devbox (416 passed).

[zianglih]

For accurate reproduction, need to make sure the flashinfer is after flashinfer-ai/flashinfer#3448 to include the bit-exact 4over6 NVFP4 quantization. That is at least https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.6.13rc1 .

# Replace FlashInfer with v0.6.13rc2:

   python3 -m pip uninstall -y flashinfer flashinfer-python flashinfer-cubin flashinfer-jit-cache
   python3 -m pip install --pre 'flashinfer-python[cu13]==0.6.13rc2' 'flashinfer-cubin==0.6.13rc2'
   python3 -m pip install --pre 'flashinfer-jit-cache==0.6.13rc2' \
     --index-url https://flashinfer.ai/whl/cu130
   flashinfer show-config

# Restore cudnn version after FlashInfer reinstall:

    pip install --no-deps nvidia-cudnn-cu13==9.16.0.29