Using FP8 inference in long response vllm rollouts

[vadimkantorov]

Hi! In our experience, vllm rollouts use 70% of the grpo iteration time (when performed in bf16)

Has anyone tried using more aggressive precision reduction for rollouts (like using FP8) for gaining speed? I wonder if some okay methods exist for online FP8 usage for gaining speed in this rollout phase

Thanks!

[vermouth1992]

The real issue was that the current vllm implementation doesn't support loading bf16 weights online and dynamically quantize them to fp8. We need this feature to land in vllm first and we only need to set the precision in vllm engine without modifying any code.

[vermouth1992]

Another approach is to perform quantization in the sharding manager. But we haven't got a chance to try it

[vadimkantorov]

Another approach is to perform quantization in the sharding manager.

Does verl do weight consolidation on one rank prior to syncing them to the vllm workers?

Probably an example is needed of scaling out the rollout phase on a separate pool of machines (and thus of efficient weight broadcast/online fp8 quantization)... Does the sharding manager handle the weight consolidation/broadcast on the verl side?

If verl uses sth like https://github.com/vllm-project/vllm/blob/ebb1ec931871ee1baca23673049865856c44ce4e/examples/offline_inference/rlhf_utils.py#L53, could it somehow reuse some fp8 quant util from vllm and construct the fp8 weights prior to calling load_weights? Although it would be more robust if vllm supported this directly...

[vermouth1992]

Another approach is to perform quantization in the sharding manager.

Does verl do weight consolidation on one rank prior to syncing them to the vllm workers?

Probably an example is needed of scaling out the rollout phase on a separate pool of machines (and thus of efficient weight broadcast/online fp8 quantization)... Does the sharding manager handle the weight consolidation/broadcast on the verl side?

If verl uses sth like https://github.com/vllm-project/vllm/blob/ebb1ec931871ee1baca23673049865856c44ce4e/examples/offline_inference/rlhf_utils.py#L53, could it somehow reuse some fp8 quant util from vllm and construct the fp8 weights prior to calling load_weights? Although it would be more robust if vllm supported this directly...

It's actually here. https://github.com/volcengine/verl/blob/main/verl/workers/sharding_manager/fsdp_vllm.py#L289. verl does weight allgather per tensor to avoid OOM.

[vermouth1992]

You can write a quantization function like

def quantize_to_fp8(name, param):
    return xxx

and modify this line to call the function

quantize_to_fp8(name, param.to(device, non_blocking=True).full_tensor())

[vadimkantorov]

I wonder if any vllm built-in quant functions exist for this... And I also wonder if at LLM(..., quantization="fp8") any full-model calibration is performed to tweak the quantized weights to match the full perf (e.g. layerwise) - but of course having a basic (but still working in reasoning generation) method would already have big impact in terms of GRPO workload speeds

[vermouth1992]

Agreed. The online quantization should be implemented inside the inference engine to support a class of models and keep the API clean in RL framework.

[vadimkantorov]

keep the API clean in RL framework.

Then I also wonder: could param allgather be fused with some fp8 quantization (to reduce the allgather param memory/communication overweight)?

It also appears that the code in https://github.com/volcengine/verl/blob/main/verl/workers/sharding_manager/fsdp_vllm.py#L289 is specialized for FSDP1?