Skip to content

[rollout, vllm, sglang, trainer] feat: quantized-layer audit - #8

Open
wengeezhang wants to merge 1 commit into
feat/mxfp8-refit-self-checkfrom
feat/quant-layer-audit
Open

wengeezhang wants to merge 1 commit into
feat/mxfp8-refit-self-checkfrom
feat/quant-layer-audit

Conversation

@wengeezhang

Copy link
Copy Markdown
Owner

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 own quantization_config). The quantized-layer audit compares them, each comparison once:

  • Training vs export (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.
  • Export vs engine (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.
  • vLLM receiver: an error raised while a bucket loads now surfaces after the transfer handshake. Raised inside the bucket callback, it left the trainer's BucketedWeightSender waiting for that bucket's ack with no timeout, and the run hung (seen on 1×B200 with the format check in raise mode). 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: the audit, and the diagram of all three checks, in the Checks section of docs/low_precision/mxfp8.md; the format check in docs/low_precision/qat.md.

API and Usage Example

No config or API change. VERL_QUANT_LAYER_AUDIT switches both comparisons: warn (the default: log and keep going), raise (stop the sync) or 0 (off). VERL_QUANT_LAYER_AUDIT_TRACE=1 logs each step of the training-side comparison.

export VERL_QUANT_LAYER_AUDIT=raise

Test

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 passed

pre-commit run passes 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 raise mode. In the fault runs, a test hook sends layer 0's o_proj in BF16 instead of FP8 after the training-vs-export audit has recorded it, so only the engine side sees the wrong format.

run training result
vLLM and SGLang, defaults MXFP8 (TE fp8_recipe=mxfp8) passes on both: the format check finished at the first sync (69 weights checked, none in another format, none unmapped), and training vs export agreed after step 1
vLLM and SGLang, defaults BF16 passes on both: training vs export is not armed, and the format check still runs (69 / 0 / 0)
SGLang, o_proj sent in BF16 MXFP8 and BF16 the format check reports that weight and no other (the 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 1
vLLM, o_proj sent in BF16, on 3ab3b20 MXFP8 and BF16 the same report, also before the write, but the run then hung until stopped (281 s and 169 s later), as described above
vLLM, o_proj sent in BF16, on ba2852e MXFP8 the run fails 5 s after the report, before step 1; the trainer's traceback ends with the format check's message
vLLM, defaults, on ba2852e MXFP8 passes, as on 3ab3b20: 69 / 0 / 0, training vs export agreed, rollout_corr/kl 0.0021

Earlier, 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 ignore differently (verl-project#8306). _translate_ignored_layers passes plain names to SGLang, which matches them as prefixes, while the trainer-side exporter matches exact module names or re: patterns. With ignore: ["model.layers.0"], SGLang keeps layer 0 in BF16 but the exporter still quantizes it. The export-vs-engine check refuses that sync in raise mode (28 weights, measured on cf32f6a); in the default warn mode 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

  • New: 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.py and verl/utils/sglang/sglang_sync_format_check.py (the parameter each synced weight lands in).
  • Training side: 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).
  • Engine side: 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: 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.

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.
@wengeezhang
wengeezhang added this pull request to stack #9 November 2, 2026 08:46
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.

1 participant