diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 083842c932..677e37ff8b 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -229,7 +229,7 @@ mkdocs.yaml @hmellor # Reinforcement learning /vllm/distributed/weight_transfer @aoshen02 -/vllm/entrypoints/serve/dev/rlhf @njhill @aoshen02 +/vllm/entrypoints/rl @njhill @aoshen02 /tests/distributed/test_weight_transfer.py @aoshen02 /tests/entrypoints/rl @DarkLight1337 @robertgshaw2-redhat @njhill @aarnphm @NickLucche @AndreasKaratzas @aoshen02 /tests/v1/worker/test_gpu_worker_weight_transfer.py @aoshen02 diff --git a/tests/entrypoints/unit_tests/test_dev_route_request_bodies.py b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py similarity index 96% rename from tests/entrypoints/unit_tests/test_dev_route_request_bodies.py rename to tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py index 78fe5eef5d..b51d08db07 100644 --- a/tests/entrypoints/unit_tests/test_dev_route_request_bodies.py +++ b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py @@ -9,7 +9,7 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from vllm.entrypoints.serve.dev.rlhf.api_router import router as rlhf_router +from vllm.entrypoints.rl.online.api_router import router as rl_router from vllm.entrypoints.serve.dev.rpc.api_router import router as rpc_router from vllm.entrypoints.serve.exception_handling.register import init_exception_handler @@ -30,7 +30,7 @@ class Engine: @pytest.fixture def client(): app = FastAPI() - app.include_router(rlhf_router) + app.include_router(rl_router) app.include_router(rpc_router) init_exception_handler(app) app.state.engine_client = Engine() diff --git a/vllm/entrypoints/serve/dev/rlhf/__init__.py b/tests/entrypoints/rl/weight_transfer_metrics/__init__.py similarity index 100% rename from vllm/entrypoints/serve/dev/rlhf/__init__.py rename to tests/entrypoints/rl/weight_transfer_metrics/__init__.py diff --git a/tests/entrypoints/rl/test_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_across_api_servers_metrics.py similarity index 100% rename from tests/entrypoints/rl/test_weight_operation_metrics.py rename to tests/entrypoints/rl/weight_transfer_metrics/test_across_api_servers_metrics.py diff --git a/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py similarity index 90% rename from tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py rename to tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py index bdddf3d247..22ec5e6112 100644 --- a/tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py +++ b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py @@ -13,9 +13,9 @@ from fastapi import FastAPI from fastapi.testclient import TestClient from prometheus_client import CollectorRegistry -from vllm.entrypoints.serve.dev.rlhf import api_router -from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics -from vllm.entrypoints.serve.dev.rlhf.metrics import WeightOperationMetrics +from vllm.entrypoints.rl.online import api_router +from vllm.entrypoints.rl.online import metrics as rl_metrics +from vllm.entrypoints.rl.online.metrics import WeightOperationMetrics from vllm.entrypoints.serve.exception_handling.register import init_exception_handler from vllm.v1.metrics.prometheus import get_prometheus_registry @@ -52,13 +52,13 @@ class Engine: @pytest.fixture def registry(monkeypatch): registry = CollectorRegistry() - monkeypatch.setattr(rlhf_metrics, "_metrics", WeightOperationMetrics(registry)) + monkeypatch.setattr(rl_metrics, "_metrics", WeightOperationMetrics(registry)) return registry @pytest.mark.parametrize("fail", [False, True]) def test_record_releases_in_flight_and_observes_duration(registry, fail): - metrics = rlhf_metrics.weight_operation_metrics() + metrics = rl_metrics.weight_operation_metrics() with ( pytest.raises(RuntimeError) if fail else nullcontext(), metrics.record("update"), @@ -108,7 +108,7 @@ async def test_cancelled_route_releases_in_flight(registry): HOLD_OPERATION = """ import sys -from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics +from vllm.entrypoints.rl.online.metrics import weight_operation_metrics with weight_operation_metrics().record("update"): print("entered", flush=True) sys.stdin.readline() @@ -163,19 +163,19 @@ def test_default_recorder_is_created_lazily_on_the_default_registry(): probe = """ from prometheus_client import REGISTRY -from vllm.entrypoints.serve.dev.rlhf import metrics as rlhf_metrics +from vllm.entrypoints.rl.online import metrics as rl_metrics NAMES = [ "vllm:rl_weight_update_operation_duration_seconds", "vllm:rl_weight_update_operations_in_flight", ] -assert rlhf_metrics._metrics is None +assert rl_metrics._metrics is None assert not set(NAMES) & set(REGISTRY._names_to_collectors) -metrics = rlhf_metrics.weight_operation_metrics() +metrics = rl_metrics.weight_operation_metrics() assert REGISTRY._names_to_collectors[NAMES[0]] is metrics.duration assert REGISTRY._names_to_collectors[NAMES[1]] is metrics.in_flight -assert rlhf_metrics.weight_operation_metrics() is metrics +assert rl_metrics.weight_operation_metrics() is metrics print("lazy registration ok") """ result = subprocess.run( diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 1bd2f98937..745e67779e 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/vllm/distributed/weight_transfer/clients.py @@ -58,7 +58,7 @@ def _json_safe_update_payload( class HTTPVLLMWeightSyncClient: """Talks to a vLLM server over the RLHF HTTP routes. - Mirrors `vllm/entrypoints/serve/dev/rlhf/api_router.py`: + Mirrors `vllm/entrypoints/rl/online/api_router.py`: `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`, `/finish_weight_update`. """ diff --git a/vllm/entrypoints/launchers/api_server/routers.py b/vllm/entrypoints/launchers/api_server/routers.py index 39d670c019..40d04862af 100644 --- a/vllm/entrypoints/launchers/api_server/routers.py +++ b/vllm/entrypoints/launchers/api_server/routers.py @@ -36,6 +36,10 @@ def register_api_routers( register_vllm_dev_api_routers(app) + from vllm.entrypoints.rl.factories import register_rl_api_routers + + register_rl_api_routers(app) + if "generate" in supported_tasks: from vllm.entrypoints.generate.api_router import ( register_generate_api_routers, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 45b6aede58..c38789a2e2 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -28,10 +28,6 @@ from vllm.config.model import ( TokenizerMode, ) from vllm.config.quantization import QuantizationConfigArgs -from vllm.distributed.weight_transfer.base import ( - WeightTransferInitRequest, - WeightTransferUpdateRequest, -) from vllm.engine.arg_utils import EngineArgs from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, @@ -40,6 +36,7 @@ from vllm.entrypoints.chat_utils import ( ) from vllm.entrypoints.generate.beam_search.offline import BeamSearchOfflineMixin from vllm.entrypoints.pooling.offline import PoolingOfflineMixin +from vllm.entrypoints.rl.offline import RLOfflineMixin from vllm.entrypoints.serve.utils.api_utils import log_non_default_args from vllm.inputs import PromptType from vllm.logger import configure_logging_if_needed, init_logger @@ -64,7 +61,9 @@ if TYPE_CHECKING: logger = init_logger(__name__) -class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): +class LLM( + BeamSearchOfflineMixin, PoolingOfflineMixin, RLOfflineMixin, OfflineInferenceMixin +): """An LLM for generating texts from given prompts and sampling parameters. This class includes a tokenizer, a language model (possibly distributed @@ -884,60 +883,6 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): """ return self.llm_engine.get_metrics() - def init_weight_transfer_engine( - self, request: WeightTransferInitRequest | dict - ) -> None: - """Initialize weight transfer for RL training. - - Args: - request: Weight transfer initialization request with backend-specific info - - """ - init_info_dict = ( - request["init_info"] if isinstance(request, dict) else request.init_info - ) - - self.llm_engine.collective_rpc( - "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} - ) - - def start_weight_update(self) -> None: - """Start a new weight update.""" - self.llm_engine.collective_rpc("start_weight_update") - - def start_draft_weight_update(self) -> None: - """Start a new weight update targeting the speculative draft model.""" - self.llm_engine.collective_rpc("start_draft_weight_update") - - def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: - """Update the weights of the model. - - Args: - request: Weight update request with backend-specific update info - - """ - update_info_dict = ( - request["update_info"] if isinstance(request, dict) else request.update_info - ) - - self.llm_engine.collective_rpc( - "update_weights", kwargs={"update_info": update_info_dict} - ) - - def finish_weight_update(self, weight_version: str | None = None) -> None: - """Finish the weight update and set its version if provided.""" - self.llm_engine.collective_rpc("finish_weight_update") - if weight_version is not None: - self.llm_engine.set_weight_version(weight_version) - - def update_weight_version(self, new_version: str) -> None: - """Set the weight version without updating weights.""" - self.llm_engine.set_weight_version(new_version) - - def get_weight_version(self) -> str: - """Return the latest committed weight version.""" - return self.llm_engine.get_weight_version() - def __repr__(self) -> str: """Return a transformers-style hierarchical view of the model.""" # Cache the result to avoid repeated collective_rpc calls diff --git a/vllm/entrypoints/rl/__init__.py b/vllm/entrypoints/rl/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/vllm/entrypoints/rl/factories.py b/vllm/entrypoints/rl/factories.py new file mode 100644 index 0000000000..a5d237ec0b --- /dev/null +++ b/vllm/entrypoints/rl/factories.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +from fastapi import FastAPI + + +def register_rl_api_routers(app: FastAPI): + from .online.api_router import router as rl_router + + app.include_router(rl_router) diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py new file mode 100644 index 0000000000..02fc63104d --- /dev/null +++ b/vllm/entrypoints/rl/offline.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) +from vllm.entrypoints.offline_utils import OfflineInferenceMixin + + +class RLOfflineMixin(OfflineInferenceMixin): + def init_weight_transfer_engine( + self, request: WeightTransferInitRequest | dict + ) -> None: + """Initialize weight transfer for RL training. + + Args: + request: Weight transfer initialization request with backend-specific info + + """ + init_info_dict = ( + request["init_info"] if isinstance(request, dict) else request.init_info + ) + + self.llm_engine.collective_rpc( + "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} + ) + + def start_weight_update(self) -> None: + """Start a new weight update.""" + self.llm_engine.collective_rpc("start_weight_update") + + def start_draft_weight_update(self) -> None: + """Start a new weight update targeting the speculative draft model.""" + self.llm_engine.collective_rpc("start_draft_weight_update") + + def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None: + """Update the weights of the model. + + Args: + request: Weight update request with backend-specific update info + + """ + update_info_dict = ( + request["update_info"] if isinstance(request, dict) else request.update_info + ) + + self.llm_engine.collective_rpc( + "update_weights", kwargs={"update_info": update_info_dict} + ) + + def finish_weight_update(self, weight_version: str | None = None) -> None: + """Finish the weight update and set its version if provided.""" + self.llm_engine.collective_rpc("finish_weight_update") + if weight_version is not None: + self.llm_engine.set_weight_version(weight_version) + + def update_weight_version(self, new_version: str) -> None: + """Set the weight version without updating weights.""" + self.llm_engine.set_weight_version(new_version) + + def get_weight_version(self) -> str: + """Return the latest committed weight version.""" + return self.llm_engine.get_weight_version() diff --git a/vllm/entrypoints/rl/online/__init__.py b/vllm/entrypoints/rl/online/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/rl/online/api_router.py similarity index 97% rename from vllm/entrypoints/serve/dev/rlhf/api_router.py rename to vllm/entrypoints/rl/online/api_router.py index 7e52e51217..ef9269749f 100644 --- a/vllm/entrypoints/serve/dev/rlhf/api_router.py +++ b/vllm/entrypoints/rl/online/api_router.py @@ -5,7 +5,7 @@ import json from http import HTTPStatus from typing import Annotated, Any, Literal -from fastapi import APIRouter, Body, FastAPI, HTTPException, Query, Request +from fastapi import APIRouter, Body, HTTPException, Query, Request from fastapi.responses import JSONResponse from vllm.distributed.weight_transfer.base import ( @@ -14,7 +14,7 @@ from vllm.distributed.weight_transfer.base import ( WeightTransferUpdateRequest, ) from vllm.engine.protocol import EngineClient -from vllm.entrypoints.serve.dev.rlhf.metrics import weight_operation_metrics +from vllm.entrypoints.rl.online.metrics import weight_operation_metrics from vllm.logger import init_logger from vllm.v1.engine import PauseMode @@ -280,7 +280,3 @@ async def get_world_size( else: world_size = parallel_config.world_size return JSONResponse(content={"world_size": world_size}) - - -def attach_router(app: FastAPI): - app.include_router(router) diff --git a/vllm/entrypoints/serve/dev/rlhf/metrics.py b/vllm/entrypoints/rl/online/metrics.py similarity index 100% rename from vllm/entrypoints/serve/dev/rlhf/metrics.py rename to vllm/entrypoints/rl/online/metrics.py diff --git a/vllm/entrypoints/serve/__init__.py b/vllm/entrypoints/serve/__init__.py index c358b9a6db..d8254cb335 100644 --- a/vllm/entrypoints/serve/__init__.py +++ b/vllm/entrypoints/serve/__init__.py @@ -50,10 +50,6 @@ def register_vllm_dev_api_routers(app: FastAPI): attach_cache_router(app) - from .dev.rlhf.api_router import attach_router as attach_rlhf_router - - attach_rlhf_router(app) - from .dev.rpc.api_router import attach_router as attach_rpc_router attach_rpc_router(app)