Skip to content

Move VLLMClient to generation module #4928

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
albertvillanova merged 2 commits into from
Jan 29, 2026
Merged

Move VLLMClient to generation module #4928

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5ac1f36
Move VLLMClient to generation module
albertvillanova Jan 29, 2026
ce6ab3f
Update docstring
albertvillanova Jan 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tests/test_vllm_client_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from transformers import AutoModelForCausalLM
from transformers.testing_utils import torch_device

from trl.extras.vllm_client import VLLMClient
from trl.generation.vllm_client import VLLMClient
from trl.import_utils import is_vllm_available
from trl.scripts.vllm_serve import chunk_list

Expand Down
2 changes: 1 addition & 1 deletion trl/experimental/gold/gold_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@

from ...data_utils import is_conversational, maybe_convert_to_chatml, pack_dataset, truncate_dataset
from ...extras.profiling import profiling_decorator
from ...extras.vllm_client import VLLMClient
from ...generation.vllm_client import VLLMClient
from ...import_utils import is_vllm_available
from ...models import prepare_deepspeed
from ...models.utils import unwrap_model_for_generation
Expand Down
2 changes: 1 addition & 1 deletion trl/experimental/online_dpo/online_dpo_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@

from ...data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template
from ...extras.profiling import profiling_context
from ...extras.vllm_client import VLLMClient
from ...generation.vllm_client import VLLMClient
from ...import_utils import is_vllm_available
from ...models.utils import create_reference_model, prepare_deepspeed, prepare_fsdp, unwrap_model_for_generation
from ...trainer.base_trainer import BaseTrainer
Expand Down
4 changes: 2 additions & 2 deletions trl/extras/vllm_client.py → trl/generation/vllm_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from transformers import is_torch_xpu_available
from urllib3.util.retry import Retry

from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available
from import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available


if is_requests_available():
Expand Down Expand Up @@ -88,7 +88,7 @@ class VLLMClient:
Use the client to generate completions and update model weights:

```python
>>> from trl.extras.vllm_client import VLLMClient
>>> from trl.generation.vllm_client import VLLMClient

>>> client = VLLMClient()
>>> client.generate(["Hello, AI!", "Tell me a joke"])
Expand Down
2 changes: 1 addition & 1 deletion trl/generation/vllm_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,9 @@

from ..data_utils import apply_chat_template, is_conversational, prepare_multimodal_messages_vllm
from ..extras.profiling import ProfilingContext, profiling_decorator
from ..extras.vllm_client import VLLMClient
from ..import_utils import is_vllm_available
from ..trainer.utils import ensure_master_addr_port
from .vllm_client import VLLMClient


if TYPE_CHECKING:
Expand Down
Loading