diff --git a/src/prime_rl/configs/orchestrator.py b/src/prime_rl/configs/orchestrator.py index f5424f377ca..922ae0d464a 100644 --- a/src/prime_rl/configs/orchestrator.py +++ b/src/prime_rl/configs/orchestrator.py @@ -11,6 +11,7 @@ HeartbeatConfig, LogConfig, PrimeMonitorConfig, + RendererConfig, TransportConfig, WandbWithExtrasConfig, ) @@ -870,6 +871,9 @@ class OrchestratorConfig(BaseConfig): # The tokenizer configuration tokenizer: TokenizerConfig = TokenizerConfig() + # The renderer configuration (only used when use_renderer=True) + renderer: RendererConfig = RendererConfig() + # The optimizer configuration (per-run LR for multi-run training) optim: OptimizerConfig = OptimizerConfig() @@ -1057,10 +1061,24 @@ class OrchestratorConfig(BaseConfig): use_token_client: Annotated[ bool, Field( - description="Whether to use the token-in-token-out (TITO) client for training across all environments. WARNING: Only use this if your environment has a linear history and the chat template has the extension property (i.e. no tokens are ever removed or inserted by the chat template)" + description="Whether to use the token-in-token-out (TITO) client for training across all environments. " + "WARNING: Only use this if your environment has a linear history and the chat template has the extension " + "property (i.e. no tokens are ever removed or inserted by the chat template). Mutually exclusive with " + "``use_renderer``." ), ] = True + use_renderer: Annotated[ + bool, + Field( + description="Whether to use the renderer client (client-side tokenization via the ``renderers`` package, " + "served by ``/v1/generate``). Mutually exclusive with ``use_token_client``. When True, the " + "``[orchestrator.renderer]`` block (name / tool_parser / reasoning_parser / pool_size) " + "applies; when False those fields must be left at their defaults. Not supported for VLMs — " + "VLMs must use the token client (TITO) so image preprocessing and chat templating stay server-side." + ), + ] = False + env_install_prerelease: Annotated[ bool, Field( @@ -1130,6 +1148,69 @@ def validate_sft_distill_mode(self): raise ValueError( "orchestrator.use_token_client must be false when orchestrator.teacher_rollout_model is configured." ) + if has_teacher and self.use_renderer: + raise ValueError( + "orchestrator.use_renderer must be false when orchestrator.teacher_rollout_model is configured " + "(external rollout uses MITO)." + ) + return self + + @model_validator(mode="after") + def validate_client_mode(self): + """The two client toggles select among three exclusive modes: + + - ``use_token_client=True`` + ``use_renderer=False`` → TITO (default) + - ``use_token_client=False`` + ``use_renderer=True`` → renderer + - ``use_token_client=False`` + ``use_renderer=False`` → MITO + + Both True is invalid: TITO and renderer are different wire protocols + (server-side templating vs client-side tokenization). + """ + if self.use_token_client and self.use_renderer: + raise ValueError( + "orchestrator.use_token_client and orchestrator.use_renderer are mutually exclusive. " + "Pick one: token client (TITO, server-side templating) or renderer client (client-side " + "tokenization)." + ) + return self + + @model_validator(mode="after") + def validate_renderer_vs_vlm(self): + """The renderer client takes plain message dicts and tokenizes + them client-side. VLMs need server-side image preprocessing and + chat templating, so they must use the token client (TITO) — fail + loudly when both are set.""" + if self.use_renderer and self.model.vlm is not None: + raise ValueError( + "orchestrator.use_renderer is not supported for VLMs. Use the token client " + "(``use_token_client=true``, the default) so image preprocessing and chat " + "templating stay on the inference server." + ) + return self + + @model_validator(mode="after") + def validate_renderer_args(self): + """``[orchestrator.renderer]`` knobs are only meaningful when + ``use_renderer=True``. Reject otherwise so callers don't silently + pass them and wonder why they're ignored.""" + if self.use_renderer: + return self + + renderer_args_set = [] + if self.renderer.name != "auto": + renderer_args_set.append(f"renderer.name={self.renderer.name!r}") + if self.renderer.tool_parser is not None: + renderer_args_set.append(f"renderer.tool_parser={self.renderer.tool_parser!r}") + if self.renderer.reasoning_parser is not None: + renderer_args_set.append(f"renderer.reasoning_parser={self.renderer.reasoning_parser!r}") + if self.renderer.pool_size is not None: + renderer_args_set.append(f"renderer.pool_size={self.renderer.pool_size!r}") + + if renderer_args_set: + raise ValueError( + "Renderer-specific args set without orchestrator.use_renderer=True: " + f"{', '.join(renderer_args_set)}. Either enable the renderer client or remove these knobs." + ) return self @model_validator(mode="after") diff --git a/src/prime_rl/configs/shared.py b/src/prime_rl/configs/shared.py index 029a78e9a01..d26c33d9a97 100644 --- a/src/prime_rl/configs/shared.py +++ b/src/prime_rl/configs/shared.py @@ -132,6 +132,61 @@ class BaseModelConfig(BaseConfig): ] = None +class RendererConfig(BaseConfig): + """Configures the client-side renderer (chat-template + response parsing). + + Only consumed when ``orchestrator.use_renderer = true``. The renderer + owns both directions: render messages → token ids on the client, and + parse model output tokens → structured ``content`` / ``reasoning_content`` + / ``tool_calls``. + """ + + name: Annotated[ + str, + Field( + description=( + "Renderer to use for chat template tokenization. " + "Options: 'auto' (detect from tokenizer), 'qwen3', 'qwen3_vl', " + "'qwen3.5', 'glm5', 'glm4.5', 'minimax-m2', 'deepseek_v3', " + "'kimi_k2', 'kimi_k25', 'nemotron3', 'gpt_oss', 'default'." + ), + ), + ] = "auto" + + tool_parser: Annotated[ + str | None, + Field( + description=( + "Name of a tool parser in renderers.parsers. Only consumed by " + "DefaultRenderer (model-specific renderers bake their own parsing " + "in). Options today: 'qwen3', 'qwen3.5', 'glm', 'deepseek_v3'." + ), + ), + ] = None + + reasoning_parser: Annotated[ + str | None, + Field( + description=( + "Name of a reasoning parser in renderers.parsers. Only consumed " + "by DefaultRenderer. Options today: 'think'." + ), + ), + ] = None + + pool_size: Annotated[ + int | None, + Field( + ge=1, + description=( + "Number of renderer slots shared across concurrent rollouts. " + "None keeps the verifiers default (1). Bump for long multi-turn " + "prompts where client-side jinja tokenization serializes." + ), + ), + ] = None + + class ElasticConfig(BaseConfig): """Configures elastic inference pool with DNS-based service discovery. diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index 46b5b5f566b..94eba0fd716 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -21,6 +21,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser from prime_rl.configs.inference import InferenceConfig +from prime_rl.inference.vllm.serving_generate import GenerateRequest from prime_rl.utils.logger import get_logger MODEL_TOOL_CALL_PARSER: dict[str, str] = { @@ -180,10 +181,62 @@ def models(request: Request) -> OpenAIServingModels: } +def generate_handler(request: Request): + return request.app.state.openai_serving_generate + + def chat_with_tokens(request: Request) -> OpenAIServingChatWithTokens | None: return request.app.state.openai_serving_chat_with_tokens +@router.post( + "/v1/generate", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.OK.value: {"content": {"application/json": {}}}, + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, +) +@with_cancellation +@load_aware_call +async def _generate(request: GenerateRequest, raw_request: Request): + handler = generate_handler(raw_request) + if handler is None: + return JSONResponse({"error": "generate endpoint not available"}, status_code=503) + result = await handler.generate(request, raw_request) + if isinstance(result, dict) and "error" in result: + return JSONResponse(result, status_code=500) + return JSONResponse(content=result.model_dump()) + + +@router.post( + "/v1/chat/completions/tokens", + dependencies=[Depends(validate_json_request)], + responses={ + HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, + }, +) +@with_cancellation +@load_aware_call +async def _chat_with_tokens(request: ChatCompletionRequestWithTokens, raw_request: Request): + handler = chat_with_tokens(raw_request) + if handler is None: + return base(raw_request).create_error_response(message="The model does not support Chat Completions API") + generator = await handler.create_chat_completion_with_tokens(request, raw_request) + if isinstance(generator, ErrorResponse): + return JSONResponse(content=generator.model_dump(), status_code=generator.error.code) + + elif isinstance(generator, ChatCompletionResponse): + return JSONResponse(content=generator.model_dump()) + + return StreamingResponse(content=generator, media_type="text/event-stream") + + @router.post("/pause") async def pause(request: Request): await engine_client(request).pause_generation(mode="keep", clear_cache=False) @@ -242,32 +295,6 @@ async def init_broadcaster(request: Request): return {"status": "ok"} -@router.post( - "/v1/chat/completions/tokens", - dependencies=[Depends(validate_json_request)], - responses={ - HTTPStatus.OK.value: {"content": {"text/event-stream": {}}}, - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, - HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - }, -) -@with_cancellation -@load_aware_call -async def _chat_with_tokens(request: ChatCompletionRequestWithTokens, raw_request: Request): - handler = chat_with_tokens(raw_request) - if handler is None: - return base(raw_request).create_error_response(message="The model does not support Chat Completions API") - generator = await handler.create_chat_completion_with_tokens(request, raw_request) - if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), status_code=generator.error.code) - - elif isinstance(generator, ChatCompletionResponse): - return JSONResponse(content=generator.model_dump()) - - return StreamingResponse(content=generator, media_type="text/event-stream") - - async def custom_init_app_state( engine_client: EngineClient, state: State, @@ -277,10 +304,16 @@ async def custom_init_app_state( """ Modifies init_app_state: 1. Call the original init_app_state to set up standard state. - 2. Replace the serving_chat with our OpenAIServingChatWithTokens wrapper. + 2. Replace ``serving_chat`` with our ``OpenAIServingChatWithTokens`` wrapper + so the ``/v1/chat/completions/tokens`` (TITO) endpoint can stream + token IDs alongside the rendered chat completion. + 3. Add ``/v1/generate`` endpoint for renderer-based token-level inference. """ await init_app_state(engine_client, state, args, supported_tasks) + state.reset_prefix_cache_after_update = getattr(args, "reset_prefix_cache_after_update", True) + + # TITO: server-side chat templating + token IDs. if "generate" in supported_tasks and state.openai_serving_chat is not None: original_chat = state.openai_serving_chat serving_chat = object.__new__(OpenAIServingChatWithTokens) @@ -290,6 +323,12 @@ async def custom_init_app_state( else: state.openai_serving_chat_with_tokens = None + # /v1/generate endpoint — tokens + optional images, no chat template + from prime_rl.inference.vllm.serving_generate import OpenAIServingGenerate + + chat_handler = state.openai_serving_chat if "generate" in supported_tasks else None + state.openai_serving_generate = OpenAIServingGenerate(engine_client, chat_handler=chat_handler) + import vllm.entrypoints.openai.api_server import vllm.v1.utils diff --git a/src/prime_rl/inference/vllm/serving_generate.py b/src/prime_rl/inference/vllm/serving_generate.py new file mode 100644 index 00000000000..594f603746d --- /dev/null +++ b/src/prime_rl/inference/vllm/serving_generate.py @@ -0,0 +1,225 @@ +"""/v1/generate endpoint — accepts pre-tokenized inputs. + +Text-only tokens in, tokens out. The Renderer does all tokenization client-side. +No Jinja rendering, no server-side chat template application. + +VLMs do not use this endpoint. The orchestrator routes VLMs to MITO +(/v1/chat/completions) where vLLM handles image preprocessing and chat +templating server-side. +""" + +from __future__ import annotations + +import asyncio +import base64 +from collections.abc import Mapping +from typing import Any +from uuid import uuid4 + +import numpy as np +from fastapi import Request +from pydantic import BaseModel +from vllm.engine.protocol import EngineClient +from vllm.inputs.engine import tokens_input +from vllm.logger import init_logger +from vllm.outputs import RequestOutput +from vllm.sampling_params import SamplingParams + +logger = init_logger(__name__) + + +# ── Request / Response schemas ─────────────────────────────────────── + + +class GenerateRequest(BaseModel): + model: str | None = None + prompt_token_ids: list[int] + + # When unset, fill from max_model_len - prompt_len at request time so we + # match /v1/chat/completions behavior. The previous 4096 hard default + # silently truncated long completions on 8k+ context runs (e.g. hendrycks + # reasoning rollouts capped at 4096 tokens, making rendered rollouts look + # shorter than main's for the same model). + max_tokens: int | None = None + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 + min_p: float = 0.0 + seed: int | None = None + n: int = 1 + stop_token_ids: list[int] | None = None + repetition_penalty: float = 1.0 + min_tokens: int = 0 + prompt_logprobs: bool = False + priority: int = 0 + # Prefix-cache invalidation salt. Must match main's + # /v1/chat/completions/tokens path: the orchestrator sets + # `extra_body["cache_salt"] = str(ckpt_step)` on every rollout + # request. vLLM's KV cache hashes include this salt, so when the + # step changes the cache misses and KV is recomputed with fresh + # weights. Without this, renderers path silently reuses stale KV + # from before the latest weight update and its logprobs drift from + # the trainer's forward pass (mismatch_kl grows 3x over training). + cache_salt: str | None = None + + +class GenerateChoiceResponse(BaseModel): + index: int + token_ids: list[int] + logprobs: list[float] + finish_reason: str | None = None + routed_experts: dict | None = None + + +class GenerateResponse(BaseModel): + id: str + model: str + prompt_token_ids: list[int] + choices: list[GenerateChoiceResponse] + usage: dict + prompt_logprobs: list[float | None] | None = None + + +# ── Handler ────────────────────────────────────────────────────────── + + +class OpenAIServingGenerate: + """Lightweight generate handler — tokens in, tokens out.""" + + def __init__(self, engine_client: EngineClient, chat_handler: Any | None = None): + self.engine_client = engine_client + self.chat_handler = chat_handler + + async def generate(self, request: GenerateRequest, raw_request: Request) -> GenerateResponse | dict: + # Pre-rendered TokensInput shape (type="token") — avoids vLLM's + # "raw prompt" deprecation that targets plain lists/strings. + engine_prompt = tokens_input(request.prompt_token_ids, cache_salt=request.cache_salt) + + # Match /v1/chat/completions: if the client didn't ask for a specific + # cap, let the model generate up to whatever room is left in context. + # vLLM v1 AsyncLLM exposes model_config directly (no async getter). + max_tokens = request.max_tokens + if max_tokens is None: + max_model_len = self.engine_client.model_config.max_model_len + max_tokens = max(1, max_model_len - len(request.prompt_token_ids)) + + sampling_params = SamplingParams( + max_tokens=max_tokens, + temperature=request.temperature, + top_p=request.top_p, + top_k=request.top_k, + min_p=request.min_p, + seed=request.seed, + n=request.n, + stop_token_ids=request.stop_token_ids or [], + repetition_penalty=request.repetition_penalty, + min_tokens=request.min_tokens, + logprobs=1, + prompt_logprobs=1 if request.prompt_logprobs else None, + skip_special_tokens=False, + ) + + request_id = f"gen-{uuid4().hex[:16]}" + routed_experts_map: dict[int, dict] = {} + final_output: RequestOutput | None = None + data_parallel_rank = None + lora_request = None + trace_headers = None + if self.chat_handler is not None: + data_parallel_rank = self.chat_handler._get_data_parallel_rank(raw_request) + trace_headers = await self.chat_handler._get_trace_headers(raw_request.headers) + lora_request = self.chat_handler._maybe_get_adapters(request) + + generator = self.engine_client.generate( + engine_prompt, + sampling_params, + request_id, + lora_request=lora_request, + trace_headers=trace_headers, + priority=request.priority, + data_parallel_rank=data_parallel_rank, + ) + + # Drain the generator without polling ``raw_request.is_disconnected`` + # per decode step. That poll is one ASGI ``receive()`` await per + # yielded token — at ~2k concurrent rollouts each producing dozens of + # tokens, it was the single largest per-step overhead on the /generate + # path (vLLM's own chat completions handler uses the CancelledError + # pattern for the same reason). Starlette cancels this coroutine on + # client disconnect, so the except branch still catches it. + try: + async for output in generator: + for comp_output in output.outputs: + if comp_output.routed_experts is not None: + routed_experts_map[comp_output.index] = _encode_routed_experts(comp_output.routed_experts) + final_output = output + except asyncio.CancelledError: + await self.engine_client.abort(request_id) + raise + + if final_output is None: + return {"error": "No output generated"} + + choices = [] + for output in final_output.outputs: + token_ids = list(output.token_ids) + logprobs_list: list[float] = [] + if output.logprobs: + for i, lp_dict in enumerate(output.logprobs): + if i < len(token_ids) and token_ids[i] in lp_dict: + logprobs_list.append(lp_dict[token_ids[i]].logprob) + else: + logprobs_list.append(0.0) + + choices.append( + GenerateChoiceResponse( + index=output.index, + token_ids=token_ids, + logprobs=logprobs_list, + finish_reason=output.finish_reason, + routed_experts=routed_experts_map.get(output.index), + ) + ) + + prompt_len = len(final_output.prompt_token_ids) + completion_len = sum(len(c.token_ids) for c in choices) + prompt_logprobs = _extract_prompt_logprobs(final_output.prompt_logprobs) + + return GenerateResponse( + id=request_id, + model=request.model or "", + prompt_token_ids=list(final_output.prompt_token_ids), + choices=choices, + usage={ + "prompt_tokens": prompt_len, + "completion_tokens": completion_len, + "total_tokens": prompt_len + completion_len, + }, + prompt_logprobs=prompt_logprobs, + ) + + +def _encode_routed_experts(arr: np.ndarray) -> dict: + return { + "data": base64.b85encode(arr.tobytes()).decode("ascii"), + "shape": list(arr.shape), + } + + +def _extract_prompt_logprobs( + prompt_logprobs: list[dict[int, Any] | None] | Mapping[int, Any] | None, +) -> list[float | None] | None: + if prompt_logprobs is None: + return None + if isinstance(prompt_logprobs, Mapping): + prompt_logprobs = [prompt_logprobs] + + extracted: list[float | None] = [] + for token_logprobs in prompt_logprobs: + if not token_logprobs: + extracted.append(None) + continue + selected = next(iter(token_logprobs.values())) + logprob = selected.logprob if hasattr(selected, "logprob") else selected.get("logprob") + extracted.append(float(logprob) if logprob is not None else None) + return extracted diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 39393492bbd..d4923146c3a 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -31,6 +31,7 @@ import pandas as pd import verifiers as vf +from renderers.base import create_renderer from transformers import AutoProcessor from prime_rl.configs.orchestrator import OrchestratorConfig @@ -118,14 +119,6 @@ async def orchestrate(config: OrchestratorConfig): # Setup rollout inference pool (handles both static and elastic modes) rollout_client_config, rollout_model_name, enable_policy_updates = setup_external_rollout_model(config, logger) - train_client_type = "openai_chat_completions_token" if config.use_token_client else "openai_chat_completions" - inference_pool = await setup_inference_pool( - rollout_client_config, - model_name=rollout_model_name, - train_client_type=train_client_type, - eval_client_type="openai_chat_completions", - ) - # Setup teacher inference pool if configured if config.teacher_model: logger.info( @@ -154,6 +147,14 @@ async def orchestrate(config: OrchestratorConfig): config.model.name, trust_remote_code=config.model.trust_remote_code, use_fast=True ) + renderer, inference_pool = await setup_rollout_inference_pool( + config=config, + rollout_client_config=rollout_client_config, + rollout_model_name=rollout_model_name, + tokenizer=tokenizer, + logger=logger, + ) + # Setup monitor (may register the run and set RUN_ID in the environment) logger.info(f"Initializing monitor (wandb={config.wandb}, prime_monitor={config.prime_monitor})") monitor = setup_monitor( @@ -480,33 +481,53 @@ async def orchestrate(config: OrchestratorConfig): # Convert rollouts to training samples parallel_preprocess_start = time.perf_counter() - # Pretokenize before VLM image cache build (which strips image data from messages) - for rollout in train_rollouts: - pretokenize_rollout_trajectory(rollout, tokenizer, processor=processor) + # Stage 1: pretokenize + (for VLM) build image cache concurrently. + # Pretokenize is a no-op when the renderer client already populated + # `tokens` on each trajectory step, but the fallback-tokenizer path + # and image-cache build are both CPU-heavy. Running them on threads + # and awaiting a single gather lets whichever finishes first free + # the event loop immediately and, with max_async_level >= 2, overlaps + # this whole stage with inference for the next batch. + async def _pretokenize_all() -> None: + await asyncio.gather( + *( + asyncio.to_thread( + pretokenize_rollout_trajectory, + rollout, + tokenizer, + processor=processor, + renderer=renderer, + ) + for rollout in train_rollouts + ) + ) - # VLM: build image cache in a thread so it doesn't block the event loop. - # This lets the scheduler continue servicing inflight rollout requests - # and — with max_async_level >= 2 — overlap with the next batch's inference. if is_vlm: - vlm_cache = await asyncio.to_thread(build_vlm_image_cache, train_rollouts, processor) mm_token_type_ids_mapping = {} if hasattr(processor, "image_token_id") and processor.image_token_id is not None: mm_token_type_ids_mapping[processor.image_token_id] = 1 if hasattr(processor, "video_token_id") and processor.video_token_id is not None: mm_token_type_ids_mapping[processor.video_token_id] = 2 - + _, vlm_cache = await asyncio.gather( + _pretokenize_all(), + asyncio.to_thread(build_vlm_image_cache, train_rollouts, processor), + ) logger.info( f"VLM timing: extract={vlm_cache.extract_time:.2f}s, preprocess={vlm_cache.preprocess_time:.2f}s " f"({vlm_cache.num_unique_images} unique images from {vlm_cache.num_unique_examples} examples)" ) else: + await _pretokenize_all() vlm_cache = None mm_token_type_ids_mapping = None # Process rollouts in parallel def process_rollout(rollout: vf.RolloutOutput, rollout_idx: int) -> list[TrainingSample] | None: return interleave_rollout( - rollout, vlm_cache=vlm_cache, cache_key=rollout_idx, mm_token_type_ids_mapping=mm_token_type_ids_mapping + rollout, + vlm_cache=vlm_cache, + cache_key=rollout_idx, + mm_token_type_ids_mapping=mm_token_type_ids_mapping, ) results = await asyncio.gather( @@ -864,5 +885,74 @@ def main(): asyncio.run(orchestrate(cli(OrchestratorConfig))) +async def setup_rollout_inference_pool( + *, + config: OrchestratorConfig, + rollout_client_config, + rollout_model_name: str, + tokenizer, + logger, +): + """Set up rollout inference. + + Routing policy is driven by ``config.use_token_client`` and + ``config.use_renderer`` (mutually exclusive — config-level validators + block both being True): + + - external teacher rollout → MITO (``openai_chat_completions``), + forced regardless of the toggles (config-level validator + rejects ``use_token_client`` / ``use_renderer`` in that case) + - ``use_renderer=True`` → renderer client (``/v1/generate``). + Not allowed for VLMs (validated at config time). + - ``use_token_client=True`` → TITO + (``openai_chat_completions_token``, ``/v1/chat/completions/tokens``). + Default. VLMs land here too. + - both False → MITO (``openai_chat_completions``). + """ + if config.teacher_rollout_model is not None: + logger.info("Using external rollout model (MITO) without renderer client") + inference_pool = await setup_inference_pool( + rollout_client_config, + model_name=rollout_model_name, + train_client_type="openai_chat_completions", + eval_client_type="openai_chat_completions", + ) + return None, inference_pool + + if config.use_renderer: + renderer = create_renderer( + tokenizer, + renderer=config.renderer.name, + tool_parser=config.renderer.tool_parser, + reasoning_parser=config.renderer.reasoning_parser, + ) + logger.info(f"Initialized {type(renderer).__name__} for {config.model.name}") + inference_pool = await setup_inference_pool( + rollout_client_config, + model_name=rollout_model_name, + train_client_type="renderer", + eval_client_type="openai_chat_completions", + renderer_name=config.renderer.name, + tool_parser=config.renderer.tool_parser, + reasoning_parser=config.renderer.reasoning_parser, + renderer_pool_size=config.renderer.pool_size, + ) + logger.info("Using direct renderer rollout client") + return renderer, inference_pool + + train_client_type = "openai_chat_completions_token" if config.use_token_client else "openai_chat_completions" + if config.use_token_client: + logger.info("Using token client (TITO) for rollouts — server-side templating, /v1/chat/completions/tokens") + else: + logger.info("Using MITO (openai_chat_completions) for rollouts") + inference_pool = await setup_inference_pool( + rollout_client_config, + model_name=rollout_model_name, + train_client_type=train_client_type, + eval_client_type="openai_chat_completions", + ) + return None, inference_pool + + if __name__ == "__main__": main() diff --git a/src/prime_rl/orchestrator/scheduler.py b/src/prime_rl/orchestrator/scheduler.py index cdd5f22bc2c..a8427b69f72 100644 --- a/src/prime_rl/orchestrator/scheduler.py +++ b/src/prime_rl/orchestrator/scheduler.py @@ -148,7 +148,10 @@ async def cancel_inflight_rollouts(self): @staticmethod def _client_identity(c: vf.ClientConfig) -> tuple[str, str | None]: - return (c.api_base_url, c.extra_headers.get("X-data-parallel-rank")) + return ( + c.api_base_url, + c.extra_headers.get("X-data-parallel-rank"), + ) async def _select_least_loaded_client(self) -> vf.ClientConfig: """Select the client with the fewest in-flight tasks. diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index b005b7278d9..3a45ee9ada5 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -199,12 +199,32 @@ def _get(tool: Any, key: str) -> Any: ] +def _tokenize_step_with_renderer( + step: vf.TrajectoryStep, + renderer, + tools: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + """Tokenize a trajectory step using a Renderer.""" + from renderers.base import build_trajectory_step + + prompt = _normalize_messages(step.get("prompt"), default_role="user") + completion = _normalize_messages(step.get("completion"), default_role="assistant") + prompt = _strip_message_content(_deserialize_tool_calls(prompt)) + completion = _strip_message_content(_deserialize_tool_calls(completion)) + return build_trajectory_step(renderer, prompt, completion, tools=tools) + + def pretokenize_rollout_trajectory( output: vf.RolloutOutput, tokenizer: PreTrainedTokenizer, processor=None, + renderer=None, ) -> bool: - """Populate missing step tokens from prompt/completion messages.""" + """Populate missing step tokens from prompt/completion messages. + + When a renderer is provided, uses it for tokenization (faster, deterministic). + Otherwise falls back to the tokenizer + apply_chat_template path. + """ logger = get_logger() tools = _convert_tools_to_oai_format(output.get("tool_defs", [])) @@ -212,17 +232,19 @@ def pretokenize_rollout_trajectory( if step["tokens"] is not None: continue - reconstructed = _tokenize_step_from_messages(step, tokenizer, tools=tools, processor=processor) - if reconstructed["prompt_prefix_len"] < reconstructed["original_prompt_len"]: - logger.debug( - f"Prompt tokenization was non-prefix for example {output['example_id']} step {step_idx}. " - f"Using longest common prefix length {reconstructed['prompt_prefix_len']} " - f"(original prompt had {reconstructed['original_prompt_len']} tokens)." - ) - - reconstructed.pop("prompt_prefix_len") - reconstructed.pop("original_prompt_len") - step["tokens"] = reconstructed + if renderer is not None: + step["tokens"] = _tokenize_step_with_renderer(step, renderer, tools=tools) + else: + reconstructed = _tokenize_step_from_messages(step, tokenizer, tools=tools, processor=processor) + if reconstructed["prompt_prefix_len"] < reconstructed["original_prompt_len"]: + logger.debug( + f"Prompt tokenization was non-prefix for example {output['example_id']} step {step_idx}. " + f"Using longest common prefix length {reconstructed['prompt_prefix_len']} " + f"(original prompt had {reconstructed['original_prompt_len']} tokens)." + ) + reconstructed.pop("prompt_prefix_len") + reconstructed.pop("original_prompt_len") + step["tokens"] = reconstructed return True @@ -384,7 +406,11 @@ def extend_sample(sample: TrainingSample, prefix_len: int, step_idx: int) -> Non new_prefix = tokens["prompt_ids"] + tokens["completion_ids"] active_samples.append((new_prefix, make_sample(tokens), step_idx)) - # Attach images once per sample using only the last merged step + # Attach images once per sample using only the last merged step. Prompt + # tokens already contain fully expanded <|image_pad|> placeholders because + # VLMs go through MITO (chat completions), which runs apply_chat_template + # server-side; the orchestrator re-tokenizes via the same processor in the + # fallback path so features and tokens stay 1:1. if vlm_cache is not None: key = output["example_id"] if cache_key is None else cache_key for _, sample, last_step_idx in active_samples: diff --git a/src/prime_rl/orchestrator/utils.py b/src/prime_rl/orchestrator/utils.py index 3b46ffce066..524f6077414 100644 --- a/src/prime_rl/orchestrator/utils.py +++ b/src/prime_rl/orchestrator/utils.py @@ -7,7 +7,6 @@ import pandas as pd import verifiers as vf -from openai.types.chat.chat_completion import ChatCompletion from rich.console import Console from rich.table import Table from verifiers.utils.client_utils import setup_openai_client @@ -85,28 +84,24 @@ async def compute_teacher_logprobs( samples: list[TrainingSample], ) -> list[list[float]]: """Compute teacher model logprobs for a batch of training samples via prefill.""" + from prime_rl.inference.vllm.serving_generate import GenerateResponse async def _compute_single(client_config: vf.ClientConfig, sample: TrainingSample) -> list[float]: client = setup_openai_client(client_config) response = await client.post( - "/chat/completions/tokens", + "/generate", + cast_to=GenerateResponse, body={ "model": model_name, - "messages": [{"role": "user", "content": ""}], - "tokens": sample.prompt_ids + sample.completion_ids, + "prompt_token_ids": sample.prompt_ids + sample.completion_ids, "max_tokens": 1, "temperature": 1.0, "top_p": 1.0, - "skip_special_tokens": False, "prompt_logprobs": True, }, - cast_to=ChatCompletion, ) - return [ - 0.0 if lp is None else float(next(iter(lp.values()))["logprob"]) - for lp in getattr(response, "prompt_logprobs", []) - ] + return [0.0 if lp is None else float(lp) for lp in response.prompt_logprobs or []] return await asyncio.gather(*[_compute_single(client, sample) for client, sample in zip(cycle(clients), samples)]) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index c1637ae3963..c1e913f19bf 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -1096,6 +1096,24 @@ def setup_model( return model +def _get_qwen3_vl_mm_token_type_ids(model: nn.Module, input_ids: Tensor) -> Tensor | None: + config = getattr(model, "config", None) + if getattr(config, "model_type", None) != "qwen3_vl": + return None + + mm_token_type_ids = torch.zeros_like(input_ids) + + image_token_id = getattr(config, "image_token_id", None) + if image_token_id is not None: + mm_token_type_ids = mm_token_type_ids.masked_fill(input_ids == image_token_id, 1) + + video_token_id = getattr(config, "video_token_id", None) + if video_token_id is not None: + mm_token_type_ids = mm_token_type_ids.masked_fill(input_ids == video_token_id, 2) + + return mm_token_type_ids + + @jaxtyped(typechecker=typechecker) def forward( model: nn.Module, @@ -1122,7 +1140,9 @@ def forward( assert image_grid_thw is not None, "pixel_values requires image_grid_thw for MRoPE computation" kwargs["pixel_values"] = pixel_values kwargs["image_grid_thw"] = image_grid_thw - kwargs["mm_token_type_ids"] = mm_token_type_ids + mm_token_type_ids = _get_qwen3_vl_mm_token_type_ids(model, input_ids) + if mm_token_type_ids is not None: + kwargs["mm_token_type_ids"] = mm_token_type_ids else: kwargs["position_ids"] = position_ids diff --git a/src/prime_rl/utils/client.py b/src/prime_rl/utils/client.py index 0953e54a898..21659dfc46b 100644 --- a/src/prime_rl/utils/client.py +++ b/src/prime_rl/utils/client.py @@ -64,8 +64,21 @@ def __init__( model_name: str, train_client_type: str = "openai_chat_completions_token", eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, ): - self._train_clients = setup_clients(client_config, client_type=train_client_type) + renderer_model_name = model_name if train_client_type == "renderer" else None + self._train_clients = setup_clients( + client_config, + client_type=train_client_type, + renderer_name=renderer_name, + renderer_model_name=renderer_model_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, + ) self._eval_clients = setup_clients(client_config, client_type=eval_client_type) self._admin_clients = setup_admin_clients(client_config) self._skip_model_check = client_config.skip_model_check @@ -112,6 +125,10 @@ async def setup_inference_pool( model_name: str, train_client_type: str = "openai_chat_completions_token", eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, ) -> InferencePool: """Create an inference pool from config (static or elastic).""" logger = get_logger() @@ -127,7 +144,14 @@ async def setup_inference_pool( from prime_rl.utils.elastic import ElasticInferencePool return await ElasticInferencePool.from_config( - client_config, model_name=model_name, train_client_type=train_client_type, eval_client_type=eval_client_type + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_name=renderer_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, ) logger.info( @@ -136,11 +160,26 @@ async def setup_inference_pool( f"api_key_var={client_config.api_key_var}, headers={client_config.headers})" ) return StaticInferencePool( - client_config, model_name=model_name, train_client_type=train_client_type, eval_client_type=eval_client_type + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_name=renderer_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, ) -def setup_clients(client_config: ClientConfig, client_type: str = "openai_chat_completions") -> list[vf.ClientConfig]: +def setup_clients( + client_config: ClientConfig, + client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + renderer_model_name: str | None = None, + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, +) -> list[vf.ClientConfig]: clients = [] client_idx = 0 for base_url in client_config.base_url: @@ -152,6 +191,11 @@ def setup_clients(client_config: ClientConfig, client_type: str = "openai_chat_c vf.ClientConfig( client_idx=client_idx, client_type=client_type, + renderer=renderer_name, + renderer_model_name=renderer_model_name, + renderer_pool_size=renderer_pool_size, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, api_base_url=base_url, api_key_var=client_config.api_key_var, timeout=client_config.timeout, diff --git a/src/prime_rl/utils/elastic.py b/src/prime_rl/utils/elastic.py index ba1fca53d5a..902f873903d 100644 --- a/src/prime_rl/utils/elastic.py +++ b/src/prime_rl/utils/elastic.py @@ -106,16 +106,25 @@ def __init__( model_name: str, train_client_type: str = "openai_chat_completions_token", eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, ): self.logger = get_logger() self.client_config = client_config self.model_name = model_name self.base_model_name = model_name # Keep original for health checks + self.renderer_model_name = model_name if train_client_type == "renderer" else None self.hostname = client_config.elastic.hostname self.port = client_config.elastic.port self.sync_interval = client_config.elastic.sync_interval self.train_client_type = train_client_type self.eval_client_type = eval_client_type + self.renderer_name = renderer_name + self.tool_parser = tool_parser + self.reasoning_parser = reasoning_parser + self.renderer_pool_size = renderer_pool_size self.router_url = client_config.router_url self._servers: dict[str, ServerState] = {} @@ -139,11 +148,22 @@ async def from_config( model_name: str, train_client_type: str = "openai_chat_completions_token", eval_client_type: str = "openai_chat_completions", + renderer_name: str = "auto", + tool_parser: str | None = None, + reasoning_parser: str | None = None, + renderer_pool_size: int | None = None, ) -> ElasticInferencePool: if client_config.elastic is None: raise ValueError("Elastic inference pool requires elastic config") pool = cls( - client_config, model_name=model_name, train_client_type=train_client_type, eval_client_type=eval_client_type + client_config, + model_name=model_name, + train_client_type=train_client_type, + eval_client_type=eval_client_type, + renderer_name=renderer_name, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + renderer_pool_size=renderer_pool_size, ) await pool.start() return pool @@ -185,7 +205,19 @@ def _rebuild_clients(self) -> None: dp_rank_count=self.client_config.dp_rank_count, extra_headers_from_state=self.client_config.extra_headers_from_state, ) - self._train_clients = setup_clients(url_config, client_type=self.train_client_type) if urls else [] + self._train_clients = ( + setup_clients( + url_config, + client_type=self.train_client_type, + renderer_name=self.renderer_name, + renderer_model_name=self.renderer_model_name, + tool_parser=self.tool_parser, + reasoning_parser=self.reasoning_parser, + renderer_pool_size=self.renderer_pool_size, + ) + if urls + else [] + ) self._eval_clients = setup_clients(url_config, client_type=self.eval_client_type) if urls else [] @property diff --git a/src/prime_rl/utils/messages.py b/src/prime_rl/utils/messages.py new file mode 100644 index 00000000000..13ebf6a6dcf --- /dev/null +++ b/src/prime_rl/utils/messages.py @@ -0,0 +1,62 @@ +"""Message utilities — normalize, deserialize tool calls, strip content.""" + +import json +from typing import Any + + +def normalize_messages(messages: Any, default_role: str) -> list[dict[str, Any]]: + if messages is None: + return [] + if isinstance(messages, str): + return [{"role": default_role, "content": messages}] + if isinstance(messages, dict): + return [dict(messages)] + if isinstance(messages, list): + normalized: list[dict[str, Any]] = [] + for message in messages: + if isinstance(message, str): + normalized.append({"role": default_role, "content": message}) + elif isinstance(message, dict): + normalized.append(dict(message)) + else: + raise TypeError(f"Unsupported message type: {type(message)}") + return normalized + raise TypeError(f"Unsupported messages container type: {type(messages)}") + + +def deserialize_tool_calls(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _deserialize_tool_call(tool_call: dict[str, Any]) -> dict[str, Any]: + function = tool_call.get("function", {}) + arguments = function.get("arguments") + if isinstance(arguments, str): + arguments = json.loads(arguments) + return { + **tool_call, + "function": {**function, "arguments": arguments}, + } + + deserialized_messages: list[dict[str, Any]] = [] + for message in messages: + if "tool_calls" not in message: + deserialized_messages.append(dict(message)) + continue + + tool_calls = message.get("tool_calls") or [] + deserialized_messages.append( + { + **message, + "tool_calls": [_deserialize_tool_call(tc) for tc in tool_calls], + } + ) + + return deserialized_messages + + +def strip_message_content(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + def _strip(message: dict[str, Any]) -> dict[str, Any]: + content = message.get("content") + if isinstance(content, str): + return {**message, "content": content.strip()} + return message + + return [_strip(message) for message in messages] diff --git a/tests/unit/inference/test_serving_generate.py b/tests/unit/inference/test_serving_generate.py new file mode 100644 index 00000000000..576d5352ea9 --- /dev/null +++ b/tests/unit/inference/test_serving_generate.py @@ -0,0 +1,93 @@ +import asyncio + +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput + +from prime_rl.inference.vllm.serving_generate import GenerateRequest, OpenAIServingGenerate + + +class _FakeChatHandler: + def __init__(self): + self.lora_request = object() + + def _get_data_parallel_rank(self, raw_request): + assert raw_request.headers["X-data-parallel-rank"] == "3" + return 3 + + async def _get_trace_headers(self, headers): + assert headers["traceparent"] == "trace" + return {"traceparent": headers["traceparent"]} + + def _maybe_get_adapters(self, request): + assert request.model == "test-model" + return self.lora_request + + +class _FakeEngineClient: + def __init__(self): + self.calls = [] + + async def generate(self, engine_prompt, sampling_params, request_id, **kwargs): + self.calls.append( + { + "engine_prompt": engine_prompt, + "sampling_params": sampling_params, + "request_id": request_id, + "kwargs": kwargs, + } + ) + yield RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=list(engine_prompt["prompt_token_ids"]), + prompt_logprobs=[None, {11: Logprob(-0.2)}, {12: Logprob(-0.4)}], + outputs=[ + CompletionOutput( + index=0, + text="done", + token_ids=[99], + cumulative_logprob=-0.6, + logprobs=[{99: Logprob(-0.6)}], + finish_reason="stop", + ) + ], + finished=True, + ) + + +class _FakeRawRequest: + def __init__(self): + self.headers = {"X-data-parallel-rank": "3", "traceparent": "trace"} + + async def is_disconnected(self): + return False + + +def test_generate_returns_prompt_logprobs_and_forwards_request_metadata(): + async def _run(): + engine_client = _FakeEngineClient() + handler = OpenAIServingGenerate(engine_client, chat_handler=_FakeChatHandler()) + request = GenerateRequest( + model="test-model", + prompt_token_ids=[10, 11, 12], + max_tokens=1, + prompt_logprobs=True, + priority=7, + ) + + response = await handler.generate(request, _FakeRawRequest()) + + assert response.prompt_token_ids == [10, 11, 12] + assert response.prompt_logprobs == [None, -0.2, -0.4] + assert response.choices[0].token_ids == [99] + assert response.choices[0].logprobs == [-0.6] + + call = engine_client.calls[0] + assert call["engine_prompt"] == {"type": "token", "prompt_token_ids": [10, 11, 12]} + assert call["sampling_params"].prompt_logprobs == 1 + assert call["kwargs"]["lora_request"] is handler.chat_handler.lora_request + assert call["kwargs"]["priority"] == 7 + assert call["kwargs"]["data_parallel_rank"] == 3 + assert call["kwargs"]["trace_headers"] == {"traceparent": "trace"} + + asyncio.run(_run()) diff --git a/tests/unit/orchestrator/test_orchestrator_setup.py b/tests/unit/orchestrator/test_orchestrator_setup.py new file mode 100644 index 00000000000..ff9bb5b79f9 --- /dev/null +++ b/tests/unit/orchestrator/test_orchestrator_setup.py @@ -0,0 +1,94 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +from prime_rl.orchestrator.orchestrator import setup_rollout_inference_pool + + +def test_setup_rollout_inference_pool_uses_plain_client_for_external_teacher_rollout(): + async def run() -> None: + tokenizer = object() + config = SimpleNamespace( + teacher_rollout_model=SimpleNamespace(), + model=SimpleNamespace(renderer="auto", name="student-model"), + ) + rollout_client_config = SimpleNamespace(base_url=["https://api.pinference.ai/api/v1"]) + logger = MagicMock() + inference_pool = object() + + with ( + patch( + "prime_rl.orchestrator.orchestrator.setup_inference_pool", new=AsyncMock(return_value=inference_pool) + ), + patch("prime_rl.orchestrator.orchestrator.create_renderer") as create_renderer_mock, + ): + renderer, returned_pool = await setup_rollout_inference_pool( + config=config, + rollout_client_config=rollout_client_config, + rollout_model_name="teacher-model", + tokenizer=tokenizer, + logger=logger, + ) + + assert renderer is None + assert returned_pool is inference_pool + create_renderer_mock.assert_not_called() + + asyncio.run(run()) + + +def test_setup_rollout_inference_pool_uses_direct_renderer_client_for_local_vllm(): + async def run() -> None: + tokenizer = object() + config = SimpleNamespace( + teacher_rollout_model=None, + use_renderer=True, + use_token_client=False, + model=SimpleNamespace(name="student-model"), + renderer=SimpleNamespace( + name="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + pool_size=None, + ), + ) + rollout_client_config = SimpleNamespace(base_url=["http://localhost:8000/v1"]) + logger = MagicMock() + renderer = object() + inference_pool = object() + + with ( + patch("prime_rl.orchestrator.orchestrator.create_renderer", return_value=renderer) as create_renderer_mock, + patch( + "prime_rl.orchestrator.orchestrator.setup_inference_pool", + new=AsyncMock(return_value=inference_pool), + ) as setup_pool_mock, + ): + returned_renderer, returned_pool = await setup_rollout_inference_pool( + config=config, + rollout_client_config=rollout_client_config, + rollout_model_name="student-model", + tokenizer=tokenizer, + logger=logger, + ) + + assert returned_renderer is renderer + assert returned_pool is inference_pool + create_renderer_mock.assert_called_once_with( + tokenizer, + renderer="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + ) + setup_pool_mock.assert_awaited_once_with( + rollout_client_config, + model_name="student-model", + train_client_type="renderer", + eval_client_type="openai_chat_completions", + renderer_name="qwen3_vl", + tool_parser=None, + reasoning_parser=None, + renderer_pool_size=None, + ) + + asyncio.run(run()) diff --git a/tests/unit/orchestrator/test_scheduler.py b/tests/unit/orchestrator/test_scheduler.py index af388ade6e3..9e73b5207da 100644 --- a/tests/unit/orchestrator/test_scheduler.py +++ b/tests/unit/orchestrator/test_scheduler.py @@ -3,6 +3,8 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch +import verifiers as vf + from prime_rl.orchestrator.scheduler import InflightRequest, Scheduler from prime_rl.utils.async_utils import safe_cancel @@ -159,3 +161,16 @@ async def update_weights(weight_dir, lora_name=None, step=0) -> None: assert scheduler.inflight_policy_update_task is None asyncio.run(run()) + + +def test_client_identity_distinguishes_base_url_and_dp_rank(): + client_a = vf.ClientConfig( + api_base_url="http://worker-a:8000/v1", + extra_headers={"X-data-parallel-rank": "0"}, + ) + client_b = vf.ClientConfig( + api_base_url="http://worker-a:8000/v1", + extra_headers={"X-data-parallel-rank": "1"}, + ) + + assert Scheduler._client_identity(client_a) != Scheduler._client_identity(client_b) diff --git a/tests/unit/orchestrator/test_teacher_logprobs.py b/tests/unit/orchestrator/test_teacher_logprobs.py new file mode 100644 index 00000000000..e126805a81a --- /dev/null +++ b/tests/unit/orchestrator/test_teacher_logprobs.py @@ -0,0 +1,62 @@ +import asyncio + +import verifiers as vf + +from prime_rl.inference.vllm.serving_generate import GenerateResponse +from prime_rl.orchestrator import utils as orchestrator_utils +from prime_rl.transport import TrainingSample + + +class _FakeOpenAIClient: + def __init__(self): + self.calls = [] + + async def post(self, path, *, cast_to, body): + self.calls.append({"path": path, "cast_to": cast_to, "body": body}) + return GenerateResponse( + id="gen-test", + model=body["model"], + prompt_token_ids=list(body["prompt_token_ids"]), + choices=[], + usage={"prompt_tokens": 3, "completion_tokens": 0, "total_tokens": 3}, + prompt_logprobs=[None, -0.7, -0.3], + ) + + +def test_compute_teacher_logprobs_uses_generate_endpoint(monkeypatch): + async def _run(): + fake_client = _FakeOpenAIClient() + monkeypatch.setattr(orchestrator_utils, "setup_openai_client", lambda _: fake_client) + + sample = TrainingSample( + prompt_ids=[1], + prompt_mask=[True], + completion_ids=[2, 3], + completion_mask=[True, True], + completion_logprobs=[-0.1, -0.2], + completion_temperatures=[1.0, 1.0], + ) + + result = await orchestrator_utils.compute_teacher_logprobs( + clients=[vf.ClientConfig()], + model_name="teacher-model", + samples=[sample], + ) + + assert result == [[0.0, -0.7, -0.3]] + assert fake_client.calls == [ + { + "path": "/generate", + "cast_to": GenerateResponse, + "body": { + "model": "teacher-model", + "prompt_token_ids": [1, 2, 3], + "max_tokens": 1, + "temperature": 1.0, + "top_p": 1.0, + "prompt_logprobs": True, + }, + } + ] + + asyncio.run(_run()) diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py new file mode 100644 index 00000000000..b62805bf3d4 --- /dev/null +++ b/tests/unit/train/test_model_forward.py @@ -0,0 +1,50 @@ +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from prime_rl.trainer.model import forward + + +class _CaptureModel(nn.Module): + def __init__(self, config: SimpleNamespace): + super().__init__() + self.config = config + self.kwargs = None + + def forward(self, **kwargs): + self.kwargs = kwargs + input_ids = kwargs["input_ids"] + return {"logits": torch.zeros(*input_ids.shape, 4)} + + +def test_forward_adds_qwen3_vl_mm_token_type_ids(): + image_token_id = 10 + video_token_id = 20 + model = _CaptureModel( + SimpleNamespace(model_type="qwen3_vl", image_token_id=image_token_id, video_token_id=video_token_id) + ) + input_ids = torch.tensor([[1, image_token_id, image_token_id, 2, video_token_id]]) + position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) + pixel_values = torch.ones(2, 3) + image_grid_thw = torch.tensor([[1, 1, 2]]) + + forward(model, input_ids, position_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw) + + assert model.kwargs is not None + assert "position_ids" not in model.kwargs + torch.testing.assert_close(model.kwargs["pixel_values"], pixel_values) + torch.testing.assert_close(model.kwargs["image_grid_thw"], image_grid_thw) + torch.testing.assert_close(model.kwargs["mm_token_type_ids"], torch.tensor([[0, 1, 1, 0, 2]])) + + +def test_forward_skips_mm_token_type_ids_for_other_vlm_models(): + model = _CaptureModel(SimpleNamespace(model_type="qwen3_5_moe", image_token_id=10)) + input_ids = torch.tensor([[1, 10, 10, 2]]) + position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) + + forward(model, input_ids, position_ids, pixel_values=torch.ones(2, 3), image_grid_thw=torch.tensor([[1, 1, 2]])) + + assert model.kwargs is not None + assert "position_ids" not in model.kwargs + assert "mm_token_type_ids" not in model.kwargs diff --git a/tests/unit/utils/test_client.py b/tests/unit/utils/test_client.py index 1c25dddf258..6b48790ef36 100644 --- a/tests/unit/utils/test_client.py +++ b/tests/unit/utils/test_client.py @@ -3,9 +3,10 @@ from unittest.mock import AsyncMock, MagicMock import httpx -import pytest +import verifiers as vf -from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter +from prime_rl.configs.shared import ClientConfig +from prime_rl.utils.client import _is_retryable_lora_error, load_lora_adapter, setup_clients def test_is_retryable_lora_error_returns_true_for_404(): @@ -48,39 +49,68 @@ def test_load_lora_adapter_succeeds_on_first_attempt(): ) -def test_load_lora_adapter_retries_on_404_then_succeeds(): - mock_client = AsyncMock() - - error_response = MagicMock() - error_response.status_code = 404 - success_response = MagicMock() - success_response.raise_for_status = MagicMock() - - call_count = 0 +def test_setup_clients_assigns_renderer_and_dp_rank_headers(): + client_config = ClientConfig( + base_url=["http://worker-a:8000/v1"], + api_key_var="PRIME_API_KEY", + headers={"X-Test": "test"}, + dp_rank_count=2, + extra_headers_from_state={"X-Session-ID": "session_id"}, + ) - async def mock_post(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise httpx.HTTPStatusError("Not found", request=MagicMock(), response=error_response) - return success_response + clients = setup_clients( + client_config, + client_type="renderer", + renderer_name="qwen3_vl", + ) - mock_client.post = mock_post + assert [client.client_type for client in clients] == ["renderer", "renderer"] + assert [client.renderer for client in clients] == ["qwen3_vl", "qwen3_vl"] + assert [client.renderer_model_name for client in clients] == [None, None] + assert [client.api_base_url for client in clients] == ["http://worker-a:8000/v1"] * 2 + assert [client.extra_headers["X-data-parallel-rank"] for client in clients] == ["0", "1"] + assert clients[0].extra_headers["X-Test"] == "test" + assert clients[0].extra_headers_from_state == {"X-Session-ID": "session_id"} - asyncio.run(load_lora_adapter([mock_client], "test-lora", Path("/test/path"))) - assert call_count == 2 +def test_setup_clients_assigns_renderer_model_name(): + client_config = ClientConfig( + base_url=["http://worker-a:8000/v1"], + api_key_var="PRIME_API_KEY", + ) + clients = setup_clients( + client_config, + client_type="renderer", + renderer_name="qwen3_vl", + renderer_model_name="Qwen/Qwen3-VL-4B-Instruct", + ) -def test_load_lora_adapter_raises_non_retryable_error_immediately(): - mock_client = AsyncMock() + assert clients[0].renderer_model_name == "Qwen/Qwen3-VL-4B-Instruct" - error_response = MagicMock() - error_response.status_code = 400 - mock_client.post.side_effect = httpx.HTTPStatusError("Bad request", request=MagicMock(), response=error_response) - with pytest.raises(httpx.HTTPStatusError) as exc_info: - asyncio.run(load_lora_adapter([mock_client], "test-lora", Path("/test/path"))) +def test_setup_clients_preserves_chat_client_defaults(): + client_config = ClientConfig( + base_url=["http://worker-a:8000/v1"], + api_key_var="PRIME_API_KEY", + ) - assert exc_info.value.response.status_code == 400 - assert mock_client.post.call_count == 1 + clients = setup_clients(client_config) + + assert clients == [ + vf.ClientConfig( + client_idx=0, + client_type="openai_chat_completions", + renderer="auto", + renderer_model_name=None, + api_key_var="PRIME_API_KEY", + api_base_url="http://worker-a:8000/v1", + timeout=client_config.timeout, + connect_timeout=client_config.connect_timeout, + max_connections=8192, + max_keepalive_connections=8192, + max_retries=10, + extra_headers={}, + extra_headers_from_state={}, + ) + ] diff --git a/tests/unit/utils/test_elastic.py b/tests/unit/utils/test_elastic.py index 0596638dc4d..a2cd40d1d17 100644 --- a/tests/unit/utils/test_elastic.py +++ b/tests/unit/utils/test_elastic.py @@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx +import verifiers as vf from prime_rl.utils.elastic import ( AdapterState, @@ -401,3 +402,49 @@ def test_get_loaded_adapter_handles_step_dash_format(): assert result is not None assert result.step == 99 + + +def test_elastic_clients_preserve_renderer_model_name_when_model_name_updates(): + with patch("prime_rl.utils.elastic.get_logger"): + client_config = MagicMock() + client_config.elastic.hostname = "test.hostname" + client_config.elastic.port = 8000 + client_config.elastic.sync_interval = 5.0 + client_config.router_url = None + client_config.timeout = 1200 + client_config.connect_timeout = 30.0 + client_config.api_key_var = "PRIME_API_KEY" + client_config.headers = {} + client_config.extra_headers_from_state = {} + client_config.dp_rank_count = 1 + + pool = ElasticInferencePool( + client_config=client_config, + model_name="Qwen/Qwen3-VL-4B-Instruct", + train_client_type="renderer", + renderer_name="qwen3_vl", + ) + pool._servers = { + "10.0.0.1": MagicMock(status="ready"), + } + + pool.update_model_name("r8-smoke") + clients = pool.train_clients + + assert clients == [ + vf.ClientConfig( + client_idx=0, + client_type="renderer", + renderer="qwen3_vl", + renderer_model_name="Qwen/Qwen3-VL-4B-Instruct", + api_key_var="PRIME_API_KEY", + api_base_url="http://10.0.0.1:8000/v1", + timeout=1200, + connect_timeout=30.0, + max_connections=8192, + max_keepalive_connections=8192, + max_retries=10, + extra_headers={}, + extra_headers_from_state={}, + ) + ]