Skip to content

[fsdp] fix: use current compressed-tensors API for QAT export - #8205

Closed
cyanseek wants to merge 1 commit into
verl-project:mainfrom
cyanseek:codex/nvfp4-compressor-api
Closed

cyanseek wants to merge 1 commit into
verl-project:mainfrom
cyanseek:codex/nvfp4-compressor-api

Conversation

@cyanseek

Copy link
Copy Markdown

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). QATQuantizer imports the removed compressors.quantized_compressors.fp4_quantized module and calls the removed compress_weight method 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

python -m pytest tests/utils/test_qat_quantizer.py -q
pre-commit run --all-files --show-diff-on-failure

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

  • Read the contribution guide.
  • Added GPU regression tests discovered by the existing GPU unit-test workflow.
  • All 14 repository pre-commit checks passed on the final patch.
  • Request upstream CI through the project channel when ready.
  • No recipe submodule change.

AI assistance was used for investigation, implementation, validation, and drafting.

@coderabbitai

coderabbitai Bot commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
  • Configuration used: Organization UI
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 764f2200-f556-404e-aefc-189700224240

📥 Commits

Reviewing files that changed from the base of the PR and between 2035f58 and 1553f03.


📒 Files selected for processing (2)
  • tests/utils/test_qat_quantizer.py
  • verl/utils/qat/quantizer.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 9 remain after this review.



📝 Summary

Summary by CodeRabbit

  • Bug Fixes
    • Improved quantized weight compression and stream round-tripping across supported floating-point formats and CPU/CUDA targets.
    • Preserved input tensors and input scaling during quantization, and ensured output weights and scales retain expected values and device placement.

Walkthrough

QATQuantizer now uses the package-level compressor API with a QuantizationScheme. New parameterized tests check stream round trips across quantization modes, input dtypes, and target devices.

Changes

QAT quantizer compression

Layer / File(s) Summary
Compressor API and round-trip validation
verl/utils/qat/quantizer.py, tests/utils/test_qat_quantizer.py
QATQuantizer creates a QuantizationScheme for Linear layers and passes weights and scales to NVFP4PackedCompressor.compress. Parameterized tests check packed outputs, decompressed FP4 values, global scales, output keys, input preservation, and target-device placement.

Priority: ➖ Normal

Severity of issue fixed: Medium

Merge Risk | ⚪ Minimal · up to 1553f

Merge Risk: ⚪ Minimal · up to 1553f

The supported CPU test workflow is unaffected, and no actionable merge-blocking issue remains after normal checks.

Security Architecture Review

Security architecture risk: 🔵 Low · up to 1553f

The existing weight-export boundary and output format appear preserved, with no demonstrated expansion of attacker access or authority. Dependency behavior during failed, interrupted, or concurrent exports remains unverified.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

  • inferred — The verified impact follows QAT-enabled FSDP weight exports into their existing rollout consumers. No new tenant, service, credential, or environmental authority is established by this change. Deployment-wide exposure cannot be determined from the available topology information.

Trust Boundaries and Controls

  • observed — The inspected transition remains model state tensors to compressed tensor mappings to existing weight loaders. The API migration does not introduce an identity-bearing input or alter the inspected caller's authority. This supports unchanged repository-visible trust boundaries, not a claim that all model-input or deployment surfaces are safe.

Resilience and Maintainability Implications

  • observed — The production export caller constructs a fresh quantizer for each export, and stream buffers are invocation-local. This limits shared-instance concurrency exposure. Direct concurrent reuse would still share the compressor and scheme; their reentrancy and exceptional-path mutation behavior remain unverified because dependency source was unavailable.

Pre-merge checks | Passed 4 | Failed 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check Passed The title clearly and concisely identifies the FSDP QAT export fix and the migration to the current compressed-tensors API.
Description check Passed The description directly explains the compatibility failure, the API changes, preserved behavior, and validation performed for the QAT export path.
Linked Issues check Passed No active directly linked issue remains. Issue #6815 is closed and provides historical context only. Therefore, no linked-issue coding requirements apply to this pull request.
Out of Scope Changes check Passed The changes are within the stated scope of PR #8205. QATQuantizer uses the current public compressed-tensors APIs, and the added tests cover the affected export behavior. No unrelated change is iden…

  • Fix all pre-merge checks with AI
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

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>
@wuxibin89 wuxibin89 closed this Oct 20, 2026
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.

2 participants