Add a vllm_llm_kwargs passthrough to the vLLM LLM() constructor

[behroozazarkhalili]

What this does

Closes #6776.

VLLMGeneration._init_vllm builds the colocate LLM(...) from a fixed set of explicit kwargs, so any engine argument TRL does not expose a field for is unreachable. Text-only Gemma 4 SFT checkpoints declare Gemma4ForConditionalGeneration, so vLLM takes the multimodal path and loading fails on a preprocessor_config.json the checkpoint never ships. hf_overrides fixes it at the vLLM level, but there was no way to pass it, so @marksverdhei is monkeypatching vllm.LLM.__init__.

This adds vllm_llm_kwargs, a dict merged over the explicit kwargs. It mirrors generation_kwargs, which already merges over the SamplingParams defaults with the same documented rule that conflicting keys override.

--vllm_llm_kwargs '{"hf_overrides": {"architectures": ["Gemma4ForCausalLM"]}}'

Why generic rather than a narrow vllm_hf_overrides

I asked this on the issue and did not get a ruling, so I built the option I argued for there and I am happy to rename if you prefer the narrow one. The generic knob covers hf_overrides and every future engine arg without another field each time, and VLLMGeneration is shared by 8 trainers, so one parameter reaches all of them. A narrow field would need adding per-argument, per-config, forever.

Coverage

grep for = LLM( across trl/ returns no construction site outside these three:

Site	Trainers reached
trl/generation/vllm_generation.py	GRPO, RLOO, Distillation, plus GOLD / IW-OPD / SDFT / SDPO / SSD
trl/experimental/online_dpo/online_dpo_trainer.py	Online DPO, which never migrated to VLLMGeneration and builds LLM(...) itself
trl/scripts/vllm_serve.py	server path, as a JSON string matching the existing speculative_config flag

Two keys are rejected, not merged

tensor_parallel_size and enable_sleep_mode raise if present. Both are read back by TRL after construction, so overriding only the engine side leaves the two out of step:

• tensor_parallel_size builds the TP process group, and drives the prompt/image all_gather_object, the PEFT barrier, and the output slicing. With a TP=2 engine and self.tensor_parallel_size == 1, every one of those branches evaluates false and each rank submits its own prompts to a shared engine. That is silent, not a crash, which is why it is worth three lines to make impossible.
• enable_sleep_mode drives the sleep/wake cycle and _llm_weights_sleeping. Disabling it engine-side while TRL still calls sleep(level=2) raises inside vLLM.

The server rejects tensor_parallel_size for the same reason: its weight-sync group world size is derived from it. Every other engine arg is pure passthrough with no read-back, so nothing else needs guarding.

Notes on the field type

dict[str, Any] | str | None and an entry in _VALID_DICT_FIELDS are both required for a dict field to be settable from the command line. _VALID_DICT_FIELDS is consumed in __post_init__, which only json-loads a value that is already a string, and argparse runs first: without | str in the annotation the field registers as type=dict and argparse rejects the JSON before __post_init__ is reached. This is the same class @qgallouedec described in #6791.

Testing

Not added to the test suite, since exercising the merge needs a real vLLM engine. Verified locally:

• --vllm_llm_kwargs '{...}' round-trips to a dict on all 9 configs, via both the CLI parser and direct dataclass construction
• the reserved-key guard raises for both keys and accepts hf_overrides
• ruff check and ruff format --check clean on all 20 files

I cannot test against a Gemma 4 text-only checkpoint. @marksverdhei, if this lands, could you confirm it removes the monkeypatch on your GH200 setup?

---

Note

Low Risk
Mostly additive config plumbing; reserved-key validation reduces risk of silent distributed/sleep desync.

Overview
Adds vllm_llm_kwargs so colocated vLLM setups can pass extra LLM(...) arguments (e.g. hf_overrides for text-only Gemma 4 checkpoints) without monkeypatching vLLM.

The option is wired through nine training configs (GRPO, RLOO, Distillation, GOLD, IW-OPD, Online DPO, SDFT, SDPO, SSD) with _VALID_DICT_FIELDS so JSON CLI values parse correctly. Trainers forward it to VLLMGeneration as llm_kwargs, which merges overrides onto TRL’s defaults before LLM(**kwargs). Online DPO merges the same way in its local LLM build; trl vllm-serve adds a JSON --llm_kwargs flag.

tensor_parallel_size and enable_sleep_mode are rejected inside the extra kwargs (and tensor_parallel_size on the server) so TRL’s TP groups, output slicing, and sleep/wake logic stay aligned with the engine.

^{Reviewed by Cursor Bugbot for commit 0a2f3a2. 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.