Skip to content

[feat] 2/n Adding NVFP4 Humans& Quantization - #1967

Draft
kailash109 wants to merge 21 commits into
NovaSky-AI:mainfrom
kailash109:nvfp4
Draft

[feat] 2/n Adding NVFP4 Humans& Quantization#1967
kailash109 wants to merge 21 commits into
NovaSky-AI:mainfrom
kailash109:nvfp4

Conversation

@kailash109

@kailash109 kailash109 commented Jul 31, 2026

Copy link
Copy Markdown

Building off of #1953 to support humans& nvfp4 recipe

Overview

The base quantization strategy is built off of the abstractions introduced in #1953

  • A new file skyrl/backends/skyrl_train/quantization/nvfp4.py introduces Nvfp4ExpertStrategy, holds the Megatron/TE recipe builders, blockwise casts for 4over6 quantization (everything goes through the NVFP4Quantizer exposed by TE), layer-wise casting (in the humans& recipe they leave the last 15% of layers in BF16, this is implemented via a high_precision_last_layers count, for now no shared expert in qwen-30b-a3b but this is part of the excluded modules)
  • skyrl/backends/skyrl_train/quantization/vllm_nvfp4.py contains some patches necessary to get nvfp4 generation working with VLLM (specifically, [Quant] Add nvfp4_per_token online MoE quantization vllm-project/vllm#48538 does not support EP, the patch does per-token nvfp4 quantization on each gpu after all-gathering across EP ranks and before combine inside vllm's FusedMoEKernel experts.apply()). Currently this works for EP runs but is really messy by patching private methods in vllm so I'll see if I can potentially upstream this
  • require te 2.17 for nvfp4 + 4over6 quantize -- added this as a new optional dependency group only enabled during nvfp4 runs in benchmark script
  • same benchmark script (examples/train/megatron/modal_expert_mxfp8_benchmark.py) added mode for nvfp4 in same e2e run as before (Qwen3-30B-A3B, dapo math 17k grpo)

Validation

Tested working two full DAPO steps as well as extra test coverage for per-tensor global scales, four-over-six quantization, BF16 final-layer exclusions, vllm init + receiver mappins

Keeping in draft for now, waiting for correctness + step time perf results from e2e validation runs, as well as need to do a deeper review of the code

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.

2 participants