[rak96]

What does this PR do?

This PR makes vLLM weight synchronization more memory-efficient for PEFT/LoRA models trained with DeepSpeed ZeRO-3.

Previously, the PEFT weight-sync path gathered all model parameters before merging adapters and syncing weights to vLLM. Under DeepSpeed ZeRO-3, this can materialize the full model on every rank and cause OOMs for large models.

This PR changes the non-FSDP PEFT path to gather, merge, sync, and unmerge one BaseTunerLayer module at a time. This avoids a full-model gather while preserving the existing parameter-name normalization and vLLM loading behavior.

It also adds minimal worker exception propagation in vllm_serve.py, so worker-side failures are sent back through the pipe and surfaced by the parent process instead of hanging.

No new dependencies are added.

Changes

• Avoid full-model parameter gathering for PEFT/LoRA + DeepSpeed ZeRO-3 vLLM weight sync.
• Keep the FSDP PEFT path unchanged.
• Preserve existing vLLM parameter-name handling for PEFT layers.
• Propagate worker exceptions from generate(), chat(), and _run_prompt_logprobs() calls.

Testing

• python -m ruff check trl/generation/vllm_generation.py trl/scripts/vllm_serve.py
• python -m ruff format --check trl/generation/vllm_generation.py trl/scripts/vllm_serve.py
• python -m pytest tests/test_vllm_client_server.py::TestChunkList tests/test_vllm_client_server.py::TestExtractLogprobs -q
• python -m pytest tests/test_vllm_client_server.py -q -m "not slow"
• Full vLLM server/GPU tests were not run locally.
• DeepSpeed ZeRO-3 + PEFT integration test was not run locally.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

@qgallouedec @AmineDiro @kashif @remi-or

---

Note

Medium Risk
Touches distributed PEFT/LoRA weight syncing under DeepSpeed ZeRO-3 and changes inter-process error handling; mistakes could cause incorrect weight updates or new runtime failures/hangs during generation.

Overview
Improves PEFT/LoRA weight synchronization to vLLM under DeepSpeed ZeRO-3 by switching from a full-model gather/merge to a module-at-a-time gather+merge+sync+unmerge flow (using BaseTunerLayer/ModulesToSaveWrapper) to avoid per-rank full-model materialization and OOMs.

Keeps the FSDP PEFT path as a full merge+sync via existing FSDP-aware traversal, but ensures unmerge happens within the gathered context.

Updates vllm_serve.py so worker-side exceptions during generate/chat (and prompt-logprob generation) are sent through the pipe and raised in the parent, preventing silent hangs and surfacing failures earlier.

^{Reviewed by Cursor Bugbot for commit 533325d. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit f8076b6. Configure here.}

[verma8076]

Checked both halves against current main.

The trl/scripts/vllm_serve.py half is fully obsolete. generate(), chat(), and _run_prompt_logprobs() don't exist anymore — #6765 ("Use vLLM's own server instead of TRL's", merged 2026-08-26) deleted the entire custom FastAPI server this part of the PR patches, taking the file from 1218 lines to 307. trl vllm-serve now just execs vLLM's own server.

The trl/generation/vllm_generation.py half's underlying bug is real and still there, but the fix needs a rewrite rather than a rebase. The full-model gather this PR avoids is still exactly present on main:

# trl/generation/vllm_generation.py, _iter_named_params(), PEFT + non-FSDP branch
with self._dist.gather_params(list(model.parameters())):
    model.merge_adapter()
    ...
    for name, param in model.named_parameters():
        ...

So DeepSpeed ZeRO-3 + PEFT still materializes every parameter on every rank before syncing, same OOM risk this PR describes. But the surrounding architecture changed since this PR was written: weight sync used to be a sync_weights() push method (self.vllm_client.update_named_param(...) / direct colocate load, matching this PR's diff), and is now an _iter_named_params() generator that yields (name, param.data) pairs for a caller to consume — gather_if_zero3 is gone too, replaced by self._dist.gather_params(...). This PR's diff won't apply, and reapplying its logic verbatim wouldn't fit the new pull-based shape even after conflict resolution — it'd need re-deriving as a module-at-a-time generator (yield per parameter inside the per-BaseTunerLayer gather, instead of pushing per-parameter inside the loop) against the current _iter_named_params().

Worth keeping open as a real, live bug report even though this specific diff isn't mergeable — the module-at-a-time approach here is the right shape for the fix, just needs to be re-expressed as a generator.

[AmineDiro]

@rak96 is it possible to solve the conflicts for review ?

[AmineDiro]

vllm_serve is obsolete now