Add a vllm_llm_kwargs passthrough for colocate vLLM generation

[behroozazarkhalili]

What this does

Closes #6776.

VLLMGeneration builds the vLLM engine itself in colocate mode, and the constructor call at trl/generation/vllm_generation.py:347 only receives the arguments TRL happens to expose. Engine settings TRL has no flag for, hf_overrides being the case in #6776, cannot be reached at all.

This adds a vllm_llm_kwargs dict that is merged into those constructor kwargs, with keys in the dict overriding TRL's defaults. It is wired through the eleven trainer configs that drive colocate generation.

Two keys are rejected rather than forwarded:

for key in ("tensor_parallel_size", "enable_sleep_mode"):
    if key in self.llm_kwargs:
        raise ValueError(f"`{key}` cannot be set in `vllm_llm_kwargs`; use `vllm_{key}` instead.")

TRL reads both back to build the tensor-parallel process group, slice generation outputs, and drive the sleep and wake cycle. Setting only the engine side would leave the two out of step, and the failure would surface far from the cause.

Scope change after merging main

This PR originally covered two paths. The server-side one is gone.

trl/scripts/vllm_serve.py used to build the engine in process through llm_worker, and this branch added an --llm_kwargs JSON flag to it. Main has since replaced that with build_command, which shells out to vllm serve. Unrecognized arguments are already forwarded verbatim:

script_args, extra_args = parser.parse_args_and_config(return_remaining_strings=True)   # :306
return command + list(extra_args or [])                                                # :270

So trl vllm-serve --model Qwen/Qwen3-0.6B --hf-overrides '{...}' reaches the engine today with no flag from this PR. Adding a JSON knob alongside that would be a second, worse spelling of the same thing, so the server half is dropped. vllm_serve.py is no longer touched by this branch.

The colocate path has no such route. LLM() is still called directly there and nothing on the command line reaches it, so the config field is the only way in and that half is kept.

Tests

Five tests in tests/test_vllm_client_server.py, all against VLLMGeneration with LLM replaced by a recorder:

• test_llm_kwargs_reach_the_llm_constructor, over two hf_overrides values
• test_llm_kwargs_override_the_defaults_trl_sets, over two gpu_memory_utilization values
• test_reserved_table_matches_the_pinned_list, which pins RESERVED_LLM_KWARGS to the seven keys spelled out in the test
• test_reserved_keys_are_rejected, parametrized over those seven keys
• test_build_leaves_the_distributed_environment_unchanged, since the constructor writes RANK, LOCAL_RANK and WORLD_SIZE

Two values per passthrough test because a constructor that hardcodes the one value a single case checks would pass that case. The reserved list is spelled out rather than read from the table because a test derived from the table loses its case when a reservation is deleted instead of failing.

llm_kwargs is the last parameter of VLLMGeneration.__init__, after generation_kwargs, so callers that pass the older parameters by position are unaffected.

Verification

ruff check and ruff format --check at the CI-pinned 0.13.3, and the pinned doc-builder at --max-len 119, pass on all 21 changed files. The tests above pass, 13 passed with 40 deselected.

Not verified here: no live vLLM engine was started, so the tests assert the kwargs reach the constructor rather than that vLLM accepts any particular key. Whether a given override is valid remains vLLM's contract.

---

Note

Medium Risk
Touches colocate vLLM engine construction across many trainers; incorrect merges could break generation or weight sync, though reserved keys and tests reduce that risk.

Overview
Adds vllm_llm_kwargs so colocate training can pass extra arguments into vLLM’s LLM() constructor (e.g. hf_overrides) when TRL builds the engine in-process.

VLLMGeneration gains an llm_kwargs parameter: TRL’s defaults are built first, then user kwargs are merged on top. _check_llm_kwargs and RESERVED_LLM_KWARGS block seven keys TRL must own (model, tensor_parallel_size, distributed_executor_backend, seed, logprobs_mode, quantization, enable_sleep_mode) so engine and trainer stay in sync.

The same config field is wired through trainer configs (GOLD, GRPO, RLOO, distillation, Online DPO, SDFT, SDPO, SSD, IW OPD, etc.) into VLLMGeneration; Online DPO also validates and merges kwargs on its direct LLM(**vllm_kwargs) path.

New unit tests mock LLM to assert passthrough, override behavior, reserved-key rejection, and that distributed env vars are restored after init.

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

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.