Add adapter-only vLLM LoRA sync for GRPO and AsyncGRPO
- Dominant language
- Python
- Stars
- 19.3k
- Forks
- 3k
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 194
Description
### Feature request
Add an adapter-only vLLM LoRA sync mode for `GRPOTrainer` and `AsyncGRPOTrainer`.
Today, PEFT models are synchronized to vLLM by merging adapters into the base model and sending full merged weights:
```python
model.merge_adapter()
send merged model weights to vLLM
model.unmerge_adapter()
```
This should remain available, but TRL should also support a LoRA-adapter sync mode where:
1. vLLM starts with the frozen base model and LoRA support enabled.
2. The trainer gathers/saves only the PEFT adapter state dict.
3. The trainer asks vLLM to load or reload the active LoRA adapter.
4. Generation uses the currently active adapter version.
This is similar to Axolotl's `trl.vllm_lora_sync: true` path and verl's `lora.merge=False` adapter path.
### Motivation
The current full-weight PEFT sync path is inefficient for LoRA and problematic for QLoRA + FSDP2.
For QLoRA + FSDP2, the base model is bitsandbytes-quantized and FSDP2 wraps/shards parameters as DTensor. PEFT's bitsandbytes merge path expects bitsandbytes parameters with quantization metadata such as `quant_state`. Calling `merge_adapter()` on DTensor-sharded bitsandbytes weights can fail.
As a result, TRL currently needs to guard this combination:
```text
GRPO/AsyncGRPO + PEFT + QLoRA + FSDP2 + full-weight vLLM sync
```
Adapter-only sync would avoid the failing operation entirely because it would not merge LoRA into quantized FSDP2 base weights. It would also reduce sync payload size and avoid full-weight NCCL sync overhead.
Expected benefits:
- Enables a viable path for QLoRA + FSDP2 + vLLM.
- Avoids `merge_adapter()` on quantized DTensor-sharded base weights.
- Sends small LoRA adapter payloads instead of full model weights.
- Reduces vLLM weight-sync overhead for LoRA/QLoRA GRPO.
- Aligns TRL with Axolotl and verl adapter-sync designs.
Related references:
- Accelerate FSDP2 QLoRA issue: https://github.com/huggingface/accelerate/issues/3874
- PEFT 4-bit `merge_adapter()` issue: https://github.com/huggingface/peft/issues/2501
- PEFT FSDP2 issue: https://github.com/huggingface/peft/issues/2344
- Axolotl FSDP + QLoRA docs: https://docs.axolotl.ai/docs/fsdp_qlora.html
- Axolotl vLLM LoRA sync docs: https://docs.axolotl.ai/docs/vllm_serving.html
- bitsandbytes FSDP-QLoRA docs: https://huggingface.co/docs/bitsandbytes/en/fsdp_qlora
### Proposed implementation sketch
Add a new configuration option, for example:
```python
vllm_lora_sync: bool = False
```
or a more general strategy enum:
```python
vllm_weight_sync_strategy: Literal["merged_weights", "lora_adapter"] = "merged_weights"
```
When adapter sync is enabled:
1. Require PEFT/LoRA model.
2. Require vLLM server mode with LoRA enabled.
3. Gather adapter weights from all trainer ranks.
4. Save adapter with `save_pretrained()` or send equivalent adapter tensors.
5. Call vLLM's LoRA load/reload endpoint, such as `/v1/load_lora_adapter`, or a TRL-managed endpoint if needed.
6. Track adapter version in the existing model-version/staleness logic.
7. Reset vLLM prefix cache after adapter update.
Fallback behavior:
- Dense/full models continue using full-weight sync.
- PEFT models continue using full merged sync unless adapter sync is enabled.
- QLoRA + FSDP2 should remain guarded unless adapter sync is enabled.
### Your contribution
I can help by submitting a PR that:
1. Adds the configuration flag.
2. Implements adapter-only sync in the shared vLLM generation/sync path.
3. Reuses the path in `GRPOTrainer` and `AsyncGRPOTrainer`.
CC: @AmineDiro
Contributor guide
Assessment
This issue has not been assessed yet.