Repository navigation
[rollout, vllm, sglang, trainer] feat: quantized-layer audit - #8
Open
wengeezhang wants to merge 1 commit into
Open
wengeezhang wants to merge 1 commit into
wengeezhang wants to merge 1 commit into
Conversation
Three parties decide which layers are quantized, each on its own: training (TE linear modules inside fp8_autocast), and by name the trainer-side exporter and the engine. The audit compares them, each comparison once, under VERL_QUANT_LAYER_AUDIT (warn by default, raise or 0): - training vs export (TrainExportAudit), when training runs a TE FP8 recipe, at the first weight sync after a training step; - export vs engine (SyncFormatCheck), for any quantized rollout, inside each engine's weight loader on the first weight sync, before each bucket is written; vLLM and SGLang share one implementation. vLLM's receiver now raises an error from a bucket's load only after acking the remaining buckets: raised inside the bucket callback, it left the trainer's sender waiting for an ack forever (seen on 1xB200 with the format check in raise mode). Split from verl-project#8327.
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.
What does this PR do?
Part 3 of 3 of verl-project#8327, on top of part 2 (the MXFP8 refit fixes and the post-sync self-check, this PR's base). With the three parts applied, the tree matches verl-project#8327 without its unit tests.
Three parties decide which layers are quantized, each on its own: training implicitly (TE linear modules inside
fp8_autocast), and by name the trainer-side exporter and the engine, both from the same quantization config (rollout.quantization, or the checkpoint's ownquantization_config). The quantized-layer audit compares them, each comparison once:TrainExportAudit), when training runs a TE FP8 recipe, at the first weight sync after a training step: each decoder layer's FP8 trace (TE's FP8 weight workspaces) against the format the sync sends for that layer's GEMM projections.SyncFormatCheck), for any quantized rollout, inside each engine's weight loader during the first weight sync, before each bucket is written: every incoming weight's format (quantized or high precision) against the engine parameter that receives it. One implementation serves vLLM and SGLang; each engine only resolves which parameter a synced name lands in.BucketedWeightSenderwaiting for that bucket's ack with no timeout, and the run hung (seen on 1×B200 with the format check inraisemode). The remaining buckets are now acked without being loaded, and the error is raised once the sender is done, as the delta path already does. This covers any error raised while a bucket loads, which hangs the same way on main.docs/low_precision/mxfp8.md; the format check indocs/low_precision/qat.md.API and Usage Example
No config or API change.
VERL_QUANT_LAYER_AUDITswitches both comparisons:warn(the default: log and keep going),raise(stop the sync) or0(off).VERL_QUANT_LAYER_AUDIT_TRACE=1logs each step of the training-side comparison.export VERL_QUANT_LAYER_AUDIT=raiseTest
Unit tests are left out of this branch for now. All of verl-project#8327's CPU tests pass against this branch's code:
python -m pytest -q tests/utils/test_sglang_quant_refit_on_cpu.py tests/utils/test_mxfp8_refit_check_on_cpu.py \ tests/utils/test_vllm_mxfp8_refit_check_on_cpu.py tests/utils/test_vllm_unquant_utils_on_cpu.py \ tests/utils/test_quant_layer_audit_on_cpu.py tests/utils/test_quant_sentinel_on_cpu.py \ tests/utils/test_quant_sync_format_check_on_cpu.py tests/workers/rollout/test_vllm_weight_update_utils_on_cpu.py # 168 passedpre-commit runpasses on the changed files.B200 validation
On verl-project#8327's code at 3ab3b20, then ba2852e (the receiver fix), on 1×B200: a tiny random Qwen3-MoE (2 layers, 8 experts, top-2) on GSM8K and an MXFP8 rollout on each engine's defaults; vLLM 0.29.0, SGLang 0.5.20, Transformer Engine 2.16.1, Megatron-Core 0.19.2. 2 GRPO steps per run, with every check in
raisemode. In the fault runs, a test hook sends layer 0'so_projin BF16 instead of FP8 after the training-vs-export audit has recorded it, so only the engine side sees the wrong format.fp8_recipe=mxfp8)o_projsent in BF16the sync sends bfloat16, the engine parameter is float8_e4m3fn) and stops before its bucket is written; the error ends SGLang's scheduler, which takes the server down, and the run fails 11–12 s later, before step 1o_projsent in BF16, on 3ab3b20o_projsent in BF16, on ba2852erollout_corr/kl0.0021Earlier, on verl-project#8314's code: Qwen3-30B-A3B on 4×B200 (training TP2 / EP2, vLLM rollout TP2, 5 GRPO steps): the audit agreed on 48 of 48 decoder layers (18,817 weights), and the engine-side comparison matched all 18,867 synced weights.
Found on main while validating (not changed here): the exporter and SGLang read
ignoredifferently (verl-project#8306)._translate_ignored_layerspasses plain names to SGLang, which matches them as prefixes, while the trainer-side exporter matches exact module names orre:patterns. Withignore: ["model.layers.0"], SGLang keeps layer 0 in BF16 but the exporter still quantizes it. The export-vs-engine check refuses that sync inraisemode (28 weights, measured on cf32f6a); in the defaultwarnmode it logs the disagreement and the FP8 bytes are loaded into the BF16 parameters. For MoE layers,re:.*mlp\.experts\..*(all experts) is the only pattern both read the same way today.Design & Code Changes
verl/utils/quant_layer_audit.py(training vs export),verl/utils/quant_sync_format_check.py(export vs engine, shared by both engines),verl/utils/vllm/vllm_sync_format_check.pyandverl/utils/sglang/sglang_sync_format_check.py(the parameter each synced weight lands in).verl/workers/engine_workers.py(creates the audit with the weight quantizer),verl/workers/engine/megatron/transformer_impl.py(the export passes through it),verl/utils/megatron_utils.py(records which modules held FP8 workspaces before they are cleared).verl/workers/rollout/vllm_rollout/utils.py(the format check around the quantized load; the receiver fix),verl/utils/sglang/sglang_quant_utils.py(the format check in the refit loader).docs/low_precision/mxfp8.md,docs/low_precision/qat.md.AI assistance disclosure
Developed with AI assistance. I reviewed every changed line, ran the tests and the B200 experiments above, and can defend the change end to end.