Repository navigation
Conversation
Contributor
wuxibin89
added a commit
that referenced
this pull request
Oct 19, 2026
…and low-precision rollout (#8287) ### What does this PR do? This PR refactors Megatron QAT and low-precision rollout around a unified trainer-side weight quantization and rollout-side refit boundary. - Add compressed-tensors-compatible trainer-side exporters for FP8_BLOCK, MXFP8, MXFP4, and NVFP4. - Add ModelOpt-backed Megatron QAT modes: `fp8`, `mxfp8`, `mxfp4`, `mxfp4_mxfp8`, and `nvfp4`. - Replace the string `rollout.quantization=fp8` API with structured `rollout.quantization.linear`, `moe`, `ignore`, and `config_file` settings. - Load and refit pre-quantized weights in vLLM and SGLang instead of quantizing BF16 weights inside rollout workers. - Add vLLM MXFP4/NVFP4 refit support, update FP8 refit handling, and update SGLang routing replay support. - Add a unified QAT guide, migration notes, examples, CPU tests, and a rollout log-prob comparison tool. - Remove the legacy `verl.utils.modelopt` implementation and TRT-LLM rollout quantization support. Related to #5972 and #4066. This PR has partial scope overlap with open low-precision work, but uses a materially different architecture: #7519 implements Transformer Engine shared-grid MXFP8 training and rollout, while this PR implements ModelOpt QAT plus a generalized trainer-side compressed-tensors export boundary for multiple formats; #7971 implements native Transformer Engine NVFP4 real quantization, while this PR implements NVFP4 fake-quant QAT and pre-quantized refit; #8205 fixes the legacy FSDP NVFP4 compressor API, while this PR refactors the Megatron and rollout path and does not add this feature to FSDP. AI assistance was used to review the branch, fix two review findings, add a regression test, and prepare this PR description. The human submitter reviewed the changes and is responsible for the contribution end to end. ### Checklist Before Starting - [x] Search for similar PRs: https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+QAT+compressed-tensors+rollout - [x] Format the PR title as `[{modules}] {type}: {description}`. ### Test Passed: ```bash git diff --check origin/main...wuxibin/qat python -m compileall -q verl/utils/qat verl/utils/vllm verl/workers/rollout/sglang_rollout verl/workers/rollout/trtllm_rollout PATH="$PWD/.venv/bin:$PATH" pre-commit run --all-files --show-diff-on-failure --color=always ``` The full pre-commit suite passed, including Ruff, Ruff format, mypy, generated-config verification, repository sanity checks, and compileall. The targeted pytest suites and GPU end-to-end tests were not run in this local environment because PyTorch, ModelOpt, and compressed-tensors were unavailable. CI/GPU validation should cover Megatron with vLLM and SGLang using FP8_BLOCK, plus MXFP8/MXFP4/NVFP4 on supported Blackwell hardware. ### API and Usage Example The previous API: ```bash actor_rollout_ref.rollout.quantization=fp8 ``` is replaced by: ```bash actor_rollout_ref.rollout.quantization.linear=FP8_BLOCK \ actor_rollout_ref.rollout.quantization.moe=FP8_BLOCK ``` Enable matching Megatron QAT with: ```bash actor_rollout_ref.actor.megatron.qat.enable=true \ actor_rollout_ref.actor.megatron.qat.mode=fp8 ``` Dense and routed-expert precision can be configured independently. For example, use BF16 dense layers and MXFP4 experts with: ```bash actor_rollout_ref.rollout.quantization.linear=null \ actor_rollout_ref.rollout.quantization.moe=MXFP4 \ actor_rollout_ref.actor.megatron.qat.enable=true \ actor_rollout_ref.actor.megatron.qat.mode=mxfp4 ``` TRT-LLM no longer accepts `rollout.quantization`; use vLLM or SGLang for supported low-precision rollout configurations. ### Design & Code Changes - Introduce `WeightQuantizer` dispatch based on rollout overrides or the model's native Hugging Face quantization config. - Export canonical HF-named tensors from Megatron, quantize them on the trainer side, and transfer the final checkpoint layout through the checkpoint engine. - Build compressed-tensors configs from per-layer presets and enforce QAT/rollout precision compatibility at startup. - Add vLLM staging and post-load processing for FP8, MXFP8, MXFP4, and NVFP4 refits. - Pass compressed-tensors model overrides to SGLang for FP8_BLOCK rollout. - Keep default high-precision layers unquantized for both generated compressed-tensors configs and native FP8 configs. - Consolidate low-precision documentation and provide `scripts/rollout_logprob_diff.py` for numerical comparison. - Remove superseded ModelOpt/vLLM monkey-patch utilities and the TRT-LLM FP8 E2E job. ### Checklist Before Submitting - [x] Read the [Contribute Guide](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md). - [x] Apply [pre-commit checks](https://github.com/verl-project/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting). - [x] Add / update [the documentation](https://github.com/verl-project/verl/tree/main/docs). - [ ] Add unit or end-to-end tests to CI to cover all code. New CPU tests and updated FP8 E2E jobs are included, but ModelOpt/compressed-tensors test dependency coverage and MXFP4/NVFP4 E2E coverage still need confirmation. - [ ] Once this PR is ready for CI, send a message in the [`ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1). - [x] Recipe submodule update is not applicable; this PR does not change the recipe submodule. --------- Co-authored-by: Cursor <cursoragent@cursor.com>
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?
FSDP NVFP4 QAT weight export fails with the compressed-tensors versions in the current lockfile (0.17 and 0.18).
QATQuantizerimports the removedcompressors.quantized_compressors.fp4_quantizedmodule and calls the removedcompress_weightmethod when the engine exports rollout weights.Import the compressor from its public namespace and pass the existing weight, block scale, and global scale through
compress(state_dict, QuantizationScheme). Preserve the streamed output names, scale fusion, passthrough weights, and W4A4 input scales.Checklist Before Starting
Test
On RTX 5090 with PyTorch 2.13.0+cu130, the new tests pass with compressed-tensors 0.17.0 and 0.18.0: 12 passed per version. Unmodified main fails to import the quantizer with both versions.
The regression covers FP32/FP16/BF16 weights, W4A16/W4A4, CPU/GPU destinations, multiple decoder layers, QKV/GateUp scale fusion, zero blocks, excluded weights, exact source preservation, and actual packed-weight decompression against an independent E2M1 rounding reference.
Additional local integration checks use a real two-rank FSDP2 QAT engine and a locally constructed tiny Qwen2. Both dependency versions pass with CPU offload on/off: two training steps, export after each step, all seven projections packed and decompressed, source parameters preserved exactly, and exported weights updated after the next optimizer step. These validate the training/export path; they do not run a vLLM generation workload.
API and Usage Example
No configuration change. Existing FSDP NVFP4 QAT rollout-weight export uses the current compressor API.
Checklist Before Submitting
AI assistance was used for investigation, implementation, validation, and drafting.