[vllm] feat: Support online quant for rollout with torchao

[jerryzh168]

What does this PR do?

Add support for torchao online quantization for vllm, configure the type of quantization by serializing a config file. See vllm PR for how to generate the quantization file (vllm-project/vllm#23014).
Requires vllm changes: vllm-project/vllm#23014 and vllm-project/vllm#26327

Test

1. generate the torchao config file (can change to other configs: https://docs.pytorch.org/ao/main/api_ref_quantization.html#inference-apis-for-quantize) torchao_config.json

from torchao.quantization import Float8DynamicActivationFloat8WeightConfig, PerRow
from torchao.core.config import config_to_dict
import json

config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())

with open("torchao_config.json", "w") as f:
    f.write(json.dumps(config_to_dict(config)))

# LLM(..., quantization="torchao", hf_overrides={"quantization_config_file": "torchao_config.json"})

this is fp8 dynamic quant:

{"_type": "Float8DynamicActivationFloat8WeightConfig", "_version": 2, "_data": {"activation_dtype": {"_type": "torch.dtype", "_data": "float8_e4m3fn"}, "weight_dtype": {"_type": "torch.dtype", "_data": "float8_e4m3fn"}, "granularity": [{"_type": "PerRow", "_version": 1, "_data": {}}, {"_type": "PerRow", "_version": 1, "_data": {}}], "mm_config": {"_type": "Float8MMConfig", "_version": 1, "_data": {"emulate": false, "use_fast_accum": true, "pad_inner_dim": false}}, "activation_value_lb": null, "activation_value_ub": null, "kernel_preference": {"_type": "KernelPreference", "_data": "AUTO"}, "set_inductor_config": true}}

2. Add following to sh examples/ppo_trainer/run_deepseek7b_llm.sh

     actor_rollout_ref.rollout.quantization=torchao \
     actor_rollout_ref.rollout.quantization_config_file=torchao_config.json \

3. Run test
   VLLM_DISABLE_COMPILE_CACHE=1 sh examples/ppo_trainer/run_deepseek7b_llm.sh

# baseline
(TaskRunner pid=539843) ("Initial validation metrics: {'val-aux/openai/gsm8k/reward/mean@1': "
(TaskRunner pid=539843)  "0.6717210007581501, 'val-core/openai/gsm8k/acc/mean@1': 0.6717210007581501, "
(TaskRunner pid=539843)  "'val-aux/num_turns/min': 2, 'val-aux/num_turns/max': 2, "
(TaskRunner pid=539843)  "'val-aux/num_turns/mean': 2.0}")
(TaskRunner pid=539843) step:105 - val-aux/openai/gsm8k/reward/mean@1:0.6717210007581501 - val-core/openai/gsm8k/acc/mean@1:0.6717210007581501 - val-aux/num_turns/min:2 - val-aux/num_turns/max:2 - val-aux/num_turns/mean:2.0



# fp8
(TaskRunner pid=3763210) validation generation end
(TaskRunner pid=3763210) ("Initial validation metrics: {'val-aux/openai/gsm8k/reward/mean@1': "
(TaskRunner pid=3763210)  "0.6739954510993177, 'val-core/openai/gsm8k/acc/mean@1': 0.6739954510993177, "
(TaskRunner pid=3763210)  "'val-aux/num_turns/min': 2, 'val-aux/num_turns/max': 2, "
(TaskRunner pid=3763210)  "'val-aux/num_turns/mean': 2.0}")
(TaskRunner pid=3763210) step:105 - val-aux/openai/gsm8k/reward/mean@1:0.6739954510993177 - val-core/openai/gsm8k/acc/mean@1:0.6739954510993177 - val-aux/num_turns/min:2 - val-aux/num_turns/max:2 - val-aux/num_turns/mean:2.0

Docs:

no docs added yet since I didn't find a place to add quantized rollout docs in https://github.com/volcengine/verl/blob/main/docs/workers/fsdp_workers.rst, happy to add later when there are more docs

We can add simple string options (e.g. fp8_tensorwise, fp8_rowwise, fp8_blockwise etc.) in the future if needed.

Reviewers:

Subscribers:

Tasks:

Tags:

Checklist Before Starting

• Search for similar PRs. https://github.com/volcengine/verl/pulls?q=sort%3Aupdated-desc+is%3Apr+is%3Aopen+quantization+
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)

[gemini-code-assist]

Code Review

This pull request introduces support for on-the-fly quantization for vLLM rollouts using torchao. The changes involve adding configuration options for quantization and implementing the quantization logic within the FSDP sharding manager.

My review identified a critical issue in the implementation where model weights would fail to load if quantization is disabled. I have provided a code suggestion to fix this. Additionally, I've pointed out that the current method for selecting layers to quantize is too specific and may miss some linear layers, which could lead to unexpected behavior.

[jerryzh168]

please let me know if the API makes sense, I can clean up both PRs after confirmation

[vadimkantorov]

I also wonder how torchao fp8 quantization compares to vllm's own impl for quantization="fp8" (and related) - and curious, why does vllm not use torchao for this (IIUC there are several "quantization backends" in vllm?), why do they have to use CUDA kernels https://github.com/vllm-project/vllm/tree/main/csrc/quantization/fp8 (why Triton is not sufficient) and why are these kernels not upstreamed to PyTorch core :) Hoping that at least on kernel level, fragmentation can be reduced by upstreaming more of these https://github.com/vllm-project/vllm/tree/main/csrc/quantization :)

And also, how does the approach in this PR compare to FlashRL approach (which also patches vllm)

[jerryzh168]

I also wonder how torchao fp8 quantization compares to vllm's own impl for quantization="fp8" (and related) - and curious, why does vllm not use torchao for this (IIUC there are several "quantization backends" in vllm?), why do they have to use CUDA kernels vllm-project/vllm@main/csrc/quantization/fp8 (why Triton is not sufficient) and why are these kernels not upstreamed to PyTorch core :) Hoping that at least on kernel level, fragmentation can be reduced by upstreaming more of these vllm-project/vllm@main/csrc/quantization :)

we haven't officially compared with them yet, but we are integrating with fbgemm kernels which should be SOTA.
why does vllm not use torchao: torchao is integrated into vllm recently: https://docs.vllm.ai/en/latest/features/quantization/torchao.html (a few months ago), and we are actively working on improving torchao for vllm users. I don't have context on vllm-project/vllm@main/csrc/quantization/fp8 though.

And also, how does the approach in this PR compare to FlashRL approach (which also patches vllm)

also haven't compared yet, but we can discuss what would be the API that makes most of sense I think

[CLAassistant]

All committers have signed the CLA.

[wuxibin89]

Please add comment above quantization and quantization_config_file.