From 6fe4249ed816917086ae35a09f6663ef8ad5dfda Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Mon, 3 Aug 2026 11:32:40 -0700 Subject: [PATCH] fix(inference): map external engine ranks --- src/prime_rl/inference/vllm/ranks.py | 39 +++++++ src/prime_rl/inference/vllm/server.py | 12 +- src/prime_rl/inference/vllm/worker/nccl.py | 27 +++-- src/prime_rl/inference/vllm/worker/nixl.py | 20 +++- src/prime_rl/utils/client.py | 108 +++++++++++++----- tests/unit/inference/test_nccl_rank.py | 90 +++++++++++++++ .../utils/test_external_engine_topology.py | 44 +++++++ 7 files changed, 299 insertions(+), 41 deletions(-) create mode 100644 src/prime_rl/inference/vllm/ranks.py create mode 100644 tests/unit/inference/test_nccl_rank.py create mode 100644 tests/unit/utils/test_external_engine_topology.py diff --git a/src/prime_rl/inference/vllm/ranks.py b/src/prime_rl/inference/vllm/ranks.py new file mode 100644 index 0000000000..8de43f4ace --- /dev/null +++ b/src/prime_rl/inference/vllm/ranks.py @@ -0,0 +1,39 @@ +def global_inference_rank( + *, + rank_offset: int, + data_parallel_index: int, + data_parallel_size: int, + worker_rank: int, + tensor_parallel_size: int, + pipeline_parallel_size: int, + inference_world_size: int, + prefill_context_parallel_size: int = 1, + engine_world_size: int | None = None, +) -> int: + """Map one vLLM worker to its rank in Prime's inference NCCL group.""" + model_parallel_size = tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size + logical_data_parallel_size = data_parallel_size + if engine_world_size is not None: + if engine_world_size <= 0 or engine_world_size % model_parallel_size: + raise ValueError( + f"engine world size {engine_world_size} is not divisible by model parallel size {model_parallel_size}" + ) + logical_data_parallel_size = engine_world_size // model_parallel_size + # Dense vLLM EngineCore processes retain their global DP index but rewrite + # data_parallel_size to one. MoE EngineCore processes preserve the logical + # size, so keep validating that value rather than masking bad discovery. + if data_parallel_size != 1 and data_parallel_size != logical_data_parallel_size: + raise ValueError( + f"data parallel size {data_parallel_size} does not match engine-derived size " + f"{logical_data_parallel_size}" + ) + if not 0 <= data_parallel_index < logical_data_parallel_size: + raise ValueError( + f"data parallel index {data_parallel_index} is outside logical data parallel size " + f"{logical_data_parallel_size}" + ) + + rank = rank_offset + data_parallel_index * model_parallel_size + worker_rank % model_parallel_size + if not 0 <= rank < inference_world_size: + raise ValueError(f"calculated inference rank {rank} is outside inference world size {inference_world_size}") + return rank diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index 28e2f572ff..5077f3003b 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -143,11 +143,21 @@ async def init_broadcaster(request: Request): timeout = data.get("timeout") rank_offset = data.get("rank_offset") inference_world_size = data.get("inference_world_size") + engine_world_size = data.get("engine_world_size") quantize_in_weight_transfer = data.get("quantize_in_weight_transfer", False) session_id = data.get("session_id", "default") await engine_client(request).collective_rpc( "init_broadcaster", - args=(host, port, rank_offset, inference_world_size, timeout, quantize_in_weight_transfer, session_id), + args=( + host, + port, + rank_offset, + inference_world_size, + timeout, + quantize_in_weight_transfer, + session_id, + engine_world_size, + ), ) return {"status": "ok"} diff --git a/src/prime_rl/inference/vllm/worker/nccl.py b/src/prime_rl/inference/vllm/worker/nccl.py index 686bb856cf..4fe743c05e 100644 --- a/src/prime_rl/inference/vllm/worker/nccl.py +++ b/src/prime_rl/inference/vllm/worker/nccl.py @@ -7,6 +7,7 @@ from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger +from prime_rl.inference.vllm.ranks import global_inference_rank from prime_rl.inference.vllm.worker.weight_transfer import ( load_weights_checkpoint_layerwise, load_weights_kernel, @@ -101,27 +102,37 @@ def init_broadcaster( timeout: int, quantize_in_weight_transfer: bool = False, session_id: str = "default", + engine_world_size: int | None = None, ) -> None: """Initialize the NCCL broadcast receiver. Args: rank_offset: Starting GPU offset for this server in the global inference group. inference_world_size: Total number of inference GPUs across all servers. + engine_world_size: Number of inference GPUs assigned to this server. """ del session_id self.quantize_in_weight_transfer = quantize_in_weight_transfer - # Use the worker's device index directly as the local rank. - # The previous dp_group-based computation broke in vLLM v1 multiprocess - # DP mode where each worker is a separate process with a singleton - # DP group (rank_in_group is always 0). - local_rank = self.device.index - global_rank_inference = rank_offset + local_rank + if engine_world_size is None: + global_rank_inference = rank_offset + self.device.index + else: + parallel_config = self.parallel_config + global_rank_inference = global_inference_rank( + rank_offset=rank_offset, + data_parallel_index=parallel_config.data_parallel_index, + data_parallel_size=parallel_config.data_parallel_size, + worker_rank=self.rank, + tensor_parallel_size=parallel_config.tensor_parallel_size, + pipeline_parallel_size=parallel_config.pipeline_parallel_size, + prefill_context_parallel_size=parallel_config.prefill_context_parallel_size, + inference_world_size=inference_world_size, + engine_world_size=engine_world_size, + ) logger.info( - f"Worker [local_rank={local_rank} rank_offset={rank_offset}] " + f"Worker [worker_rank={self.rank} rank_offset={rank_offset}] " f"-> [global_rank={global_rank_inference} inference_world_size={inference_world_size}]" ) - self.nccl_broadcast_receiver = NCCLWeightBroadcastReceiver( host=host, port=port, diff --git a/src/prime_rl/inference/vllm/worker/nixl.py b/src/prime_rl/inference/vllm/worker/nixl.py index d6990211a1..5959b3a6bc 100644 --- a/src/prime_rl/inference/vllm/worker/nixl.py +++ b/src/prime_rl/inference/vllm/worker/nixl.py @@ -18,6 +18,7 @@ from vllm.config import set_current_vllm_config from vllm.logger import init_logger +from prime_rl.inference.vllm.ranks import global_inference_rank from prime_rl.inference.vllm.worker.weight_transfer import update_mla_absorbed_weights from prime_rl.trainer.rl.broadcast.nixl.agent import MemDesc, NixlAgent, make_agent_name, set_ucx_env_defaults from prime_rl.trainer.rl.broadcast.nixl.cuda_malloc_memory import ( @@ -96,9 +97,24 @@ def init_broadcaster( timeout: int, quantize_in_weight_transfer: bool = False, session_id: str = "default", + engine_world_size: int | None = None, ) -> None: - del inference_world_size, quantize_in_weight_transfer - global_rank = rank_offset + self.device.index + del quantize_in_weight_transfer + if engine_world_size is None: + global_rank = rank_offset + self.device.index + else: + parallel_config = self.parallel_config + global_rank = global_inference_rank( + rank_offset=rank_offset, + data_parallel_index=parallel_config.data_parallel_index, + data_parallel_size=parallel_config.data_parallel_size, + worker_rank=self.rank, + tensor_parallel_size=parallel_config.tensor_parallel_size, + pipeline_parallel_size=parallel_config.pipeline_parallel_size, + prefill_context_parallel_size=parallel_config.prefill_context_parallel_size, + inference_world_size=inference_world_size, + engine_world_size=engine_world_size, + ) server_url = f"{host}:{port}" set_ucx_env_defaults() self.nixl_agent = NixlAgent(make_agent_name("inference", global_rank)) diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 2d35f122ef..d39cce6b62 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -515,6 +515,8 @@ async def init_nccl_broadcast( timeout: int, inference_world_size: int | None = None, quantize_in_weight_transfer: bool = False, + *, + engine_world_sizes: list[int] | None = None, ) -> None: """Initialize NCCL broadcast on all inference servers. @@ -524,32 +526,46 @@ async def init_nccl_broadcast( """ logger = get_logger() + has_explicit_engine_world_sizes = engine_world_sizes is not None if inference_world_size is None: - inference_world_size = len(admin_clients) + if engine_world_sizes is not None: + inference_world_size = sum(engine_world_sizes) + else: + inference_world_size = len(admin_clients) logger.warning( f"inference_world_size not provided, defaulting to {inference_world_size} (one GPU per admin client)" ) - gpus_per_server = inference_world_size // len(admin_clients) + if engine_world_sizes is None: + if inference_world_size % len(admin_clients) != 0: + raise ValueError("inference_world_size must be divisible by the number of admin clients") + engine_world_sizes = [inference_world_size // len(admin_clients)] * len(admin_clients) + if len(engine_world_sizes) != len(admin_clients): + raise ValueError("one engine world size is required for each admin client") + rank_offsets = _rank_offsets(engine_world_sizes, inference_world_size) logger.info( f"Initializing NCCL broadcast: {len(admin_clients)} servers, " - f"inference_world_size={inference_world_size}, gpus_per_server={gpus_per_server}" + f"inference_world_size={inference_world_size}, engine_world_sizes={engine_world_sizes}" ) - async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> None: + async def _init_nccl_broadcast( + admin_client: AsyncClient, + rank_offset: int, + engine_world_size: int, + ) -> None: + payload = { + "host": host, + "port": port, + "rank_offset": rank_offset, + "inference_world_size": inference_world_size, + "timeout": timeout, + "quantize_in_weight_transfer": quantize_in_weight_transfer, + } + if has_explicit_engine_world_sizes: + payload["engine_world_size"] = engine_world_size try: - response = await admin_client.post( - "/init_broadcaster", - json={ - "host": host, - "port": port, - "rank_offset": rank_offset, - "inference_world_size": inference_world_size, - "timeout": timeout, - "quantize_in_weight_transfer": quantize_in_weight_transfer, - }, - ) + response = await admin_client.post("/init_broadcaster", json=payload) response.raise_for_status() except httpx.HTTPStatusError as e: if e.response.status_code == 404: @@ -558,8 +574,10 @@ async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> N await asyncio.gather( *[ - _init_nccl_broadcast(admin_client, client_num * gpus_per_server) - for client_num, admin_client in enumerate(admin_clients) + _init_nccl_broadcast(admin_client, rank_offset, engine_world_size) + for admin_client, rank_offset, engine_world_size in zip( + admin_clients, rank_offsets, engine_world_sizes, strict=True + ) ] ) @@ -571,31 +589,61 @@ async def init_nixl_broadcast( timeout: int, inference_world_size: int, session_id: str, + *, + engine_world_sizes: list[int] | None = None, ) -> None: """Configure every vLLM worker for NIXL + ModelExpress pulls.""" - workers_per_server = inference_world_size // len(admin_clients) - - async def initialize(admin_client: AsyncClient, rank_offset: int) -> None: + has_explicit_engine_world_sizes = engine_world_sizes is not None + if engine_world_sizes is None: + if inference_world_size % len(admin_clients) != 0: + raise ValueError("inference_world_size must be divisible by the number of admin clients") + engine_world_sizes = [inference_world_size // len(admin_clients)] * len(admin_clients) + if len(engine_world_sizes) != len(admin_clients): + raise ValueError("one engine world size is required for each admin client") + rank_offsets = _rank_offsets(engine_world_sizes, inference_world_size) + + async def initialize(admin_client: AsyncClient, rank_offset: int, engine_world_size: int) -> None: + payload = { + "host": host, + "port": port, + "rank_offset": rank_offset, + "inference_world_size": inference_world_size, + "timeout": timeout, + "quantize_in_weight_transfer": False, + "session_id": session_id, + } + if has_explicit_engine_world_sizes: + payload["engine_world_size"] = engine_world_size await _admin_post( admin_client, "/init_broadcaster", timeout_s=max(ADMIN_TIMEOUT_S, timeout), - json={ - "host": host, - "port": port, - "rank_offset": rank_offset, - "inference_world_size": inference_world_size, - "timeout": timeout, - "quantize_in_weight_transfer": False, - "session_id": session_id, - }, + json=payload, ) await asyncio.gather( - *[initialize(admin_client, index * workers_per_server) for index, admin_client in enumerate(admin_clients)] + *[ + initialize(admin_client, rank_offset, engine_world_size) + for admin_client, rank_offset, engine_world_size in zip( + admin_clients, rank_offsets, engine_world_sizes, strict=True + ) + ] ) +def _rank_offsets(engine_world_sizes: list[int], inference_world_size: int) -> list[int]: + if not engine_world_sizes or any(isinstance(size, bool) or size <= 0 for size in engine_world_sizes): + raise ValueError("engine world sizes must be positive integers") + if sum(engine_world_sizes) != inference_world_size: + raise ValueError("engine world sizes do not match inference_world_size") + offsets: list[int] = [] + offset = 0 + for world_size in engine_world_sizes: + offsets.append(offset) + offset += world_size + return offsets + + async def prefill_logprobs(openai: AsyncOpenAI, model: str, token_ids: list[int]) -> list[float]: """Prefill-score ``token_ids`` under ``model`` via ``/inference/v1/generate`` + ``prompt_logprobs`` (the prime-rl server-side extension in diff --git a/tests/unit/inference/test_nccl_rank.py b/tests/unit/inference/test_nccl_rank.py new file mode 100644 index 0000000000..dc2955814d --- /dev/null +++ b/tests/unit/inference/test_nccl_rank.py @@ -0,0 +1,90 @@ +import pytest + +from prime_rl.inference.vllm.ranks import global_inference_rank + + +def test_dense_aggregate_uses_engine_size_when_vllm_rewrites_dp_size(): + actual = { + global_inference_rank( + rank_offset=0, + data_parallel_index=dp_index, + data_parallel_size=1, + worker_rank=tp_rank, + tensor_parallel_size=2, + pipeline_parallel_size=1, + inference_world_size=8, + engine_world_size=8, + ) + for dp_index in range(4) + for tp_rank in range(2) + } + + assert actual == set(range(8)) + + +def test_already_global_moe_worker_ranks_are_not_double_counted(): + actual = { + global_inference_rank( + rank_offset=0, + data_parallel_index=dp_index, + data_parallel_size=4, + worker_rank=dp_index * 2 + tp_rank, + tensor_parallel_size=2, + pipeline_parallel_size=1, + inference_world_size=8, + ) + for dp_index in range(4) + for tp_rank in range(2) + } + + assert actual == set(range(8)) + + +def test_rank_offset_composes_with_pipeline_parallel_rank(): + actual = { + global_inference_rank( + rank_offset=8, + data_parallel_index=dp_index, + data_parallel_size=2, + worker_rank=model_parallel_rank, + tensor_parallel_size=2, + pipeline_parallel_size=2, + inference_world_size=16, + ) + for dp_index in range(2) + for model_parallel_rank in range(4) + } + + assert actual == set(range(8, 16)) + + +def test_global_inference_rank_rejects_out_of_bounds_rank(): + with pytest.raises(ValueError, match="outside inference world size"): + global_inference_rank( + rank_offset=2, + data_parallel_index=3, + data_parallel_size=4, + worker_rank=0, + tensor_parallel_size=2, + pipeline_parallel_size=1, + inference_world_size=8, + ) + + +def test_prefill_context_parallel_ranks_are_unique(): + actual = { + global_inference_rank( + rank_offset=0, + data_parallel_index=0, + data_parallel_size=1, + worker_rank=worker_rank, + tensor_parallel_size=1, + pipeline_parallel_size=1, + prefill_context_parallel_size=2, + inference_world_size=2, + engine_world_size=2, + ) + for worker_rank in range(2) + } + + assert actual == {0, 1} diff --git a/tests/unit/utils/test_external_engine_topology.py b/tests/unit/utils/test_external_engine_topology.py new file mode 100644 index 0000000000..591bd8c78a --- /dev/null +++ b/tests/unit/utils/test_external_engine_topology.py @@ -0,0 +1,44 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock + +from prime_rl.utils.client import _rank_offsets, init_nccl_broadcast + + +def test_rank_offsets_support_heterogeneous_engines(): + assert _rank_offsets([1, 3, 2], inference_world_size=6) == [0, 1, 4] + + +def test_nccl_broadcast_forwards_explicit_engine_sizes(): + clients = [AsyncMock(), AsyncMock()] + for client in clients: + response = MagicMock() + response.raise_for_status = MagicMock() + client.post.return_value = response + + asyncio.run( + init_nccl_broadcast( + clients, + host="127.0.0.1", + port=29519, + timeout=1200, + inference_world_size=4, + engine_world_sizes=[1, 3], + ) + ) + + assert [call.kwargs["json"]["rank_offset"] for client in clients for call in client.post.await_args_list] == [0, 1] + assert [call.kwargs["json"]["engine_world_size"] for client in clients for call in client.post.await_args_list] == [ + 1, + 3, + ] + + +def test_nccl_broadcast_preserves_legacy_payload(): + client = AsyncMock() + response = MagicMock() + response.raise_for_status = MagicMock() + client.post.return_value = response + + asyncio.run(init_nccl_broadcast([client], "127.0.0.1", 29519, 1200, 1)) + + assert "engine_world_size" not in client.post.await_args.kwargs["json"]