FP8 RL MXFP8 Weight Sync - #2072
Open
YJHMITWEB wants to merge 15 commits into
Open
Conversation
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.
# 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.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces support for serialized FP8 weight synchronization (blockwise and MXFP8) and persistent FP8 parameters across the training and rollout pipeline, alongside robustness improvements to token-based batching. The review feedback highlights several critical runtime safety issues, including potential AttributeError and KeyError crashes when retrieving padding flags from microbatch metadata, as well as potential PyTorch RuntimeError exceptions from calling .view() on non-contiguous tensors during quantization and caching. Addressing these issues with the suggested contiguity and metadata guards will ensure the stability of the FP8 pipeline.
Comment on lines
+834
to
+835
| if not microbatch.metadata.get("is_padding_batch"): | ||
| all_loss_fn_outputs.extend(outputs) |
Contributor
There was a problem hiding this comment.
If microbatch.metadata is None, calling .get() on it will raise an AttributeError. Adding a guard to check if metadata is not None prevents potential runtime crashes.
Suggested change
| if not microbatch.metadata.get("is_padding_batch"): | |
| all_loss_fn_outputs.extend(outputs) | |
| is_padding = microbatch.metadata.get("is_padding_batch") if microbatch.metadata else False | |
| if not is_padding: | |
| all_loss_fn_outputs.extend(outputs) |
Collaborator
There was a problem hiding this comment.
This won't happen as we defined metadata for all entries in skyrl/backends/skyrl_train/workers/worker_utils.py
Comment on lines
+1348
to
+1352
| validate_padding_head( | ||
| [m["is_padding_batch"] for m in micro_buffer], | ||
| num_padding_microbatches, | ||
| "forward_backward", | ||
| ) |
Contributor
There was a problem hiding this comment.
Accessing m["is_padding_batch"] directly on the microbatch object will raise a KeyError because "is_padding_batch" is stored in the metadata dictionary rather than as a direct key of the TrainingInputBatch dict. Retrieve this flag from m.metadata instead.
Suggested change
| validate_padding_head( | |
| [m["is_padding_batch"] for m in micro_buffer], | |
| num_padding_microbatches, | |
| "forward_backward", | |
| ) | |
| validate_padding_head( | |
| [bool(m.metadata.get("is_padding_batch", False)) if m.metadata else False for m in micro_buffer], | |
| num_padding_microbatches, | |
| "forward_backward", | |
| ) |
Collaborator
There was a problem hiding this comment.
micro_buffer above already contains this key.
Comment on lines
+1410
to
+1411
| if m_batch["is_padding_batch"]: | ||
| continue |
Contributor
There was a problem hiding this comment.
Accessing m_batch["is_padding_batch"] directly will raise a KeyError for the same reason. Retrieve the padding status from m_batch.metadata to prevent a runtime crash.
Suggested change
| if m_batch["is_padding_batch"]: | |
| continue | |
| is_padding = bool(m_batch.metadata.get("is_padding_batch", False)) if m_batch.metadata else False | |
| if is_padding: | |
| continue |
Collaborator
There was a problem hiding this comment.
nope, same reason as above
On Blackwell the trainer computes with Transformer Engine's native MXFP8 recipe (1x32 groups, E8M0 scales), but the serialized FP8 weight sync always shipped blockwise (128x128, FP32 scales), so the rollout served a different quantized representation than the one training computed with. Add an MXFP8 wire format and make fp8_weight_sync_mode="auto" resolve from the policy's resolved fp8_recipe. Resolution keys off the recipe, never the architecture: an explicit blockwise recipe on Blackwell keeps a blockwise wire, because train/rollout agreement is the point. Blockwise output stays byte-identical and no vLLM patch is required. - mx_cast_to_fp8 / batched_mx_cast_to_fp8 emit E4M3 codes plus per-32-col E8M0 uint8 scales, bitwise-identical to TE's MXFP8Quantizer. The payload ships in compressed-tensors layout; engines boot with quantization=compressed-tensors from an injected config and the layerwise-reload bracket loads it unmodified. - All routing lives in quantization_utils: resolve_auto_fp8_recipe, resolve_auto_wire_format, wire_to_engine_quantization. Config resolution, engine bootstrap and launchers share one answer to "how does auto route?"; an unresolved "auto" recipe is rejected with an actionable message instead of silently defaulting to blockwise. - The MXFP8 ignore list extends blockwise's with the kernel floors (K % 32 == 0, N >= 128), which exclude the GDN in_proj and vision projections. Stripping them makes a real engine refuse to build. - CUDA casts dispatch to TE's MXFP8Quantizer and 3D expert stacks flatten to one [E*N, K] call; the torch path stays as the CPU fallback and the parity oracle. SKYRL_MX_CAST_BACKEND=python|te|auto. - Cache the TRT-LLM MoE prepare permutation across weight syncs. All four tensors are 1-byte dtypes, so the per-expert shuffle collapses to one index_select per output. The permutation is learned per shape key on first use and must reproduce the original bitwise on two validation inputs before it is trusted; any mismatch, including a future vLLM layout change, falls back permanently. SKYRL_TRTLLM_MOE_PREPARE_CACHE=0 disables it. - The runtime block-scale environment pins apply to the blockwise wire only; MXFP8 has no amax-epsilon or scale-mode concept. - FP8 KV cache (kv_cache_dtype=fp8_e4m3) now composes with the serialized wires. vLLM 0.26 corrupts attention scales twice on this path: the compressed-tensors KV method copies dummy-load placeholder scales verbatim at boot (no sentinel default; FlashInfer then bakes the float mirrors into captured attention plans), and the post-wake reset covers only the k/v tensors, so q wakes as 0.0. Result: NaN generations on the quantized-Q path, silently wrong logprobs (mean diff 1.24) on the bf16-Q path. vllm_compat now normalizes every scale (tensors, float mirrors, CPU copies) to 1.0 — the wire's contract, since it ships no calibration — at boot before graph capture, after every wake, and after every weight sync. A plain-vLLM control confirmed the fp8-KV kernels themselves are healthy; pinned by a new b200 E2E row (dense mxfp8 + fp8 KV) and validated at the 35B tp1/ep8 production shape. - Validate GDN in_proj TP-shard alignment at worker init: Megatron's own guard checks only the GLOBAL fused in_proj dim (gated_delta_net.py), but TE quantizes the local TP shard and requires every dim % 32 (MXFP8Quantizer::create_tensor). A misaligned config (e.g. Qwen3.5's 12352 = 32*386 at TP=4) previously passed validation and died as a TE C++ assert mid model-build; it now raises with the arithmetic and the fix before the model is constructed. Yield padding microbatches first. Megatron's overlap_grad_reduce latches bucket communication off the last microbatch's backward, and an all-masked padding microbatch there produces a degenerate backward that makes finish_grad_sync assert. Both ordering-sensitive consumers are corrected with it: reorder_and_combine_batches slices from the head, and _reorder_megatron_forward_output offsets by the padding count so a rank holding padding never reads padding outputs as real logprobs. Padding outputs are also kept out of loss_fn_outputs. Numerics are unchanged. - Verify the padding-first layout at runtime instead of trusting it: outputs carry their microbatch's is_padding_batch metadata, and validate_padding_head refuses any run where padding is not exactly the head — at the reorder head slice, the megatron forward chunk offset (plus a chunk-count check), and forward_backward (overlap_grad_reduce needs a real LAST microbatch or finish_grad_sync asserts). A yield-order revert now fails loudly with the contract named instead of crashing inside Megatron or silently reading padding logprobs as real samples. Signed-off-by: YJHMITWEB <24315061+YJHMITWEB@users.noreply.github.com>
jinghanyao1-hub
force-pushed
the
fp8-rl-mxfp8-weight-sync
branch
from
August 20, 2026 17:20
14bec7f to
1d1dccc
Compare
…-sync # Conflicts: # skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py
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 4580e91. Configure here.
| groups = cols // MXFP8_GROUP_SIZE | ||
| if tuple(scales.shape) != (rows, groups): | ||
| scales = scales[:rows, :groups].contiguous() | ||
| return codes.view(rows, cols), scales |
There was a problem hiding this comment.
TE MXFP8 scales skip uint8 contract
High Severity
The TE fast path in _te_mx_cast_2d returns _rowwise_scale_inv as-is after an optional pad crop, while mx_cast_to_fp8 documents and implements vLLM compressed-tensors scales as biased uint8 in logical [rows, cols // 32] layout. TE commonly stores E8M0 as float8_e8m0fnu in a padded scale matrix, so the wire can ship the wrong dtype and/or extra pad elements. The receiver then loads scales that do not match scale_dtype: uint8, which can fail the sync or silently dequantize MXFP8 weights incorrectly.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 4580e91. Configure here.
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
2 participants
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.
MXFP8 Weight Synchronization for Blackwell Rollouts
This PR is based on PR #1898 and #1899, and extends the FP8 RL path with MXFP8 weight synchronization on Blackwell.
In PR #1898, the only FP8 wire is blockwise (128×128). On Blackwell, that means the trainer computes in TE's MXFP8 recipe while the rollout is served with blockwise weights. That design does converge as validated by our long-run tests, since the weight sync itself guarantees the correctness of the quantized weights. In this PR, for consistency, we make the wire match the trainer. User can set
fp8_weight_sync_mode=auto, and it follows the trainer's resolved fp8_recipe, so an MXFP8 trainer gets an MXFP8 rollout automatically.We also make the weight-sync faster. The cast dispatches to TE's
MXFP8Quantizer(inquantize.py).On vLLM side, we add a permutation cache for the MXFP8 MoE prepare (FlashInfer's TRT-LLM-Gen path), which otherwise re-derives a fixed relocation with a per-expert Python loop on every sync (
_shuffle_mxfp8_moe_weights, flashinfer_utils.py#L355-L427: per-expertreorder_rows_for_gated_act_gemm+shuffle_matrix_a+shuffle_matrix_sf_a— pure byte relocations, identical for fixed shapes, so we learn the row permutation once and replay it as oneindex_selectper tensor. The speedup: 8.13 s → 1.87 s, compared with 4.50 s for blockwise.FP8 KV cache works with this wire too
PR #1899 only covers the blockwise weight-sync and the corresponding fp8 kv cache. With MXFP8 weight-sync, the attention scales are not reset properly — in particular
_q_scaleremains 0 after the engine wakes from sleep, because vLLM's post-wake reset only restores_k_scale/_v_scale(gpu_model_runner.py#L980-L1021).On B200 vLLM quantizes the query by dividing it by
_q_scale(flashinfer.py#L1658-L1677), so a zero scale means NaN in the run. On top of that, the compressed-tensors path our wire lands on copies the boot-time placeholder scales without validation(
compressed_tensors.py#L1183-L1213), and those values get frozen into the CUDA graphs at capture.We fix this in
vllm_compat.py: our wire ships no scale calibration, so the correct value is always 1.0, and we force all attention scales to 1.0 at boot (before CUDA-graph capture), after every wake, and after every weight sync.Quality Checks
Qwen3.5-35B-A3B runs (with full MXFP8): https://wandb.ai/sky-posttraining-uc-berkeley/qwen35_35b_a3b_fp8_amaxeps_20260705?nw=nwuserjinghanyao1
Reproduce key configs
B200, showing only the difference from the blockwise B200 configuration:
Note
High Risk
Touches Megatron FP8 training, mixed-dtype weight transfer, and vLLM load/KV-scale patches. Incorrect quantization, packing alignment, or scale resets can silently corrupt logprobs or crash MoE/GDN builds.
Overview
Makes the FP8 RL rollout wire match the trainer recipe:
fp8_weight_sync_mode=autofollows the resolvedfp8_recipe, so Blackwell MXFP8 training ships MXFP8 (E4M3 + E8M0) weights to vLLM instead of re-quantizing a BF16 export or using a blockwise wire.Adds a serialized FP8 stack (Qwen3.5 spec, blockwise/MXFP8 casts with a TE MXFP8 fast path, batched MoE tensors) and boots vLLM with dummy load + the matching
fp8/compressed-tensorsconfig. NCCL/IPC send mixed-dtype chunks; the worker loads fused MoE via a compact batched prefix.Also: recipe-aware sequence packing;
fp8_paramoptimizer-master init (Hopper only; MXFP8 forbids it); vLLM patches that force attention KV scales to 1.0 at boot/wake/sync; a learned permutation cache for MXFP8 TRT-LLM MoE prepare; DAPO example scripts for Hopper blockwise and Blackwell MXFP8.Reviewed by Cursor Bugbot for commit 4580e91. Bugbot is set up for automated code reviews on this repo. Configure here.