[feat] 1/n Unified Quantization Framework - #1953
Conversation
![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 expert-only MXFP8 quantization on Blackwell GPUs (SM100/SM103) using Transformer Engine and vLLM, including new quantization strategies, model layouts, and benchmark scripts. Feedback on the changes highlights a potential RuntimeError in blockwise_cast_to_fp8 due to non-contiguous tensors, a missing copy import in the inference server utilities, and a misleading error message in the MoE weight splitting logic.
End-to-end FP8 for Megatron RL: FP8 compute recipes, persistent FP8 params (fp8_param) with exact optimizer-master initialization, and serialized blockwise-FP8 rollout weight sync into vLLM (FP8 codes + one FP32 scale per 128x128 block; batched MoE expert tensors stay fused for vLLM's 3D loader). Includes the review-round changes for NovaSky-AI#1898: - workers/megatron/quantization/ package and weight_sync/fp8/ split (quantize.py, vllm_format.py, models/). - Generic ModelFp8Spec registry (models/base.py): matches / should_quantize / ignored_layers / moe_expert_spec per model; the vLLM worker extension derives its fused-loader targets from the same specs. Qwen3.5 is the first registered spec (models/README.md documents adding one). - fp8_recipe="auto": architecture-native recipe defaults (blockwise on Hopper, native MXFP8 on SM100+), recipe-aware sequence alignment (MXFP8 1x32 tiles), and a symmetric per-arch block-scale env contract (NVTE_FP8_BLOCK_SCALING_FP32_SCALES / VLLM_USE_DEEP_GEMM_E8M0) validated at startup. - NVTE_FP8_BLOCK_AMAX_EPSILON provenance documented (escape hatch for blockwise-on-Blackwell; neither default path needs it). - FP8 GPU CI rows for the logprobs roundtrip test (full_fp8 dense/MoE + fp8_param, H100-validated) and FP8 configuration docs in the config dataclass docstrings.
… Blackwell - packing_utils: MXFP8 quantizes SP all-gather inputs in 1x32 tiles, so packed sequences align to 32*tp*cp at any TP; blockwise keeps 128*tp*cp (tp>1) and 16*cp local slabs at TP=1. - New distributed/megatron/quantization_utils.py holds the recipe/arch helpers (is_fp8_enabled, is_mxfp8_recipe, is_blackwell_or_newer, resolve_auto_fp8_recipe) that are not packing-specific. The resolver warns when a non-mxfp8 recipe is configured on SM100+, where TE emulates blockwise on the MX datapath. - MegatronConfig gains top-level fp8, fp8_recipe, fp8_param and fp8_amax_compute_algo fields; they fold into transformer_config_kwargs via setdefault, so an explicitly configured kwarg still wins. - Drop the NVTE_FP8_BLOCK_AMAX_EPSILON patch and its plumbing: it only ever applied to blockwise-emulated-on-Blackwell, which native MXFP8 replaces. Sync-side casts keep a fixed 1e-10 scale floor so all-zero blocks cannot degenerate. - Qwen3.5 FP8 spec ignores every vision block linear (attn.proj plus both MLP linears), not just attn.proj: vLLM builds the vision tower even for text-only runs and those dims stop being 128-divisible once TP-sharded. vLLM engines now build at inference TP=1/2/4 under blockwise FP8. - examples/train/fp8/: runnable Hopper blockwise, Hopper blockwise + fp8_param, and Blackwell MXFP8 recipes for dense 9B and MoE 35B-A3B, each with a colocated/non-colocated toggle. - Inline should_use_serialized_fp8 at its two call sites, and mirror the FP8 tests to the source layout (weight_sync/fp8/, workers/megatron/quantization/).
Two blank lines before the top-level forward_backward_payload helper; the pre-commit black hook fails on main without them.
fp8_weight_sync_mode already names the feature, so the value only needs to name the quantization scheme it puts on the wire. serialized_blockwise also read ambiguously: serial-vs-parallel rather than serialization, and it said nothing the key did not already say. blockwise pairs with the training recipe of the same name, so a config makes the train/rollout contract visible: fp8_recipe=blockwise with fp8_weight_sync_mode=blockwise quantizes both sides the same way, while a Blackwell run resolving to mxfp8 does not. A future wire format becomes fp8_weight_sync_mode=mxfp8.
|
hey @kailash109, just checking in, we were aiming to get #1898 in a state where adding MXFP8 is configurable there as well, which I think has some overlap with the changes for the unified quantization framework here it would be super helpful to see if your changes needed here (and for nvfp4) are expressible in the design from #1898, since that PR is quite close to being merged |
# Conflicts: # skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py # tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py
…edExperts nesting vLLM 0.26 turned FusedMoE into a factory returning a MoERunner whose RoutedExperts submodule registers the expert parameters, adding one segment to every runtime name. The batched-MoE loader now tries the nested name when the flat pre-0.26 name is absent and reports both candidates when neither resolves.
megatron_worker passes derive_metadata_from_chunks to whichever sender is active; the delta sender did not accept it, so every Megatron delta sync raised TypeError even with FP8 off. Delta checkpoints cannot represent serialized-FP8 wire chunks, so the flag is rejected explicitly rather than ignored.
A driver without a visible CUDA device resolved "auto" to blockwise and baked it into the config every worker receives, so Blackwell workers behind a CPU-only Ray head ran TE-emulated blockwise instead of native mxfp8. A blind process now leaves "auto" in place; each Megatron worker resolves against its own device and re-runs the device/recipe validation the driver had to skip.
Give the FP8 knobs the same treatment as the MEGATRON_* parallelism aliases so each example script shows its full FP8 configuration in one place. The resolved command lines are unchanged.
bc3b51a to
0012761
Compare
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Reviewed by Cursor Bugbot for commit 0012761. Configure here.
| return 32 * tp_size * cp_size | ||
| if tp_size > 1: | ||
| return 128 * tp_size * cp_size | ||
| return 16 * cp_size |
There was a problem hiding this comment.
Auto recipe packing alignment mismatch
High Severity
_fp8_token_align treats unresolved fp8_recipe="auto" as blockwise because only the literal "mxfp8" string selects the 32-token path. A GPU-less driver intentionally leaves "auto" unresolved while Blackwell workers resolve it to mxfp8, so controller-side packed layout (e.g. PackedDataCollator) and worker preprocess_packed_seqs disagree on align_size. That silently mis-reads multi-subsequence offsets and corrupts loss/grads.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 0012761. Configure here.
|
hey @erictang000 i've redone the changes here on top of #1898 's code, I'm re-running my e2e RL benchmark of mxfp8 vs bf16 to verify that the perf gains are still maintained this PR is already built off of the serialization path introduced in #1898 and so it was not too big of a refactor to rewrite the ModelFp8Spec from that PR on top of my abstraction's base classes (ModelQuantizationLayout and QuantizationStrategy as the interfaces for serialization of quantized params + defining which modules per module to be quantized). I think these are nicely extensible to supporting other quantization schemes in the future, and my mxfp8 implementation is already built on top of these in the original PR code and so that's unchanged Will keep you posted once i have the training results |
3 participants
Overview
This PR extends from #1940 and #1898 to propose a general unified framework for implementing quantization schemes in SkyRL -- the blockwise fp8 and mxfp8 implementations are refactored to fit in this implementation
follow-up PR implementing nvfp4 recipe on top of this: #1967
Additions
Added to skyrl/backends/skyrl_train/quantization:
QuantizedModelLayoutandQuantizationStrategyclasses: the QuantizedModelLayout class defines a model-specific weight structure, providing model-type lookup/validation, a method for splitting packed Bridge tensors (ie. qwen gate/up expert weights), and expert export layout required from Mbridge. QuantizationStrategy defines a complete quantization spec, providing ashould_quantize()method to select weight categories for quantization, methods to build the TE config, megatron provider, and runtime env for the particular quantization format, as well as a serialize_weight() functionality which emits the packed weights during weight syncs (this replaces the original logic ofiter_serialized_fp8_tensors()during weight syncs).These abstractions make it very easy to implement new quantization strategies with arbitrary packing/scale logic + targetting arbitary layers/modules of models. Rather than have to plumb through new configs for TE/Megatron and add specific per-layer quantization logic + additional weight serialization in the WeightExtractor, everything is automatically handled through the QuantizationStrategy class (and/or QuantizedModelLayout if adding quantization support for new model). This also enables significant reuse of existing quantization frameworks if supported models/vLLM weight recasting is already implemented.
Current WIP is implementing the nvfp4 recipe from humans& on top of these classes.
Validation
All CPU tests pass including new tests added to test emitted weight/scales from existing FP8/MXFP8 strategy + qwen3 layout.
e2e training run results:
Note
High Risk
Changes core Megatron FP8/MXFP8 training, optimizer master initialization, and the trainer→vLLM weight sync path; numerical mismatches or failed MoE loads would directly affect rollout policy correctness.
Overview
Introduces a unified quantization layer (
skyrl/backends/skyrl_train/quantization/) so Megatron training, Bridge export, and vLLM rollout sync share one model-layout + strategy split. Blockwise FP8 and expert MXFP8 are refactored intoQuantizationStrategyimplementations with Qwen MoE layouts, replacing ad-hoc serialization in the weight extractor.Rollout sync can use
generator.inference_engine.serialized_weight_sync_mode(blockwiseormxfp8) to send FP8 payloads and scales over NCCL broadcast or CUDA IPC; vLLM engine defaults and the inference worker wrap load tagged batched-MoE tensors into FusedMoE. Delta checkpoint sync is explicitly rejected for serialized modes. IPC splits mixed-dtype chunks; broadcast can derive metadata per chunk.Megatron gains expert-only MXFP8 training (
trainer.policy.model.expert_mxfp8), optional persistent expert FP8 params with deferred load and FP32 master init (fp8_param.py),fp8_recipe=autoresolution, and recipe-aware sequence packing alignment. Top-levelmegatron_config.fp8*fields fold intotransformer_config_kwargs.Adds FP8 DAPO example scripts, a Modal expert-MXFP8 benchmark, and docs for expert MXFP8 and quantization-aware weight sync.
Reviewed by Cursor Bugbot for commit 0012761. Bugbot is set up for automated code reviews on this repo. Configure here.