feat(sglang): NVFP4 quantized rollout and refit - #13
Closed
xiuhu17 wants to merge 1 commit into
Closed
Conversation
Owner
Author
|
Opened against the wrong repo — reopened upstream as NVIDIA-NeMo#3426. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
scheme: nvfp4alongside the existingmxfp8path underpolicy.generation.sglang_cfg.quantization, so SGLang rollouts can runNVFP4 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_projgets one shared global scaleinstead of two divergent ones.
nvfp4_setup.py— offline HF→NVFP4 checkpoint conversion with cachefingerprinting, 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 containerexpansion. SGLang picks one quantization method per
FusedMoEmodule, so asingle skipped expert tensor has to pull its whole container into BF16.
post_process_weights()is replaced bybegin_weight_update()/end_weight_update(). Quantized weights carryderived scales that have to be re-materialized after the last bucket
lands, which the old single-shot post-process could not express.
utils/patches.pybackports the endpoints to the pinned SGLang release.MegatronSGLangHfWeightIterator— quantizes during the AutoBridge walkand 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.
v2.17for the NVFP4 quantizer API.The DTensor/FSDP policy stays BF16-only and now says so explicitly
(
dtensor_sglang_refit.refit_sglang_colocatedraises 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.lockwas hand-resolved for the TE bump during therebase onto
mainand should be regenerated withuv lockin the projectcontainer before this merges.