Skip to content

feat(sglang): NVFP4 quantized rollout and refit - #13

Closed
xiuhu17 wants to merge 1 commit into
zhw/mxfp8_supportfrom
zhw/nvfp4_support2
Closed

feat(sglang): NVFP4 quantized rollout and refit#13
xiuhu17 wants to merge 1 commit into
zhw/mxfp8_supportfrom
zhw/nvfp4_support2

Conversation

@xiuhu17

@xiuhu17 xiuhu17 commented Jul 30, 2026

Copy link
Copy Markdown
Owner

Adds scheme: nvfp4 alongside the existing mxfp8 path under
policy.generation.sglang_cfg.quantization, so SGLang rollouts can run
NVFP4 weights while training stays BF16/FP8.

Stacked on zhw/mxfp8_support; review only the top commit.

What this does

  • nvfp4_quantization_core.py — TE-backed NVFP4 weight quantization.
    Two-level scaling (per-16-element E4M3 block scales + a per-tensor global
    encode/decode scale), row padding to TE's 16-row alignment, and gate/up
    pair quantization so a fused gate_up_proj gets one shared global scale
    instead of two divergent ones.
  • nvfp4_setup.py — offline HF→NVFP4 checkpoint conversion with cache
    fingerprinting, so SGLang boots from a checkpoint whose tensor layout and
    high-precision exclusions match what the online refit will send.
  • quantization_utils.py — the scheme-agnostic pieces both schemes need:
    scheme validation, the dynamic high-precision policy
    (extra_high_precision_layers_hf, modules_to_not_convert,
    num_layers_at_start_in_bf16 / ..._end_in_bf16), and fused-MoE container
    expansion. SGLang picks one quantization method per FusedMoE module, so a
    single skipped expert tensor has to pull its whole container into BF16.
  • Refit session APIpost_process_weights() is replaced by
    begin_weight_update() / end_weight_update(). Quantized weights carry
    derived scales that have to be re-materialized after the last bucket
    lands, which the old single-shot post-process could not express.
    utils/patches.py backports the endpoints to the pinned SGLang release.
  • MegatronSGLangHfWeightIterator — quantizes during the AutoBridge walk
    and keeps every companion tensor produced from one source weight (and each
    NVFP4 gate/up pair) inside the same bucket, so a bucket boundary can never
    split a weight from its scales.
  • TransformerEngine is bumped to v2.17 for the NVFP4 quantizer API.

The DTensor/FSDP policy stays BF16-only and now says so explicitly
(dtensor_sglang_refit.refit_sglang_colocated raises on any other scheme)
rather than silently sending BF16 to a quantized engine.

Prior art

This follows the NVFP4 RL work in Miles:

Status

Draft. Unit tests for the quantization core, the offline converters, the
skip/high-precision policy, the weight iterator, and the SGLang patches are
included and pass. uv.lock was hand-resolved for the TE bump during the
rebase onto main and should be regenerated with uv lock in the project
container before this merges.

@xiuhu17

xiuhu17 commented Jul 30, 2026

Copy link
Copy Markdown
Owner Author

Opened against the wrong repo — reopened upstream as NVIDIA-NeMo#3426.

@xiuhu17 xiuhu17 closed this Jul 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant