From 321ff13d3d721e124f336a4dea59db9640f9e944 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:01:52 +0200 Subject: [PATCH 01/13] Serve training records through native Chat Completions --- configs/advanced/minimax-m2.5/swe.toml | 2 +- configs/advanced/nemotron-3-super/swe.toml | 2 +- configs/ci/integration/alphabet_sort.toml | 2 +- .../integration/reverse-text-lora/resume.toml | 2 +- .../integration/reverse-text-lora/start.toml | 2 +- .../integration/reverse-text-moe/start.toml | 2 +- .../reverse-text-rl-opd/start.toml | 2 +- .../reverse-text-rl-sft/start.toml | 8 +- .../ci/integration/reverse-text/resume.toml | 2 +- .../ci/integration/reverse-text/start.toml | 2 +- configs/ci/nightly-fft/hendrycks-sanity.toml | 2 +- configs/ci/nightly-fft/reverse-text.toml | 2 +- configs/ci/nightly-fft/wordle.toml | 2 +- configs/debug/algo/echo.toml | 2 +- configs/debug/algo/grpo.toml | 2 +- configs/debug/algo/max_rl.toml | 2 +- configs/debug/algo/mixed_grpo_opd.toml | 2 +- configs/debug/algo/opd.toml | 4 +- configs/debug/algo/opd_lora.toml | 4 +- configs/debug/algo/rae.toml | 2 +- configs/debug/algo/self_distill.toml | 2 +- configs/debug/algo/sft_distill.toml | 10 +- configs/debug/algo/sft_distill_lora.toml | 10 +- deps/verifiers | 2 +- docs/algorithms.md | 49 ++--- examples/advanced/glm-4.5-air/search.toml | 2 +- examples/advanced/glm-4.5-air/swe-budget.toml | 2 +- examples/advanced/glm-4.5-air/swe.toml | 2 +- examples/advanced/glm-4.5-air/terminal.toml | 2 +- examples/advanced/glm-5.3/swe.toml | 2 +- examples/advanced/intellect-3.1/rl.toml | 2 +- examples/basic/hendrycks-sanity/rl.toml | 2 +- examples/basic/reverse-text/rl.toml | 2 +- examples/basic/wordle/rl.toml | 2 +- examples/extra/dynamo/rl.toml | 2 +- .../src/prime_rl/configs/inference.py | 7 +- .../src/prime_rl/configs/orchestrator.py | 41 ---- .../src/prime_rl/configs/shared.py | 4 +- src/prime_rl/inference/patches.py | 39 ---- src/prime_rl/inference/vllm/server.py | 20 +- src/prime_rl/inference/vllm/serving_chat.py | 188 ++++++++++++++++++ src/prime_rl/orchestrator/algo/echo.py | 30 ++- src/prime_rl/orchestrator/clients.py | 58 +----- src/prime_rl/orchestrator/dispatcher.py | 6 +- src/prime_rl/orchestrator/envs.py | 10 +- src/prime_rl/orchestrator/orchestrator.py | 4 - src/prime_rl/orchestrator/train_sink.py | 4 +- src/prime_rl/orchestrator/trajectories.py | 12 +- tests/unit/orchestrator/test_algorithms.py | 35 ++++ tests/unit/orchestrator/test_clients.py | 37 +--- tests/unit/orchestrator/test_qwen3_vl_e2e.py | 176 ++++++++-------- tests/unit/test_configs.py | 117 +++-------- 52 files changed, 463 insertions(+), 468 deletions(-) create mode 100644 src/prime_rl/inference/vllm/serving_chat.py diff --git a/configs/advanced/minimax-m2.5/swe.toml b/configs/advanced/minimax-m2.5/swe.toml index f8f43b6a6a9..191666e822b 100644 --- a/configs/advanced/minimax-m2.5/swe.toml +++ b/configs/advanced/minimax-m2.5/swe.toml @@ -64,5 +64,5 @@ env.agent.runtime.labels = ["minimax-swe", "swe-bench-verified"] [inference.vllm] tensor_parallel_size = 8 -[orchestrator.renderer] +[inference.renderer] name = "minimax-m2" diff --git a/configs/advanced/nemotron-3-super/swe.toml b/configs/advanced/nemotron-3-super/swe.toml index ea1f57dd7f2..692987ca8ce 100644 --- a/configs/advanced/nemotron-3-super/swe.toml +++ b/configs/advanced/nemotron-3-super/swe.toml @@ -72,7 +72,7 @@ num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 -[orchestrator.renderer] +[inference.renderer] name = "nemotron-3" truncate_history_thinking = false diff --git a/configs/ci/integration/alphabet_sort.toml b/configs/ci/integration/alphabet_sort.toml index 29b7a295c3e..74aa649b71e 100644 --- a/configs/ci/integration/alphabet_sort.toml +++ b/configs/ci/integration/alphabet_sort.toml @@ -43,6 +43,6 @@ env.agent.runtime.type = "subprocess" # Mirror of Qwen/Qwen3-0.6B; PI's template patch always re-emits prior # blocks. Match that with renderer thinking_retention. -[orchestrator.renderer] +[inference.renderer] name = "auto" thinking_retention = "all" diff --git a/configs/ci/integration/reverse-text-lora/resume.toml b/configs/ci/integration/reverse-text-lora/resume.toml index 57cf9d52695..a0963a0bdff 100644 --- a/configs/ci/integration/reverse-text-lora/resume.toml +++ b/configs/ci/integration/reverse-text-lora/resume.toml @@ -41,5 +41,5 @@ env.agent.runtime.type = "subprocess" [inference.vllm] enable_lora = true gpu_memory_utilization = 0.7 -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/configs/ci/integration/reverse-text-lora/start.toml b/configs/ci/integration/reverse-text-lora/start.toml index 66f3598075d..57ab76bff93 100644 --- a/configs/ci/integration/reverse-text-lora/start.toml +++ b/configs/ci/integration/reverse-text-lora/start.toml @@ -38,5 +38,5 @@ env.agent.runtime.type = "subprocess" [inference.vllm] enable_lora = true gpu_memory_utilization = 0.7 -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/configs/ci/integration/reverse-text-moe/start.toml b/configs/ci/integration/reverse-text-moe/start.toml index f371592ee1f..7d71fe39e1c 100644 --- a/configs/ci/integration/reverse-text-moe/start.toml +++ b/configs/ci/integration/reverse-text-moe/start.toml @@ -38,5 +38,5 @@ env.agent.runtime.type = "subprocess" [inference.vllm] max_model_len = 2048 gpu_memory_utilization = 0.7 -[orchestrator.renderer] +[inference.renderer] name = "default" diff --git a/configs/ci/integration/reverse-text-rl-opd/start.toml b/configs/ci/integration/reverse-text-rl-opd/start.toml index 0a769d8f0b5..2b77babd674 100644 --- a/configs/ci/integration/reverse-text-rl-opd/start.toml +++ b/configs/ci/integration/reverse-text-rl-opd/start.toml @@ -36,7 +36,7 @@ type = "opd" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" base_url = "http://localhost:8001/v1" -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text-rl-sft/start.toml b/configs/ci/integration/reverse-text-rl-sft/start.toml index 200ac877714..0099fed2455 100644 --- a/configs/ci/integration/reverse-text-rl-sft/start.toml +++ b/configs/ci/integration/reverse-text-rl-sft/start.toml @@ -37,12 +37,8 @@ type = "sft" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" base_url = "http://localhost:8001/v1" -# The student's renderer tokenizes the teacher's messages into the student's -# token space (backfill). Use `prime-qwen3` — it matches this PrimeIntellect -# Qwen3 model's chat template token-for-token (single-turn + system), so the CE -# target stays faithful, without the empty `` block the stock -# `qwen3` renderer injects (which this model's template never emits). -[orchestrator.renderer] +# Match the policy chat template. The frozen generation server also owns its renderer. +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/ci/integration/reverse-text/resume.toml b/configs/ci/integration/reverse-text/resume.toml index fed5fefbf96..6433ab2442e 100644 --- a/configs/ci/integration/reverse-text/resume.toml +++ b/configs/ci/integration/reverse-text/resume.toml @@ -40,5 +40,5 @@ env.agent.runtime.type = "subprocess" # PrimeIntellect Qwen3 fine-tune - select its renderer explicitly (the -SFT # suffix misses the MODEL_RENDERER_MAP exact match). -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/configs/ci/integration/reverse-text/start.toml b/configs/ci/integration/reverse-text/start.toml index aa73930f8f1..1eb95ee4aee 100644 --- a/configs/ci/integration/reverse-text/start.toml +++ b/configs/ci/integration/reverse-text/start.toml @@ -37,5 +37,5 @@ env.agent.runtime.type = "subprocess" # PrimeIntellect Qwen3 fine-tune - select its renderer explicitly (the -SFT # suffix misses the MODEL_RENDERER_MAP exact match). -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/configs/ci/nightly-fft/hendrycks-sanity.toml b/configs/ci/nightly-fft/hendrycks-sanity.toml index ee22eb035b7..d1c2a0ee409 100644 --- a/configs/ci/nightly-fft/hendrycks-sanity.toml +++ b/configs/ci/nightly-fft/hendrycks-sanity.toml @@ -50,6 +50,6 @@ attn = "flash_attention_2" [inference.vllm] max_model_len = 8192 data_parallel_size = 4 -[orchestrator.renderer] +[inference.renderer] name = "default" reasoning_parser = "think" diff --git a/configs/ci/nightly-fft/reverse-text.toml b/configs/ci/nightly-fft/reverse-text.toml index 474d3278347..8c93107f119 100644 --- a/configs/ci/nightly-fft/reverse-text.toml +++ b/configs/ci/nightly-fft/reverse-text.toml @@ -36,5 +36,5 @@ max_completion_tokens = 128 [inference] -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/configs/ci/nightly-fft/wordle.toml b/configs/ci/nightly-fft/wordle.toml index e9e7517a48e..2b2b6c25bdd 100644 --- a/configs/ci/nightly-fft/wordle.toml +++ b/configs/ci/nightly-fft/wordle.toml @@ -36,5 +36,5 @@ max_completion_tokens = 1024 [inference.vllm] data_parallel_size = 6 -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/configs/debug/algo/echo.toml b/configs/debug/algo/echo.toml index 354640189d7..62fc8c5ddeb 100644 --- a/configs/debug/algo/echo.toml +++ b/configs/debug/algo/echo.toml @@ -39,7 +39,7 @@ env.agent.runtime.type = "subprocess" max_completion_tokens = 512 # Fine-tune inherits the PrimeIntellect Qwen3 template byte-for-byte. -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" [trainer.optim] diff --git a/configs/debug/algo/grpo.toml b/configs/debug/algo/grpo.toml index a4cefbcbc43..df7d3d20b10 100644 --- a/configs/debug/algo/grpo.toml +++ b/configs/debug/algo/grpo.toml @@ -17,7 +17,7 @@ group_size = 16 [orchestrator.train.algo] type = "grpo" -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/max_rl.toml b/configs/debug/algo/max_rl.toml index 39f4ac030c6..da6264e511a 100644 --- a/configs/debug/algo/max_rl.toml +++ b/configs/debug/algo/max_rl.toml @@ -17,7 +17,7 @@ group_size = 16 [orchestrator.train.algo] type = "max_rl" -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/mixed_grpo_opd.toml b/configs/debug/algo/mixed_grpo_opd.toml index ec199ca24c7..7ca7f5943e7 100644 --- a/configs/debug/algo/mixed_grpo_opd.toml +++ b/configs/debug/algo/mixed_grpo_opd.toml @@ -27,7 +27,7 @@ group_size = 16 [orchestrator.train.algo] type = "grpo" -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/opd.toml b/configs/debug/algo/opd.toml index 4201b2a7877..d51dfff168b 100644 --- a/configs/debug/algo/opd.toml +++ b/configs/debug/algo/opd.toml @@ -28,8 +28,8 @@ type = "opd" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" base_url = "http://localhost:8001/v1" -# Teacher rolls out over plain chat-completions (no tokens); the renderer backfills them. -[orchestrator.renderer] +# The inference server returns exact tokens with the policy chat template. +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/opd_lora.toml b/configs/debug/algo/opd_lora.toml index 5df1cf23d38..e28e4df12b8 100644 --- a/configs/debug/algo/opd_lora.toml +++ b/configs/debug/algo/opd_lora.toml @@ -28,8 +28,8 @@ type = "opd" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" base_url = "http://localhost:8001/v1" -# Teacher rolls out over plain chat-completions (no tokens); the renderer backfills them. -[orchestrator.renderer] +# The inference server returns exact tokens with the policy chat template. +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/rae.toml b/configs/debug/algo/rae.toml index 887b542c1af..d1f1ec21b11 100644 --- a/configs/debug/algo/rae.toml +++ b/configs/debug/algo/rae.toml @@ -17,7 +17,7 @@ group_size = 1 [orchestrator.train.algo] type = "rae" -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/self_distill.toml b/configs/debug/algo/self_distill.toml index 006f95889e9..070aed544e7 100644 --- a/configs/debug/algo/self_distill.toml +++ b/configs/debug/algo/self_distill.toml @@ -25,7 +25,7 @@ group_size = 1 type = "opsd" demo_key = "answer" -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/sft_distill.toml b/configs/debug/algo/sft_distill.toml index 92ceb124237..deff62491e6 100644 --- a/configs/debug/algo/sft_distill.toml +++ b/configs/debug/algo/sft_distill.toml @@ -1,6 +1,6 @@ # Start the frozen reference server first (on a separate GPU): # CUDA_VISIBLE_DEVICES=1 uv run inference \ -# --vllm.model PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --vllm.model PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL --renderer.name prime-qwen3 \ # --server.port 8001 --vllm.gpu-memory-utilization 0.5 --vllm.enforce-eager # Then: # uv run rl @ configs/debug/algo/sft_distill.toml @@ -29,12 +29,8 @@ type = "sft" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" base_url = "http://localhost:8001/v1" -# The student's renderer tokenizes the teacher's messages into the student's -# token space (backfill). Use `prime-qwen3` — it matches this PrimeIntellect -# Qwen3 model's chat template token-for-token (single-turn + system), keeping the -# CE target faithful, without the empty `` block the stock `qwen3` -# renderer injects (which this model's template never emits). -[orchestrator.renderer] +# Match the policy chat template. The frozen generation server also owns its renderer. +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/configs/debug/algo/sft_distill_lora.toml b/configs/debug/algo/sft_distill_lora.toml index 30080bab985..d875484daee 100644 --- a/configs/debug/algo/sft_distill_lora.toml +++ b/configs/debug/algo/sft_distill_lora.toml @@ -1,6 +1,6 @@ # Start the frozen reference server first (on a separate GPU): # CUDA_VISIBLE_DEVICES=1 uv run inference \ -# --vllm.model PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL \ +# --vllm.model PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL --renderer.name prime-qwen3 \ # --server.port 8001 --vllm.gpu-memory-utilization 0.5 --vllm.enforce-eager # Then: # uv run rl @ configs/debug/algo/sft_distill_lora.toml @@ -29,12 +29,8 @@ type = "sft" name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL" base_url = "http://localhost:8001/v1" -# The student's renderer tokenizes the teacher's messages into the student's -# token space (backfill). Use `prime-qwen3` — it matches this PrimeIntellect -# Qwen3 model's chat template token-for-token (single-turn + system), keeping the -# CE target faithful, without the empty `` block the stock `qwen3` -# renderer injects (which this model's template never emits). -[orchestrator.renderer] +# Match the policy chat template. The frozen generation server also owns its renderer. +[inference.renderer] name = "prime-qwen3" [orchestrator.train.sampling] diff --git a/deps/verifiers b/deps/verifiers index e888a719adf..00c1534182f 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit e888a719adf328c87f271fba39fbc812ca99677b +Subproject commit 00c1534182f5eab7beab1f5230662b960c45e6d3 diff --git a/docs/algorithms.md b/docs/algorithms.md index 0923d6e67eb..15ca5c7350c 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -490,7 +490,7 @@ reject_max = 0.05 ## Multi-Turn Trajectories -For multi-turn rollouts (tool use, browser environments, long conversations), `prime-rl` records each LLM request/response as an independent **trajectory step** and merges them at training time using best-effort interleaving — with [renderers](#renderers) as the mechanism that keeps the merge safe by construction. +For multi-turn rollouts (tool use, browser environments, long conversations), `prime-rl` records each LLM request/response as an independent **trajectory step** and merges them at training time using best-effort interleaving. Verifiers checks the server-returned token IDs before sharing a prefix. ### Extension Property @@ -518,47 +518,32 @@ steps 4–5: extension holds → merged into Sample 2 result: 2 training samples instead of 5 ``` -The orchestrator enforces an **exact prefix invariant**: the prompt at turn $t$ must be the exact concatenation of prior messages exactly as the LLM originally generated them. If turn 2's prompt is `U1, A1', U2` while `A1' ≠ A1`, the orchestrator can't safely merge — either choice produces logprob drift between trainer and inference. Starting a fresh sample is the only correct behavior, so that's what happens. +Verifiers enforces an **exact prefix invariant** when merging records: the prompt at turn $t$ must be the exact concatenation of prior messages exactly as the LLM originally generated them. If turn 2's prompt is `U1, A1', U2` while `A1' ≠ A1`, the orchestrator can't safely merge — either choice produces logprob drift between trainer and inference. Starting a fresh sample is the only correct behavior, so that's what happens. ### Renderers -Best-effort interleaving works because the renderer guarantees the exact-prefix invariant *by construction* — it never re-renders prior turns, so it can't lose tokens to chat-template normalization, BPE retokenization drift, or thinking stripping. A renderer turns a model's chat template into a Python object that can: +Evaluation and training use the same OpenAI-compatible client. Training requests ask the +inference server for exact prompt and completion token IDs, sampled-token logprobs, and +training metadata. Each call is a valid training record in its own right. Verifiers merges +records only when their token prefixes match exactly; a re-rendered history starts a new +sample with its context masked out. -- `render_ids(messages)` — tokenize messages to ids the inference engine accepts. -- `parse_response(completion_ids)` — recover structured `(content, reasoning_content, tool_calls)` from sampled ids. -- `bridge_to_next_turn(prev_prompt_ids, prev_completion_ids, new_messages)` — extend the previous turn's tokens verbatim with the new environment turn, instead of re-rendering history. - -When `bridge_to_next_turn` succeeds, the trainer sees the exact token stream the sampler produced; when it can't be proven safe (e.g. the renderer is `DefaultRenderer` and the template's stop sequence is unknown), it returns `None` and the orchestrator falls back to a full re-render — which triggers the new-sample fallback above. - -A common source of breakage in the absence of a hand-coded renderer is models like Qwen3 whose chat templates strip past `` blocks across user turns: - -```python -from transformers import AutoTokenizer -tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") -messages = [ - {"role": "user", "content": "U1"}, - {"role": "assistant", "content": "R1A1"}, - {"role": "user", "content": "U2"}, -] -tok.apply_chat_template(messages[:1], tokenize=False) -# <|im_start|>user -# U1<|im_end|> - -tok.apply_chat_template(messages, tokenize=False) -# <|im_start|>user\nU1<|im_end|>\n<|im_start|>assistant\nA1<|im_end|>\n<|im_start|>user\nU2<|im_end|> -# (the R1 from turn 2 is gone) -``` - -Hand-coded renderers ship for `qwen3`, `qwen3-vl`, `qwen3.5`, `glm-5`, `glm-4.5`, `minimax-m2`, `deepseek-v3`, `kimi-k2`, `kimi-k2.5`, `nemotron-3`, `gpt-oss`; anything else falls back to `DefaultRenderer` (a generic `apply_chat_template` wrapper). Pick one via: +The inference server uses the `renderers` library to attribute prompt tokens to messages +and message bodies. ECHO uses this attribution to select observations, including when +tokenization changes across turns. Multimodal records keep inline image URLs and exact +placeholder ranges; the trainer builds image features from those URLs. Native vLLM handles +sampling, and the inference renderer parses responses and tool calls. Sampling masks and +expert routing travel alongside the tokens for replay in the trainer. ```toml -[orchestrator.renderer] -name = "auto" # detect from tokenizer; pass an explicit name for fine-tunes +[inference.renderer] +name = "auto" # detect from tokenizer; use an explicit name for fine-tunes ``` A renderer outside the `renderers` package loads as a custom renderer: `name = "custom"` with `import_path = "my_module.Class"` (or `path/to/file.py:Class`), see [Custom renderers](training.md#dataset-format). -For the full design rationale (failure modes ruled out, empirical token-identity comparison against `apply_chat_template`, when to write a hand-coded renderer), see [the renderers writeup on the Prime Intellect blog](https://www.primeintellect.ai/blog/renderers) — the canonical reference. +RL and distillation losses require sampling logprobs. CE-only frozen generation can omit +logprobs, but its endpoint must still return exact tokens and training metadata. ### Discontinuous Trajectories diff --git a/examples/advanced/glm-4.5-air/search.toml b/examples/advanced/glm-4.5-air/search.toml index 7f8d119b582..693fff0f4f2 100644 --- a/examples/advanced/glm-4.5-air/search.toml +++ b/examples/advanced/glm-4.5-air/search.toml @@ -64,7 +64,7 @@ group_size = 16 initial_inflight = 256 max_inflight = 512 -[orchestrator.renderer] +[inference.renderer] name = "glm-4.5" thinking_retention = "all" diff --git a/examples/advanced/glm-4.5-air/swe-budget.toml b/examples/advanced/glm-4.5-air/swe-budget.toml index 59c62f56b9d..13e5376e3b0 100644 --- a/examples/advanced/glm-4.5-air/swe-budget.toml +++ b/examples/advanced/glm-4.5-air/swe-budget.toml @@ -82,7 +82,7 @@ num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 -[orchestrator.renderer] +[inference.renderer] name = "glm-4.5" thinking_retention = "all" diff --git a/examples/advanced/glm-4.5-air/swe.toml b/examples/advanced/glm-4.5-air/swe.toml index 00bb54c93b9..5c2997c0e2b 100644 --- a/examples/advanced/glm-4.5-air/swe.toml +++ b/examples/advanced/glm-4.5-air/swe.toml @@ -82,7 +82,7 @@ num_output_tokens_weight = 0.0 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 -[orchestrator.renderer] +[inference.renderer] name = "glm-4.5" thinking_retention = "all" diff --git a/examples/advanced/glm-4.5-air/terminal.toml b/examples/advanced/glm-4.5-air/terminal.toml index 97c447560da..53c9db7e845 100644 --- a/examples/advanced/glm-4.5-air/terminal.toml +++ b/examples/advanced/glm-4.5-air/terminal.toml @@ -65,7 +65,7 @@ group_size = 16 initial_inflight = 256 max_inflight = 512 -[orchestrator.renderer] +[inference.renderer] name = "glm-4.5" thinking_retention = "all" diff --git a/examples/advanced/glm-5.3/swe.toml b/examples/advanced/glm-5.3/swe.toml index 3d209e189c2..2004438a01e 100644 --- a/examples/advanced/glm-5.3/swe.toml +++ b/examples/advanced/glm-5.3/swe.toml @@ -123,7 +123,7 @@ max_inflight = 512 [orchestrator.model] name = "zai-org/GLM-5.3-BF16" -[orchestrator.renderer] +[inference.renderer] name = "glm-5.3" clear_thinking = false diff --git a/examples/advanced/intellect-3.1/rl.toml b/examples/advanced/intellect-3.1/rl.toml index 21ee145a6dc..cd8f2b9c64c 100644 --- a/examples/advanced/intellect-3.1/rl.toml +++ b/examples/advanced/intellect-3.1/rl.toml @@ -109,5 +109,5 @@ env.agent.runtime.type = "subprocess" [inference.vllm] tensor_parallel_size = 8 -[orchestrator.renderer] +[inference.renderer] name = "default" diff --git a/examples/basic/hendrycks-sanity/rl.toml b/examples/basic/hendrycks-sanity/rl.toml index 174beaf96c0..ec3e4b95f39 100644 --- a/examples/basic/hendrycks-sanity/rl.toml +++ b/examples/basic/hendrycks-sanity/rl.toml @@ -41,6 +41,6 @@ seq_len = 16384 [inference.vllm] max_model_len = 8192 -[orchestrator.renderer] +[inference.renderer] name = "default" reasoning_parser = "think" diff --git a/examples/basic/reverse-text/rl.toml b/examples/basic/reverse-text/rl.toml index bda8e9bb63b..ac165ea019b 100644 --- a/examples/basic/reverse-text/rl.toml +++ b/examples/basic/reverse-text/rl.toml @@ -45,5 +45,5 @@ lr = 3e-6 # PrimeIntellect Qwen3 fine-tune - select its renderer explicitly (the -SFT # suffix misses the MODEL_RENDERER_MAP exact match). -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/examples/basic/wordle/rl.toml b/examples/basic/wordle/rl.toml index c8af6b0db4c..07e0e4da11d 100644 --- a/examples/basic/wordle/rl.toml +++ b/examples/basic/wordle/rl.toml @@ -33,5 +33,5 @@ max_completion_tokens = 1024 [inference.vllm] data_parallel_size = 6 -[orchestrator.renderer] +[inference.renderer] name = "prime-qwen3" diff --git a/examples/extra/dynamo/rl.toml b/examples/extra/dynamo/rl.toml index c8199e1175b..86a660183e3 100644 --- a/examples/extra/dynamo/rl.toml +++ b/examples/extra/dynamo/rl.toml @@ -37,6 +37,6 @@ base_url = "http://127.0.0.1:8000/v1" [orchestrator.model.client.dynamo] enabled = true -[orchestrator.renderer] +[inference.renderer] name = "qwen3" enable_thinking = false diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index da2869df955..dd04dd7ebfd 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -5,6 +5,7 @@ from pydantic import ConfigDict, Field, model_validator from pydantic_config import BaseConfig +from renderers import AutoRendererConfig, RendererConfig from prime_rl.configs.shared import EnvVars, LogConfig, SlurmConfig from prime_rl.utils.config import default_output_dir, find_package_resource @@ -410,6 +411,9 @@ def num_nodes(self) -> int: class InferenceConfig(BaseConfig): + renderer: RendererConfig = AutoRendererConfig() + """Server-side prompt rendering and message/image attribution for training records.""" + server: ServerConfig = ServerConfig() router: RouterConfig | None = Field(default_factory=VllmRouterConfig) @@ -439,7 +443,7 @@ class InferenceConfig(BaseConfig): """Auto-set for disaggregated P/D: emit the NIXL transfer connector. Persisted into the per-node config (which drops ``deployment``) so the connector is still built per worker. Not meant to be set by hand.""" enable_return_sampling_mask: bool = False - """Return per-token sampling masks (``sampling_mask``) on ``/inference/v1/generate`` responses via vLLM's native ``--return-sampling-mask`` (>= 0.28). The ``rl`` entrypoint enables this field for truncated policy sampling. Standalone inference must set it explicitly because no orchestrator sampling config is available. The field persists into per-node configs and selects the V2 model runner before vLLM starts. Capture is engine-wide: vLLM rejects requests with ``temperature <= 0`` or without ``top_k > 0`` while it is on.""" + """Return per-token sampling masks (``sampling_mask``) on training Chat Completions and ``/inference/v1/generate`` responses via vLLM's native ``--return-sampling-mask`` (>= 0.28). The ``rl`` entrypoint enables this field for truncated policy sampling. Standalone inference must set it explicitly because no orchestrator sampling config is available. The field persists into per-node configs and selects the V2 model runner before vLLM starts. Capture is engine-wide: vLLM rejects requests with ``temperature <= 0`` or without ``top_k > 0`` while it is on.""" enable_fp32_lm_head: bool = True """Run the lm_head projection in fp32 via a native bf16×bf16 → fp32 GEMM (``torch.mm`` with ``out_dtype=torch.float32``). Stabilizes logprob precision under FP8/bf16 inference, matching SGLang's ``--enable-fp32-lm-head``. Implemented natively by vLLM's LogitsProcessor, which reads ``head_dtype`` off the HF config, so this flag injects ``hf_overrides = {"head_dtype": "float32"}``.""" @@ -573,6 +577,7 @@ def to_namespace(self) -> Namespace: host=self.server.host, port=self.server.port, liveness_timeout_seconds=self.server.liveness_timeout_seconds, + prime_renderer=self.renderer, ) extra_fields = self.vllm.model_extra or {} diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index 485f75bc53e..3a64ca7cb9c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -5,7 +5,6 @@ import verifiers.v1 as vf from pydantic import AliasChoices, BaseModel, Field, SerializeAsAny, TypeAdapter, ValidationError, model_validator from pydantic.fields import FieldInfo -from renderers import AutoRendererConfig, RendererConfig from verifiers.v1.configs.agent import agent_config_fields from prime_rl.configs.algorithm import ( @@ -477,12 +476,6 @@ class OrchestratorConfig(BaseConfig): tokenizer: TokenizerConfig = TokenizerConfig() - renderer: RendererConfig = AutoRendererConfig() - """Typed renderer config (``renderers.RendererConfig`` discriminated union), required — - training is renderer-only. Defaults to ``"auto"``, which resolves from - ``tokenizer.name_or_path`` via ``MODEL_RENDERER_MAP``. RL/OPD roll out through the renderer - client; SFT uses it to backfill tokens for its chat-completions teacher.""" - eval: RLOnlineEvalConfig | None = None """Evaluation configuration.""" @@ -641,40 +634,6 @@ def any_policy_sourced(self) -> bool: """True when at least one train env samples rollouts from the live policy.""" return any(env.algo.sampling.source == "policy" for env in self.train.source) - @model_validator(mode="after") - def validate_renderer_auto_resolves(self): - """Reject the silent DefaultRenderer fallback at config time. - - When ``renderer.name='auto'`` and the model isn't in - ``MODEL_RENDERER_MAP``, ``create_renderer`` would fall back to - ``DefaultRenderer``. That fallback doesn't fix the - position-dependent chat-template bug the renderer client exists - to solve, and rejects envs that pass tools (the rollout dies - with "RendererPool does not support tools") unless - ``DefaultRendererConfig.tool_parser`` is configured. Surface at - config time so ``--dry-run`` reports the error. - """ - if self.renderer.name != "auto": - return self - from renderers.base import MODEL_RENDERER_MAP - - model_id = self.tokenizer.name or self.model.name - if model_id in MODEL_RENDERER_MAP: - return self - raise ValueError( - f"orchestrator.renderer.name='auto' but " - f"{model_id!r} is not in renderers.base.MODEL_RENDERER_MAP, so it " - f"would silently fall back to DefaultRenderer. Pick one: " - f"(a) [orchestrator.renderer] name='default' — for fine-tunes / " - f"vendored mirrors with custom chat templates (DefaultRenderer " - f"calls apply_chat_template); set tool_parser= if the env " - f"uses tools. " - f"(b) [orchestrator.renderer] name= — " - f"if {model_id!r} is template-identical to a mapped family " - f"(and ideally also add it upstream to " - f"renderers.base.MODEL_RENDERER_MAP)." - ) - @model_validator(mode="after") def resolve_batching(self): group_sizes = [source.group_size for source in self.train.source] or [self.train.group_size] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 1ce168817a9..cb96c11d65b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -4,7 +4,7 @@ from typing import Annotated, Literal, TypeAlias from pydantic import AfterValidator, Field, model_validator -from verifiers.v1.configs.client import BaseClientConfig +from verifiers.v1.configs.client import ClientConfig as VerifiersClientConfig from prime_rl.utils.config import BaseConfig @@ -222,7 +222,7 @@ class DynamoConfig(BaseConfig): """Dynamo frontend URL used to discover inference workers for RL control.""" -class ClientConfig(BaseClientConfig): +class ClientConfig(VerifiersClientConfig): wait_for_ready_timeout: int = 3600 """Seconds to wait at startup for the inference pool to become ready.""" diff --git a/src/prime_rl/inference/patches.py b/src/prime_rl/inference/patches.py index bdc4e8c8ad8..dfaf1651578 100644 --- a/src/prime_rl/inference/patches.py +++ b/src/prime_rl/inference/patches.py @@ -26,7 +26,6 @@ def apply_shared_vllm_patches(): monkey_patch_online_fp8_parameter_cast() monkey_patch_deepseek_v4_allowed_layer_types() monkey_patch_tokenize_params_validation() - monkey_patch_strip_routed_experts_from_chat() monkey_patch_dp_coordinator_startup_timeout() monkey_patch_minimax_m2_for_lora() monkey_patch_packed_sampling_masks() @@ -480,44 +479,6 @@ def _is_pd_decode_request(request) -> bool: return bool(params.get("remote_engine_id")) and not params.get("do_remote_decode") -def monkey_patch_strip_routed_experts_from_chat(): - """Drop routed_experts from chat-completions responses. - - routed_experts are only consumed via the serialized ``/generate`` - (serving_tokens) path used for router-replay training, which encodes them as a - ``{data, shape, start}`` object the PD router can merge. The stock - chat-completions path instead encodes them as a base64 ``np.save`` *string*, - which the PD router rejects ("prefill routed_experts must be an object with - base64 data and shape") and fails every eval rollout (evals go through chat - completions). ``enable_return_routed_experts`` is a server-wide model-config - flag with no per-request toggle, so strip the field on the chat path here. - """ - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat - from vllm.logger import init_logger - - logger = init_logger(__name__) - - if getattr(OpenAIServingChat.chat_completion_full_generator, "_prime_rl_strips_routed_experts", False): - return - - _original = OpenAIServingChat.chat_completion_full_generator - - async def _strip(result_generator): - async for res in result_generator: - for output in res.outputs: - output.routed_experts = None - yield res - - async def _patched(self, request, result_generator, *args, **kwargs): - return await _original(self, request, _strip(result_generator), *args, **kwargs) - - _patched._prime_rl_strips_routed_experts = True - OpenAIServingChat.chat_completion_full_generator = _patched - logger.info( - "Stripped routed_experts from chat-completions responses (PD router merges only the /generate object form)." - ) - - def _patch_qwen35_moe_lora_format(): """Force Qwen3.5-MoE onto vLLM's 2D per-expert LoRA format. diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index 7fea3bb4e6c..5fa0b931409 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -1,6 +1,7 @@ import asyncio import gc from argparse import Namespace +from threading import Lock import uvloop from fastapi import APIRouter, Request @@ -127,14 +128,7 @@ async def custom_init_app_state( args: Namespace, supported_tasks: tuple, ): - """ - Modifies init_app_state: - 1. Call the original init_app_state to set up standard state, including - vLLM 0.20's ``serving_tokens`` for ``/inference/v1/generate``. - 2. Replace ``serving_tokens`` with ``PrimeRlServingTokens`` so DP-rank - routing and ``routed_experts`` export survive the migration off the - legacy ``/v1/generate`` endpoint. - """ + """Attach training metadata to the native vLLM serving interfaces.""" await init_app_state(engine_client, state, args, supported_tasks) state.liveness_timeout_seconds = args.liveness_timeout_seconds @@ -150,6 +144,16 @@ async def custom_init_app_state( prime_serving.__dict__.update(upstream.__dict__) state.serving_tokens = prime_serving + if "generate" in supported_tasks and state.openai_serving_chat is not None: + from prime_rl.inference.vllm.serving_chat import PrimeRlServingChat + + upstream = state.openai_serving_chat + chat = object.__new__(PrimeRlServingChat) + chat.__dict__.update(upstream.__dict__) + chat.training_renderer_config = args.prime_renderer + chat.training_renderer_lock = Lock() + state.openai_serving_chat = chat + # Under RL load the API server holds millions of live objects, and CPython's # default thresholds rerun full collections over them every few seconds. # Move everything alive after startup out of GC and collect less often. diff --git a/src/prime_rl/inference/vllm/serving_chat.py b/src/prime_rl/inference/vllm/serving_chat.py new file mode 100644 index 00000000000..9ee867a9fc6 --- /dev/null +++ b/src/prime_rl/inference/vllm/serving_chat.py @@ -0,0 +1,188 @@ +"""Native Chat Completions with exact tokens and optional training attribution.""" + +from __future__ import annotations + +import asyncio +import json +from threading import Lock + +from renderers import Renderer, RendererConfig, create_renderer, load_tokenizer +from renderers.base import ToolCallParseStatus, _resolve_renderer_config, is_multimodal +from renderers.client import _build_mm_features +from vllm.entrypoints.generate.base.protocol import FunctionCall, ToolCall +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, +) +from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat +from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item +from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls +from vllm.inputs import mm_input, tokens_input +from vllm.multimodal.inputs import MultiModalKwargsItems, PlaceholderRange + +from prime_rl.inference.vllm.routed_experts import serialize_routed_experts + + +class TrainingChoice(ChatCompletionResponseChoice): + routed_experts: dict | None = None # type: ignore[assignment] + sampling_mask: list[list[int]] | None = None + + +class TrainingCompletion(ChatCompletionResponse): + choices: list[TrainingChoice] + training_metadata: dict | None = None + + +class PrimeRlServingChat(OpenAIServingChat): + training_renderer_config: RendererConfig + training_renderer_lock: Lock + training_renderer: Renderer | None = None + training_tokenizer = None + training_template_kwargs: dict | None = None + + async def render_chat_request(self, request: ChatCompletionRequest): + if not getattr(request, "return_training_metadata", False): + return await super().render_chat_request(request) + if request.stream or not request.return_token_ids or not request.include_reasoning: + raise ValueError("training metadata requires stream=false, return_token_ids=true, include_reasoning=true") + if request.continue_final_message or request.chat_template is not None: + raise ValueError("training records require complete messages and the server's configured renderer") + error = await self._check_model(request) + if error is not None: + return error + self._preflight(request.n or 1) + # Tokenizer/processor calls are synchronous. One worker-owned renderer is + # sufficient; no rollout state or previously sampled tokens are kept here. + return await asyncio.to_thread(self._render_training_prompt, request) + + def _render_training_prompt(self, request: ChatCompletionRequest): + # Hold the lock in the worker thread even if the HTTP request is cancelled. + with self.training_renderer_lock: + if self.training_tokenizer is None: + self.training_tokenizer = load_tokenizer(self.model_config.tokenizer) + self.training_renderer_config = _resolve_renderer_config( + self.training_tokenizer, self.training_renderer_config + ) + kwargs = {**self.default_chat_template_kwargs, **(request.chat_template_kwargs or {})} + template_fields = self.training_renderer_config.template_field_names() + if request.reasoning_effort is not None: + if "reasoning_effort" in template_fields: + kwargs.setdefault("reasoning_effort", request.reasoning_effort) + elif "enable_thinking" in template_fields: + kwargs.setdefault("enable_thinking", request.reasoning_effort != "none") + if self.training_renderer is None or kwargs != self.training_template_kwargs: + self.training_renderer = create_renderer( + self.training_tokenizer, self.training_renderer_config, chat_template_kwargs=kwargs + ) + self.training_template_kwargs = kwargs + renderer = self.training_renderer + messages = [dict(message) for message in request.messages] + for message in messages: + if "reasoning" in message: + message["reasoning_content"] = message.pop("reasoning") + tools = [tool.model_dump(exclude_none=True) for tool in request.tools] if request.tools else None + rendered = renderer.render(messages, tools=tools, add_generation_prompt=request.add_generation_prompt) + if len(rendered.token_ids) >= self.model_config.max_model_len: + raise ValueError( + f"Prompt length ({len(rendered.token_ids)}) exceeds maximum context length " + f"({self.model_config.max_model_len})" + ) + request.stop_token_ids = list( + dict.fromkeys([*(request.stop_token_ids or []), *renderer.get_stop_token_ids()]) + ) + mmd = rendered.multi_modal_data + # Decode with this request's renderer even if another request replaces the cache. + request._prime_response_renderer = (renderer, tools) + request._prime_training_metadata = { + "message_spans": rendered.message_token_spans(), + "is_content": rendered.is_content, + "mm_token_type_id_map": renderer.mm_token_type_id_map if is_multimodal(renderer) else None, + "mm_placeholders": sorted( + (span.offset, span.length) for ranges in mmd.mm_placeholders.values() for span in ranges + ) + if mmd is not None + else None, + } + engine_input = tokens_input(rendered.token_ids, cache_salt=request.cache_salt) + if mmd is not None: + features = _build_mm_features(renderer, mmd) + engine_input = mm_input( + prompt_token_ids=rendered.token_ids, + mm_kwargs=MultiModalKwargsItems( + { + modality: [decode_mm_kwargs_item(item) for item in items] + for modality, items in features["kwargs_data"].items() + } + ), + mm_hashes=features["mm_hashes"], + mm_placeholders={ + modality: [PlaceholderRange(**span) for span in spans] + for modality, spans in features["mm_placeholders"].items() + }, + cache_salt=request.cache_salt, + ) + return messages, [engine_input] + + async def chat_completion_full_generator(self, request, result_generator, *args, **kwargs): + experts = {} + masks = {} + + async def capture(): + async for result in result_generator: + for output in result.outputs: + routing = serialize_routed_experts(output.routed_experts, request.routed_experts_prompt_start) + if routing is not None: + experts[output.index] = routing + # Native vLLM emits .npy here; the PD router consumes compact arrays. + output.routed_experts = None + if output.sampling_mask is not None: + masks[output.index] = output.sampling_mask.token_ids + yield result + + response_renderer = getattr(request, "_prime_response_renderer", None) + if response_renderer is not None: + kwargs["parser"] = None + response = await super().chat_completion_full_generator(request, capture(), *args, **kwargs) + if not isinstance(response, ChatCompletionResponse): + return response + if response_renderer is not None: + renderer, tools = response_renderer + for choice in response.choices: + parsed = await asyncio.to_thread(renderer.parse_response, choice.token_ids, tools=tools) + choice.message.content = parsed.content or None + choice.message.reasoning = parsed.reasoning_content + choice.message.tool_calls = [ + ToolCall( + **({"id": call.id} if call.id else {}), + function=FunctionCall( + name=call.name, + arguments=call.arguments + if isinstance(call.arguments, str) + else json.dumps(call.arguments or {}), + ), + ) + for call in parsed.tool_calls + if call.status == ToolCallParseStatus.OK + ] + if ( + choice.message.tool_calls + and choice.finish_reason == "stop" + and request.tool_choice in (None, "auto", "required") + ): + choice.finish_reason = "tool_calls" + return TrainingCompletion( + **response.model_dump(exclude={"choices"}), + choices=[ + maybe_filter_parallel_tool_calls( + TrainingChoice( + **choice.model_dump(exclude={"routed_experts"}), + routed_experts=experts.get(choice.index), + sampling_mask=masks.get(choice.index), + ), + request, + ) + for choice in response.choices + ], + training_metadata=getattr(request, "_prime_training_metadata", None), + ) diff --git a/src/prime_rl/orchestrator/algo/echo.py b/src/prime_rl/orchestrator/algo/echo.py index f1fb31600e6..f9c390ff852 100644 --- a/src/prime_rl/orchestrator/algo/echo.py +++ b/src/prime_rl/orchestrator/algo/echo.py @@ -3,6 +3,8 @@ from functools import partial from typing import TYPE_CHECKING, Callable +from verifiers.v1.graph import _node_key + from prime_rl.configs.algorithm import EchoAlgoConfig from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm from prime_rl.orchestrator.trajectories import iter_trainable_branches @@ -40,8 +42,7 @@ async def score_episode(self, episode: vf.Episode) -> None: def _weight_observations(self, trace: vf.Trace) -> None: """Write graph-native ``ce`` weights over the env-provided - observation tokens of later turns. Provenance is structural under v1: - within a branch, the non-sampled nodes that follow the first model + observation tokens of later turns. Non-sampled messages after a model response (tool output, user feedback) are the env-provided observations — each such node's tokens get its message role's weight, narrowed by the optional user filter. The initial prompt (before the @@ -57,6 +58,19 @@ def _weight_observations(self, trace: vf.Trace) -> None: fall back to weighting the whole non-sampled span.""" trainable_branches = [branch for branch, _ in iter_trainable_branches(trace)] filter_masks = self._filter_masks(trace, trainable_branches) if self.filter_fn is not None else None + # A re-rendered history may have different tokens and therefore fork the + # physical graph. Match message paths to retain actual response provenance. + paths = {} + path_by_node = {} + sampled_paths = set() + weighted_nodes = {} + for node in trace.nodes: + parent = path_by_node[id(trace.nodes[node.parent])] if node.parent is not None else None + key = _node_key(parent, node.message, node.tools) + path = paths.setdefault(key, len(paths)) + path_by_node[id(node)] = path + if node.sampled: + sampled_paths.add(path) for branch_idx, branch in enumerate(trainable_branches): weights = [0.0] * len(branch.token_ids) offset = 0 @@ -64,7 +78,14 @@ def _weight_observations(self, trace: vf.Trace) -> None: for node in branch.nodes: span = len(node.token_ids) role = node.message.role - if seen_response and not node.sampled and role in self.role_weights: + path = path_by_node[id(node)] + is_response = path in sampled_paths + if ( + seen_response + and not is_response + and role in self.role_weights + and weighted_nodes.get(path, id(node)) == id(node) + ): weight = self.role_weights[role] keep_mask = filter_masks[branch_idx] if filter_masks is not None else None # Per-token content granularity when the renderer attributed it; otherwise @@ -75,7 +96,8 @@ def _weight_observations(self, trace: vf.Trace) -> None: continue if keep_mask is None or keep_mask[i]: weights[i] = weight - if node.sampled: + weighted_nodes[path] = id(node) + if is_response: seen_response = True offset += span offset = 0 diff --git a/src/prime_rl/orchestrator/clients.py b/src/prime_rl/orchestrator/clients.py index c6d2de97f87..1fb69962231 100644 --- a/src/prime_rl/orchestrator/clients.py +++ b/src/prime_rl/orchestrator/clients.py @@ -9,16 +9,9 @@ import verifiers.v1 as vf from httpx import AsyncClient from openai import AsyncOpenAI -from renderers import RendererConfig from tenacity import AsyncRetrying, retry, retry_if_exception, stop_after_attempt, stop_after_delay, wait_exponential from verifiers.v1.clients.base import build_async_openai -from verifiers.v1.configs.client import ( - BaseClientConfig, - EvalClientConfig, - TrainClientConfig, - resolve_api_key, - resolve_headers, -) +from verifiers.v1.configs.client import resolve_api_key, resolve_headers from prime_rl.configs.algorithm import FrozenModelConfig from prime_rl.configs.shared import ClientConfig, PolicyClientConfig @@ -27,7 +20,7 @@ class PrefillScorer: """Prefill-scores token ids against a pool's endpoint, lazily resolving - a single OpenAI client from the pool's train client config.""" + a single OpenAI client from the pool's endpoint config.""" def __init__(self) -> None: self._client: AsyncOpenAI | None = None @@ -51,18 +44,8 @@ def __init__( self, client_config: ClientConfig, model_name: str, - train_client_type: str = "openai_chat_completions", - eval_client_type: str = "openai_chat_completions", - renderer_config: RendererConfig | None = None, ): - renderer_model_name = model_name if train_client_type == "renderer" else None - self.train_client = setup_client( - client_config, - client_type=train_client_type, - renderer_config=renderer_config, - renderer_model_name=renderer_model_name, - ) - self.eval_client = setup_client(client_config, client_type=eval_client_type) + self.client = setup_client(client_config) self._scorer = PrefillScorer() # Managed routed deployments set admin_base_url so engine admin traffic # bypasses the client-facing router. External and frozen clients do not. @@ -76,7 +59,7 @@ def __init__( async def score(self, token_ids: list[int]) -> list[float]: """Prefill-score ``token_ids`` under this endpoint's model (one logprob per token, 0.0 for the leading token).""" - return await self._scorer.score(self.train_client, self.model_name, token_ids) + return await self._scorer.score(self.client, self.model_name, token_ids) async def aclose(self) -> None: await self._scorer.aclose() @@ -229,40 +212,17 @@ async def check_inference_ready(client_config: ClientConfig, model_name: str) -> await admin.aclose() -async def connect_frozen_client( - config: FrozenModelConfig, *, renderer_config: RendererConfig | None = None -) -> InferenceClient: +async def connect_frozen_client(config: FrozenModelConfig) -> InferenceClient: """Connect to an externally hosted frozen model and wait for it.""" get_logger().info(f"Initializing frozen model pool (model={config.name}, base_url={config.base_url})") - if renderer_config is not None: - clients = InferenceClient( - config, model_name=config.name, train_client_type="renderer", renderer_config=renderer_config - ) - else: - clients = InferenceClient(config, model_name=config.name) + clients = InferenceClient(config, model_name=config.name) await check_inference_ready(config, config.name) return clients -def setup_client( - client_config: ClientConfig, - client_type: str = "openai_chat_completions", - renderer_config: RendererConfig | None = None, - renderer_model_name: str | None = None, -) -> vf.ClientConfig: - """Build a v1 client config for the base URL. ``client_type`` - ``renderer`` → token-in/out (``TrainClientConfig``, with the renderer the env - server should use forwarded as a serialized config so it doesn't fall back to the - default renderer); otherwise plain chat-completions (``EvalClientConfig``).""" - is_renderer = client_type == "renderer" - config_cls = TrainClientConfig if is_renderer else EvalClientConfig - renderer_extra: dict = {} - if is_renderer: - renderer_extra = { - "renderer": renderer_config, - "renderer_model_name": renderer_model_name, - } - return config_cls(**client_config.model_dump(include=set(BaseClientConfig.model_fields)), **renderer_extra) +def setup_client(client_config: ClientConfig) -> vf.ClientConfig: + """Build the shared evaluation and training endpoint config.""" + return vf.ClientConfig(**client_config.model_dump(include=set(vf.ClientConfig.model_fields))) def setup_admin_clients(client_config: ClientConfig) -> list[AsyncClient]: diff --git a/src/prime_rl/orchestrator/dispatcher.py b/src/prime_rl/orchestrator/dispatcher.py index 4a4222f82af..c52f1022abc 100644 --- a/src/prime_rl/orchestrator/dispatcher.py +++ b/src/prime_rl/orchestrator/dispatcher.py @@ -545,17 +545,13 @@ async def schedule_group_episode(self, group_id: uuid.UUID, group: GroupState) - ready, no permits). Returns True after issuing one task — the caller loops to keep scheduling. """ - # Train rollouts use the env's generation source via the - # renderer/token train client. Eval always evaluates the policy and - # goes through the eval client (chat-completions) so eval scores stay - # comparable. if group.kind == "eval": clients, model_name = self.policy_clients, self.policy.model_name live_sourced = True else: clients, model_name, live_sourced = self._train_generation_for(group.env_name) - client = clients.eval_client if group.kind == "eval" else clients.train_client + client = clients.client env_collection = self.train_envs if group.kind == "train" else self.eval_envs if env_collection is None: diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 15a65d1c993..6dbcd0862b5 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -26,7 +26,6 @@ from typing import Generic, TypeVar import verifiers.v1 as vf -from renderers import RendererConfig from verifiers.v1.serve import EnvClient from prime_rl.configs.algorithm import FrozenModelConfig @@ -144,14 +143,12 @@ def __init__( address: str | None, address_file: Path, clients: InferenceClient, - renderer_config: RendererConfig | None, algorithm: Algorithm, ): super().__init__(config, address, address_file) # Train rollouts are generated from `clients`: the policy, or the frozen - # `sampling.source` connected in setup() with the renderer (token-in/out) client. + # `sampling.source` connected in setup(). self.clients = clients - self.renderer_config = renderer_config self.connected: InferenceClient | None = None self.algorithm = algorithm self.uses_live_policy = config.algo.sampling.source == "policy" @@ -159,6 +156,7 @@ def __init__( **config.sampling.model_dump(exclude_none=True), # Only policy rollouts need sampling logprobs for importance ratios. logprobs=True if self.uses_live_policy else None, + return_training_metadata=True, ) # Truncated policy sampling must ship the sampling masks the trainer replays. self.requires_sampling_masks = config.sampling.truncates_distribution() and self.uses_live_policy @@ -167,7 +165,7 @@ async def setup(self) -> None: async def connect_source() -> None: source = self.config.algo.sampling.source if isinstance(source, FrozenModelConfig): - self.connected = await connect_frozen_client(source, renderer_config=self.renderer_config) + self.connected = await connect_frozen_client(source) self.clients = self.connected await asyncio.gather(connect_source(), self.algorithm.setup()) @@ -243,7 +241,6 @@ def __init__( config_dir: Path, *, clients, - renderer_config=None, ): self._envs: dict[str, TrainEnv] = {} for config in configs: @@ -253,7 +250,6 @@ def __init__( addresses[("train", config.resolved_name)], env_address_file(config_dir, "train", config.resolved_name), clients, - renderer_config, build_algorithm(config.algo, clients), ) self._envs[env.name] = env diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index f0299bb3cf9..f1577f18d00 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -195,9 +195,6 @@ async def setup(self) -> None: self.clients = InferenceClient( config.model.client, model_name=config.model.name, - train_client_type="renderer", - eval_client_type="openai_chat_completions", - renderer_config=config.renderer, ) self.admin_plane = setup_admin_plane(config.model.client, config.model.name) @@ -228,7 +225,6 @@ async def setup(self) -> None: config.env_addresses, config_dir, clients=self.clients, - renderer_config=config.renderer, ) if config.eval is not None: self.eval_envs = EvalEnvs(config.eval.source, config.env_addresses, config_dir) diff --git a/src/prime_rl/orchestrator/train_sink.py b/src/prime_rl/orchestrator/train_sink.py index 5c835a22e4e..3c0ee1451d3 100644 --- a/src/prime_rl/orchestrator/train_sink.py +++ b/src/prime_rl/orchestrator/train_sink.py @@ -246,7 +246,9 @@ async def process_group(self, group_id: str) -> None: samples_by_trace: dict[str, list[TrainingSample]] = {} temperature = env.config.sampling.temperature for trace in survivors: - samples = await asyncio.to_thread(trace_to_samples, trace, env_name=env_name) + samples = await asyncio.to_thread( + trace_to_samples, trace, env_name=env_name, require_logprobs=env.algorithm.action_loss_type != "ce" + ) for sample in samples: sample.temperatures = [temperature] * len(sample.token_ids) if env.requires_sampling_masks and sample.sampling_mask is None: diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 4d584577921..ff5a5ced167 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -7,9 +7,9 @@ training sample directly. Token-length readers (`completion_len`, `total_tokens`, `num_turns`) live on `vf.Trace` itself. -Training is renderer-only across every mode (RL/OPD student, SFT teacher), so every node -always carries its tokens — no backfill needed. Multimodal RL keeps the inline image URLs on -the messages and pairs them with vLLM's expanded image-token runs here. +The inference endpoint supplies exact token records for RL and frozen-model generation. +No client-side tokenization or token backfill is needed. Multimodal RL keeps the inline image +URLs on the messages and pairs them with the expanded image-token runs here. """ from __future__ import annotations @@ -147,7 +147,7 @@ def _loss_weights(branch: vf.Branch, name: str, trained_nodes: set[int]) -> list return weights if any(weights) else None -def trace_to_samples(trace: vf.Trace, *, env_name: str = "") -> list[TrainingSample]: +def trace_to_samples(trace: vf.Trace, *, env_name: str = "", require_logprobs: bool = True) -> list[TrainingSample]: """Convert a v1 `Trace` into `TrainingSample`s — one per branch. Each `trace.branches` entry is already a flat token sequence (`branch.token_ids` / @@ -160,6 +160,10 @@ def trace_to_samples(trace: vf.Trace, *, env_name: str = "") -> list[TrainingSam samples: list[TrainingSample] = [] trained_loss_nodes: dict[str, set[int]] = {"rl": set(), "ce": set(), "ref_kl": set()} for branch, mask in iter_trainable_branches(trace): + if require_logprobs: + for node in branch.nodes: + if node.sampled and len(node.logprobs) != sum(node.mask): + raise ValueError("RL samples require a logprob for every sampled token") token_ids = branch.token_ids mm_token_type_ids: list[int] | None = None mm_refs: MMRefs | None = None diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index f7bf2c26cb6..c59b919a37a 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -443,6 +443,41 @@ def test_echo_weights_observations_by_role(): asyncio.run(algo.score_episode(episode)) assert trace_to_samples(trace)[0].ce_weights is None + # Later calls can re-render the same message history into a new physical branch. + episode = _two_turn_episode() + trace = episode.traces[0] + original = list(trace.nodes) + for i, node in enumerate(original): + trace.nodes.append( + node.model_copy( + update={ + "parent": None if node.parent is None else len(original) + node.parent, + "sampled": i == len(original) - 1, + "token_ids": [token + 100 for token in node.token_ids], + "mask": node.mask if i == len(original) - 1 else [False] * len(node.token_ids), + "logprobs": node.logprobs if i == len(original) - 1 else [], + } + ) + ) + # Keep the first response's genuine sample; the second response is sampled in + # the re-rendered branch. Its tool observation still follows a real response. + del trace.nodes[2 : len(original)] + trace.nodes[3].parent = 2 + trace.nodes[4].parent = 3 + trace.nodes[5].parent = 4 + algo = _echo_algorithm() + asyncio.run(algo.score_episode(episode)) + samples = trace_to_samples(trace) + assert len(samples) == 2 + assert samples[0].token_ids == [1, 2, 3, 4] + assert samples[1].token_ids == [101, 102, 103, 104, 105, 106, 107, 108] + assert samples[1].mask == [False] * 6 + [True, True] + assert samples[1].ce_weights == [0.0] * 4 + [0.1, 0.1, 0.0, 0.0] + trace.nodes[-1].logprobs = [] + with pytest.raises(ValueError, match="logprob for every sampled token"): + trace_to_samples(trace) + assert trace_to_samples(trace, require_logprobs=False) + def test_echo_weights_only_content_tokens_when_is_content_present(): # The observation node [5,6] carries per-token is_content: the first token is diff --git a/tests/unit/orchestrator/test_clients.py b/tests/unit/orchestrator/test_clients.py index 4bc686dcce3..5302fe1404b 100644 --- a/tests/unit/orchestrator/test_clients.py +++ b/tests/unit/orchestrator/test_clients.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx -from verifiers.v1.configs.client import EvalClientConfig +from verifiers.v1.configs.client import ClientConfig as VFClientConfig from prime_rl.configs.shared import ClientConfig from prime_rl.orchestrator.clients import ( @@ -86,9 +86,7 @@ def test_admin_plane_initializes_nccl(): asyncio.run(admin_plane.aclose()) -def test_setup_client_creates_renderer_client(): - from renderers import Qwen3VLRendererConfig - +def test_setup_client_preserves_endpoint_headers(): client_config = ClientConfig( base_url="http://worker-a:8000/v1", api_key_var="PRIME_API_KEY", @@ -97,16 +95,9 @@ def test_setup_client_creates_renderer_client(): timeout={"read": 12}, ) - renderer_settings = Qwen3VLRendererConfig() - client = setup_client( - client_config, - client_type="renderer", - renderer_config=renderer_settings, - ) + client = setup_client(client_config) - assert client.type == "train" - assert client.renderer == renderer_settings - assert client.renderer_model_name is None + assert isinstance(client, VFClientConfig) assert client.base_url == "http://worker-a:8000/v1" assert "X-data-parallel-rank" not in client.headers assert client.headers["X-Test"] == "test" @@ -127,24 +118,6 @@ def test_check_health_retries_non_success_status(): assert client.get.await_count == 2 -def test_setup_client_assigns_renderer_model_name(): - from renderers import Qwen3VLRendererConfig - - client_config = ClientConfig( - base_url="http://worker-a:8000/v1", - api_key_var="PRIME_API_KEY", - ) - - client = setup_client( - client_config, - client_type="renderer", - renderer_config=Qwen3VLRendererConfig(), - renderer_model_name="Qwen/Qwen3-VL-4B-Instruct", - ) - - assert client.renderer_model_name == "Qwen/Qwen3-VL-4B-Instruct" - - def test_setup_client_preserves_chat_client_defaults(): client_config = ClientConfig( base_url="http://worker-a:8000/v1", @@ -153,7 +126,7 @@ def test_setup_client_preserves_chat_client_defaults(): client = setup_client(client_config) - assert client == EvalClientConfig( + assert client == VFClientConfig( api_key_var="PRIME_API_KEY", base_url="http://worker-a:8000/v1", headers={}, diff --git a/tests/unit/orchestrator/test_qwen3_vl_e2e.py b/tests/unit/orchestrator/test_qwen3_vl_e2e.py index 654673a2f1d..1c3e8f1d9ee 100644 --- a/tests/unit/orchestrator/test_qwen3_vl_e2e.py +++ b/tests/unit/orchestrator/test_qwen3_vl_e2e.py @@ -1,95 +1,44 @@ -"""End-to-end request-shape check for raw Qwen3-VL inference.""" +"""Qwen3-VL training records retain inline images and exact placeholder ranges.""" from __future__ import annotations -import asyncio import base64 -import json -from io import BytesIO +import io from pathlib import Path -from typing import Any +from threading import Lock +from types import SimpleNamespace -import httpx import pytest -_HF_CACHE = Path("~/.cache/huggingface/hub").expanduser() _MODEL = "Qwen/Qwen3-VL-4B-Instruct" - - -def _model_cached() -> bool: - safe = "models--" + _MODEL.replace("/", "--") - snapshots = _HF_CACHE / safe / "snapshots" - if not snapshots.is_dir(): - return False - return any(p.is_dir() for p in snapshots.iterdir()) - - +_HF_CACHE = Path("~/.cache/huggingface/hub").expanduser() pytestmark = pytest.mark.skipif( - not _model_cached(), + not (_HF_CACHE / ("models--" + _MODEL.replace("/", "--")) / "snapshots").is_dir(), reason=f"{_MODEL}: HF snapshot not cached locally", ) -class _FakeOpenAI: - """Minimal AsyncOpenAI stand-in that captures POST bodies. - - ``renderers.client.generate`` calls ``client.post(absolute_url, - body=...)``; we capture the body for assertions and return a canned - generate response so the parse-side of the flow runs. - """ - - def __init__(self, image_pad_id: int): - self.calls: list[dict[str, Any]] = [] - self.base_url = "http://fake-host:8000/v1" - self.image_pad_id = image_pad_id - - async def post(self, path, *, cast_to=dict, body=None, options=None): - self.calls.append({"path": path, "body": body, "options": options}) - # Reply with two sampled tokens + <|im_end|>. The renderer's - # parse_response slices the content tokens. - prompt_ids = list(body["token_ids"]) - pad_index = prompt_ids.index(self.image_pad_id) - prompt_ids[pad_index : pad_index + 1] = [self.image_pad_id] * 4 - payload = { - "request_id": "qwen-vl-e2e", - "prompt_token_ids": prompt_ids, - "mm_placeholders": {"image": [{"offset": pad_index, "length": 4}]}, - "choices": [ - { - "index": 0, - "token_ids": [50, 60, 151645], - "logprobs": { - "content": [ - {"token": "token_id:50", "logprob": -0.1}, - {"token": "token_id:60", "logprob": -0.2}, - {"token": "token_id:151645", "logprob": -0.3}, - ] - }, - "finish_reason": "stop", - }, - ], - } - return httpx.Response(200, content=json.dumps(payload).encode()) - - -def test_generate_qwen3_vl_sends_raw_content_and_uses_expanded_prompt_ids(): +def test_generate_qwen3_vl_e2e_preserves_inline_images_and_expanded_prompt_ids(): + import verifiers.v1 as vf from PIL import Image + from renderers import Qwen3VLRendererConfig from renderers.base import load_tokenizer - from renderers.client import generate from renderers.qwen3_vl import Qwen3VLRenderer + from transformers import AutoProcessor + from verifiers.v1.dialects import ChatDialect + from verifiers.v1.graph import prepare_turn + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest - tokenizer = load_tokenizer(_MODEL) - renderer = Qwen3VLRenderer(tokenizer) - - image_pad_id = tokenizer.convert_tokens_to_ids("<|image_pad|>") + from prime_rl.inference.vllm.serving_chat import PrimeRlServingChat + from prime_rl.orchestrator.trajectories import trace_to_samples - fake = _FakeOpenAI(image_pad_id) - - # ── Build a user message with an image (OpenAI content-part shape). ─ - img = Image.new("RGB", (224, 224), color=(64, 128, 255)) - buffer = BytesIO() - img.save(buffer, format="PNG") - image_url = f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode()}" + tokenizer = load_tokenizer(_MODEL) + processor = AutoProcessor.from_pretrained(_MODEL) + renderer = Qwen3VLRenderer(tokenizer, processor=processor) + image = Image.new("RGB", (224, 224), color=(64, 128, 255)) + buffer = io.BytesIO() + image.save(buffer, format="PNG") + image_url = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() messages = [ { "role": "user", @@ -99,26 +48,63 @@ def test_generate_qwen3_vl_sends_raw_content_and_uses_expanded_prompt_ids(): ], } ] - - result = asyncio.run( - generate( - client=fake, - renderer=renderer, - messages=messages, - model=_MODEL, - sampling_params={"max_tokens": 16}, - # Explicit cap so generate() skips the /v1/models discovery round-trip. - max_prompt_len=1_000_000, - process_multimodal=False, - ) + request = ChatCompletionRequest(model=_MODEL, messages=messages, return_token_ids=True) + serving = object.__new__(PrimeRlServingChat) + serving.model_config = SimpleNamespace(tokenizer=_MODEL, max_model_len=1_000_000) + serving.chat_template = None + serving.chat_template_content_format = "auto" + serving.default_chat_template_kwargs = {} + serving.training_renderer_lock = Lock() + serving.training_tokenizer = tokenizer + serving.training_renderer = renderer + serving.training_template_kwargs = {} + serving.training_renderer_config = Qwen3VLRendererConfig() + _, (engine_input,) = serving._render_training_prompt(request) + + assert engine_input["type"] == "multimodal" + assert len(engine_input["mm_hashes"]["image"]) == 1 + (placeholder,) = engine_input["mm_placeholders"]["image"] + pad_ids = engine_input["prompt_token_ids"][placeholder.offset : placeholder.offset + placeholder.length] + assert pad_ids and all(token == tokenizer.convert_tokens_to_ids("<|image_pad|>") for token in pad_ids) + (item,) = engine_input["mm_kwargs"]["image"] + assert set(item) == {"pixel_values", "image_grid_thw"} + expected = processor.image_processor(images=[image], return_tensors="pt") + assert item["image_grid_thw"].data.tolist() == expected["image_grid_thw"][0].tolist() + + dialect = ChatDialect() + response = dialect.parse_response( + { + "id": "vlm-record", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "prompt_token_ids": engine_input["prompt_token_ids"], + "training_metadata": request._prime_training_metadata, + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "A picture."}, + "token_ids": [50, 60, 151645], + "logprobs": { + "content": [{"token": f"token_id:{token}", "logprob": -0.1} for token in [50, 60, 151645]] + }, + } + ], + } ) - - assert len(fake.calls) == 1 - body = fake.calls[0]["body"] - assert "features" not in body - assert body["content_parts"] == [{"type": "image_url", "url": image_url}] - assert body["token_ids"].count(image_pad_id) == 1 - assert result["renderer_prompt_ids"] == body["token_ids"] - assert result["prompt_ids"].count(image_pad_id) == 4 - assert result["mm_placeholders"] == {"image": [{"offset": body["token_ids"].index(image_pad_id), "length": 4}]} - assert result["completion_ids"] == [50, 60, 151645] + assert response.tokens is not None + assert "multi_modal_data" not in request._prime_training_metadata + assert response.tokens.mm_placeholders == [(placeholder.offset, placeholder.length)] + assert len(response.tokens.is_content) == len(engine_input["prompt_token_ids"]) + trace = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="image")), + ) + prompt, _ = dialect.parse_request({"messages": messages}) + prepare_turn(trace, prompt.messages).commit(response) + (sample,) = trace_to_samples(trace) + assert sample.token_ids == [*engine_input["prompt_token_ids"], 50, 60, 151645] + assert sample.mm_refs is not None + (image_ref,) = sample.mm_refs.images + assert (image_ref.url, image_ref.offset, image_ref.length) == (image_url, placeholder.offset, placeholder.length) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index dd8daee41b5..53a2930bcf8 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -440,7 +440,6 @@ def test_client_defaults_survive_partial_overrides( def test_env_algo_inherits_the_group_algo(): config = OrchestratorConfig.model_validate( { - "renderer": {"name": "qwen3"}, # echo needs the renderer's role attribution "train": { "algo": {"type": "echo", "roles": {"user": {"alpha": 0.1}}}, "source": [ @@ -573,7 +572,6 @@ def test_policy_sources_accept_different_top_p_values(): with pytest.warns(UserWarning, match="defaulting top_k"): config = OrchestratorConfig.model_validate( { - "renderer": {"name": "qwen3"}, "train": { "source": [ { @@ -632,7 +630,6 @@ def test_policy_sources_reject_mixed_top_k_capture(): with pytest.raises(ValidationError, match="cannot mix top_k > 0 and top_k = -1"): OrchestratorConfig.model_validate( { - "renderer": {"name": "qwen3"}, "train": { "source": [ { @@ -731,34 +728,11 @@ def test_disaggregated_inference_inherits_deployment_gpus_per_node(): RLConfig.model_validate(config) -def test_orchestrator_vlm_requires_renderer(): +def test_inference_vlm_requires_renderer(): with pytest.raises(ValidationError, match="renderer"): - OrchestratorConfig.model_validate( - { - "model": { - "name": "Qwen/Qwen3-VL-4B-Instruct", - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", - }, - }, - "renderer": None, - } - ) - - config = OrchestratorConfig.model_validate( - { - "model": { - "name": "Qwen/Qwen3-VL-4B-Instruct", - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", - }, - }, - } - ) - - assert config.renderer is not None + InferenceConfig.model_validate({"vllm": {"model": "Qwen/Qwen3-VL-4B-Instruct"}, "renderer": None}) + config = InferenceConfig.model_validate({"vllm": {"model": "Qwen/Qwen3-VL-4B-Instruct"}}) + assert config.renderer.name == "auto" def test_trainer_rejects_vlm_cp_with_ring(): @@ -790,7 +764,7 @@ def test_shared_model_name_propagates_to_subconfigs(): { "model": {"name": model_name}, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, "inference": {}, } ) @@ -807,7 +781,7 @@ def test_shared_tokenizer_propagates_when_subconfigs_unset(): "model": {"name": "my-model"}, "tokenizer": {"name": "my-tokenizer"}, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) assert config.trainer.tokenizer.name == "my-tokenizer" @@ -824,7 +798,7 @@ def test_shared_and_sub_tokenizer_name_conflict_raises(): "model": {"name": "my-model"}, "tokenizer": {"name": "shared-tok"}, "trainer": {"tokenizer": {"name": "trainer-tok"}}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) @@ -835,7 +809,7 @@ def test_tokenizer_name_falls_back_to_model_name_when_unset(): "model": {"name": "my-model"}, "tokenizer": {"trust_remote_code": True}, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) assert config.trainer.tokenizer.name == "my-model" @@ -863,7 +837,6 @@ def test_explicit_subconfig_tokenizer_name_survives_shared_model_propagation(): "model": {"name": "M"}, "trainer": {}, "orchestrator": { - "renderer": {"name": "default"}, "tokenizer": {"name": "explicit-orch-tok"}, }, } @@ -882,7 +855,7 @@ def test_tokenizer_chat_template_mismatch_raises(): RLConfig.model_validate( { "trainer": {"tokenizer": {"chat_template": "A"}}, - "orchestrator": {"renderer": {"name": "default"}, "tokenizer": {"chat_template": "B"}}, + "orchestrator": {"tokenizer": {"chat_template": "B"}}, } ) @@ -892,7 +865,7 @@ def test_shared_seq_len_propagates_to_subconfigs(): { "seq_len": 4096, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) assert config.trainer.model.seq_len == 4096 @@ -908,7 +881,7 @@ def test_shared_and_sub_seq_len_conflict_raises(): { "seq_len": 4096, "trainer": {"model": {"seq_len": 8192}}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) @@ -920,7 +893,7 @@ def test_shared_and_sub_model_name_conflict_raises(): { "model": {"name": "X"}, "trainer": {"model": {"name": "Y"}}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) @@ -932,7 +905,7 @@ def test_shared_and_sub_max_steps_conflict_raises(): { "max_steps": 100, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}, "max_steps": 200}, + "orchestrator": {"max_steps": 200}, } ) @@ -947,7 +920,7 @@ def test_trainer_chat_template_cascades_to_inference(): { "model": {"name": "Qwen/Qwen3-0.6B"}, "trainer": {"tokenizer": {"chat_template": "TPL"}}, - "orchestrator": {"renderer": {"name": "default"}, "tokenizer": {"chat_template": "TPL"}}, + "orchestrator": {"tokenizer": {"chat_template": "TPL"}}, "inference": {}, } ) @@ -975,7 +948,7 @@ def test_shared_wandb_fields_propagate_to_subconfigs(): } }, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) for component in (config.trainer.monitors.wandb, config.orchestrator.monitors.wandb): @@ -997,7 +970,7 @@ def test_shared_monitor_disable_and_prime_propagate(): "model": {"name": "Qwen/Qwen3-0.6B"}, "monitors": {"wandb": "None", "file": "None", "prime": {"name": "shared-prime"}}, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) assert config.trainer.monitors.wandb is None and config.trainer.monitors.file is None @@ -1013,7 +986,7 @@ def test_empty_shared_ckpt_block_does_not_conflict_with_subconfig_ckpt(): { "ckpt": {}, # empty block, no field set "trainer": {"ckpt": {"interval": 50}}, - "orchestrator": {"renderer": {"name": "default"}, "ckpt": {"interval": 50}}, + "orchestrator": {"ckpt": {"interval": 50}}, } ) assert config.trainer.ckpt is not None @@ -1027,7 +1000,7 @@ def test_shared_and_subconfig_disjoint_fields_coexist(): { "model": {"name": "Qwen/Qwen3-0.6B"}, "trainer": {"model": {"attn": "flash_attention_3"}}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, } ) assert config.trainer.model.name == "Qwen/Qwen3-0.6B" @@ -1060,19 +1033,6 @@ def test_run_dir_propagates_through_cli(tmp_path): assert config.orchestrator.monitors.wandb is not None and config.orchestrator.monitors.wandb.name == "my-exp" -def test_orchestrator_renderer_auto_rejects_unmapped_model(): - """Default ``renderer`` (AutoRendererConfig) must reject models not in MODEL_RENDERER_MAP.""" - with pytest.raises(ValidationError, match="silently fall back to DefaultRenderer"): - OrchestratorConfig.model_validate({"model": {"name": "not-a-real-org/not-a-real-model"}}) - - -def test_orchestrator_renderer_auto_accepts_mapped_model(): - """The default Qwen model is in MODEL_RENDERER_MAP and should validate cleanly.""" - config = OrchestratorConfig.model_validate({"model": {"name": "Qwen/Qwen3-0.6B"}}) - assert config.renderer is not None - assert config.renderer.name == "auto" - - def test_sft_renderer_auto_accepts_prime_qwen_model(): config = SFTConfig.model_validate({"model": {"name": "PrimeIntellect/Qwen3-0.6B"}}) assert config.renderer.name == "auto" @@ -1093,40 +1053,15 @@ def test_sft_allows_unused_default_renderer_for_fake_data(): assert config.renderer.name == "default" -def test_orchestrator_explicit_renderer_skips_unmapped_check(): - """Explicit renderer.name bypasses the auto-resolution check — user opted in.""" - config = OrchestratorConfig.model_validate( - { - "model": {"name": "not-a-real-org/not-a-real-model"}, - "renderer": {"name": "qwen3"}, - } - ) - assert config.renderer is not None +def test_inference_renderer_reaches_server_namespace(): + config = InferenceConfig.model_validate({"vllm": {"model": "org/custom-model"}, "renderer": {"name": "qwen3"}}) + assert config.to_namespace().prime_renderer == config.renderer assert config.renderer.name == "qwen3" -def test_orchestrator_renderer_none_rejected(): - """A renderer is required (training is renderer-only): the non-optional type rejects None.""" - with pytest.raises(ValidationError, match="renderer"): - OrchestratorConfig.model_validate( - { - "model": {"name": "not-a-real-org/not-a-real-model"}, - "renderer": None, - } - ) - - -def test_orchestrator_explicit_default_renderer_with_unmapped_model(): - """renderer.name='default' is an explicit opt-in to DefaultRenderer and must pass.""" - config = OrchestratorConfig.model_validate( - { - "model": {"name": "not-a-real-org/not-a-real-model"}, - "renderer": {"name": "default", "tool_parser": "qwen3"}, - } - ) - assert config.renderer is not None - assert config.renderer.name == "default" - assert config.renderer.tool_parser == "qwen3" +def test_orchestrator_accepts_custom_inference_model(): + config = OrchestratorConfig.model_validate({"model": {"name": "org/custom-model"}}) + assert config.model.name == "org/custom-model" def test_shared_model_name_resolves_inference_parsers(): @@ -1138,7 +1073,7 @@ def test_shared_model_name_resolves_inference_parsers(): { "model": {"name": "Qwen/Qwen3-Coder-30B-A3B-Instruct"}, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, "inference": {}, } ) @@ -1154,7 +1089,7 @@ def test_explicit_inference_parser_wins_over_auto(): { "model": {"name": "Qwen/Qwen3-Coder-30B-A3B-Instruct"}, "trainer": {}, - "orchestrator": {"renderer": {"name": "default"}}, + "orchestrator": {}, "inference": {"vllm": {"tool_call_parser": "hermes"}}, } ) From d84aeaeefcc5c16b3f4ee425878d6eefc7733e53 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:49:52 +0200 Subject: [PATCH 02/13] Use the restructured inference stack and correct renderer import --- src/prime_rl/inference/vllm/serving_chat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/inference/vllm/serving_chat.py b/src/prime_rl/inference/vllm/serving_chat.py index 9ee867a9fc6..3d737ddfa3b 100644 --- a/src/prime_rl/inference/vllm/serving_chat.py +++ b/src/prime_rl/inference/vllm/serving_chat.py @@ -6,8 +6,8 @@ import json from threading import Lock -from renderers import Renderer, RendererConfig, create_renderer, load_tokenizer -from renderers.base import ToolCallParseStatus, _resolve_renderer_config, is_multimodal +from renderers import Renderer, RendererConfig, create_renderer +from renderers.base import ToolCallParseStatus, _resolve_renderer_config, is_multimodal, load_tokenizer from renderers.client import _build_mm_features from vllm.entrypoints.generate.base.protocol import FunctionCall, ToolCall from vllm.entrypoints.openai.chat_completion.protocol import ( From 1ca61bef629c23d682464cf5862b4a1801e7b82d Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Sat, 26 Sep 2026 10:16:06 +0200 Subject: [PATCH 03/13] Fix renderer defaults and optional training attribution --- src/prime_rl/inference/vllm/serving_chat.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/inference/vllm/serving_chat.py b/src/prime_rl/inference/vllm/serving_chat.py index 3d737ddfa3b..c119dc37772 100644 --- a/src/prime_rl/inference/vllm/serving_chat.py +++ b/src/prime_rl/inference/vllm/serving_chat.py @@ -64,8 +64,14 @@ def _render_training_prompt(self, request: ChatCompletionRequest): self.training_renderer_config = _resolve_renderer_config( self.training_tokenizer, self.training_renderer_config ) - kwargs = {**self.default_chat_template_kwargs, **(request.chat_template_kwargs or {})} template_fields = self.training_renderer_config.template_field_names() + # vLLM defaults may include template options for unrelated models. + kwargs = { + key: value + for key, value in self.default_chat_template_kwargs.items() + if self.training_renderer_config._allow_opaque_template_kwargs or key in template_fields + } + kwargs.update(request.chat_template_kwargs or {}) if request.reasoning_effort is not None: if "reasoning_effort" in template_fields: kwargs.setdefault("reasoning_effort", request.reasoning_effort) @@ -96,7 +102,7 @@ def _render_training_prompt(self, request: ChatCompletionRequest): request._prime_response_renderer = (renderer, tools) request._prime_training_metadata = { "message_spans": rendered.message_token_spans(), - "is_content": rendered.is_content, + "is_content": rendered.is_content or None, "mm_token_type_id_map": renderer.mm_token_type_id_map if is_multimodal(renderer) else None, "mm_placeholders": sorted( (span.offset, span.length) for ranges in mmd.mm_placeholders.values() for span in ranges From ecb34e9d747ab7de83417d4a8794594374b68b97 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Fri, 2 Oct 2026 08:19:05 -0700 Subject: [PATCH 04/13] Validate frozen training vocabularies and constrain Chat replay metadata --- docs/algorithms.md | 5 ++ src/prime_rl/inference/vllm/server.py | 11 +++++ src/prime_rl/inference/vllm/serving_chat.py | 23 ++++++---- src/prime_rl/orchestrator/envs.py | 46 +++++++++++++++++-- src/prime_rl/orchestrator/orchestrator.py | 2 +- tests/integration/test_reverse_text_rl_sft.py | 2 + 6 files changed, 74 insertions(+), 15 deletions(-) diff --git a/docs/algorithms.md b/docs/algorithms.md index 15ca5c7350c..2ca72503177 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -544,6 +544,11 @@ A renderer outside the `renderers` package loads as a custom renderer: `name = " RL and distillation losses require sampling logprobs. CE-only frozen generation can omit logprobs, but its endpoint must still return exact tokens and training metadata. +Before frozen generation starts, its endpoint's `/v1/tokenizer` vocabulary must match +the policy tokenizer's token-to-ID mapping. Different vocabularies are rejected; frozen +tokens are not retokenized for the policy. +If a frozen endpoint's router does not forward `/v1/tokenizer`, set its `admin_base_url` +to the inference servers so their vocabularies can be checked directly. ### Discontinuous Trajectories diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index 5fa0b931409..e90e0ef9123 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -6,6 +6,7 @@ import uvloop from fastapi import APIRouter, Request from fastapi.responses import JSONResponse +from renderers.base import load_tokenizer from starlette.datastructures import State from vllm.engine.protocol import EngineClient from vllm.entrypoints.launchers.api_server.app_state import init_app_state @@ -41,6 +42,16 @@ def models(request: Request) -> OpenAIServingModels: } +@router.get("/v1/tokenizer") +def tokenizer(request: Request): + """Expose the token-to-ID mapping used for training Chat Completions.""" + chat = request.app.state.openai_serving_chat + with chat.training_renderer_lock: + if chat.training_tokenizer is None: + chat.training_tokenizer = load_tokenizer(chat.model_config.tokenizer) + return chat.training_tokenizer.get_vocab() + + @router.post("/pause") async def pause(request: Request): logger.debug("Received /pause request (mode=keep, clear_cache=False)") diff --git a/src/prime_rl/inference/vllm/serving_chat.py b/src/prime_rl/inference/vllm/serving_chat.py index c119dc37772..6b2766be965 100644 --- a/src/prime_rl/inference/vllm/serving_chat.py +++ b/src/prime_rl/inference/vllm/serving_chat.py @@ -85,8 +85,10 @@ def _render_training_prompt(self, request: ChatCompletionRequest): renderer = self.training_renderer messages = [dict(message) for message in request.messages] for message in messages: - if "reasoning" in message: - message["reasoning_content"] = message.pop("reasoning") + if (reasoning := message.pop("reasoning", None)) is not None: + message["reasoning_content"] = reasoning + if message.get("reasoning_content") is None: + message.pop("reasoning_content", None) tools = [tool.model_dump(exclude_none=True) for tool in request.tools] if request.tools else None rendered = renderer.render(messages, tools=tools, add_generation_prompt=request.add_generation_prompt) if len(rendered.token_ids) >= self.model_config.max_model_len: @@ -131,26 +133,29 @@ def _render_training_prompt(self, request: ChatCompletionRequest): return messages, [engine_input] async def chat_completion_full_generator(self, request, result_generator, *args, **kwargs): + training = getattr(request, "return_training_metadata", False) experts = {} masks = {} async def capture(): async for result in result_generator: for output in result.outputs: - routing = serialize_routed_experts(output.routed_experts, request.routed_experts_prompt_start) - if routing is not None: - experts[output.index] = routing - # Native vLLM emits .npy here; the PD router consumes compact arrays. + if training: + routing = serialize_routed_experts(output.routed_experts, request.routed_experts_prompt_start) + if routing is not None: + experts[output.index] = routing + if output.sampling_mask is not None: + masks[output.index] = output.sampling_mask.token_ids + # Replay arrays belong only in explicitly requested training records. output.routed_experts = None - if output.sampling_mask is not None: - masks[output.index] = output.sampling_mask.token_ids + output.sampling_mask = None yield result response_renderer = getattr(request, "_prime_response_renderer", None) if response_renderer is not None: kwargs["parser"] = None response = await super().chat_completion_full_generator(request, capture(), *args, **kwargs) - if not isinstance(response, ChatCompletionResponse): + if not training or not isinstance(response, ChatCompletionResponse): return response if response_renderer is not None: renderer, tools = response_renderer diff --git a/src/prime_rl/orchestrator/envs.py b/src/prime_rl/orchestrator/envs.py index 6dbcd0862b5..2c9981939c1 100644 --- a/src/prime_rl/orchestrator/envs.py +++ b/src/prime_rl/orchestrator/envs.py @@ -22,8 +22,9 @@ import asyncio import time from collections.abc import Callable, Iterator, Sequence +from contextlib import AsyncExitStack from pathlib import Path -from typing import Generic, TypeVar +from typing import TYPE_CHECKING, Generic, TypeVar import verifiers.v1 as vf from verifiers.v1.serve import EnvClient @@ -31,10 +32,13 @@ from prime_rl.configs.algorithm import FrozenModelConfig from prime_rl.configs.orchestrator import EnvConfig, EvalSourceConfig, TrainSourceConfig from prime_rl.orchestrator.algo import Algorithm, build_algorithm -from prime_rl.orchestrator.clients import InferenceClient, connect_frozen_client +from prime_rl.orchestrator.clients import InferenceClient, connect_frozen_client, setup_admin_clients from prime_rl.utils.logger import format_time, get_logger from prime_rl.utils.pathing import env_address_file +if TYPE_CHECKING: + from transformers.tokenization_utils import PreTrainedTokenizer + # Max wait for the env server to answer health. Generous because the launcher spawns # servers concurrently with the orchestrator, and a server imports its env package # before serving. @@ -161,12 +165,44 @@ def __init__( # Truncated policy sampling must ship the sampling masks the trainer replays. self.requires_sampling_masks = config.sampling.truncates_distribution() and self.uses_live_policy - async def setup(self) -> None: + async def setup(self, tokenizer: PreTrainedTokenizer) -> None: async def connect_source() -> None: source = self.config.algo.sampling.source if isinstance(source, FrozenModelConfig): - self.connected = await connect_frozen_client(source) - self.clients = self.connected + clients = await connect_frozen_client(source) + try: + async with AsyncExitStack() as stack: + admin_clients = [ + await stack.enter_async_context(client) for client in setup_admin_clients(source) + ] + # Frozen samples train their native IDs directly, so every ID must retain its meaning. + policy_vocab = tokenizer.get_vocab() + for client in admin_clients: + response = await client.get("/v1/tokenizer", timeout=source.wait_for_ready_timeout) + if response.status_code != 200: + raise ValueError( + f"Frozen generation requires GET /v1/tokenizer; {client.base_url} returned " + f"HTTP {response.status_code}. If the router does not forward this endpoint, " + "set admin_base_url to the inference engine URLs." + ) + vocab = response.json() + if ( + not isinstance(vocab, dict) + or not vocab + or any(type(index) is not int for index in vocab.values()) + ): + raise ValueError( + f"Frozen generation endpoint {client.base_url} returned an invalid vocabulary" + ) + if vocab != policy_vocab: + raise ValueError( + f"Frozen generation endpoint {client.base_url} must use the same token-to-ID mapping " + "as the policy tokenizer; training across different tokenizers is not supported." + ) + except BaseException: + await clients.aclose() + raise + self.clients = self.connected = clients await asyncio.gather(connect_source(), self.algorithm.setup()) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index f1577f18d00..9e7efbc2ac1 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -282,7 +282,7 @@ async def setup(self) -> None: get_logger().success(f"Policy inference pool ready after {format_time(time.perf_counter() - t0)}") # Build + ready pools for each env's frozen generation source and the # algorithm's frozen reference model - await asyncio.gather(*(env.setup() for env in self.train_envs)) + await asyncio.gather(*(env.setup(self.tokenizer) for env in self.train_envs)) get_logger().info(f"Initializing weight broadcast ({config.weight_broadcast})") t0 = time.perf_counter() diff --git a/tests/integration/test_reverse_text_rl_sft.py b/tests/integration/test_reverse_text_rl_sft.py index a5dd4c0346b..202c2b32943 100644 --- a/tests/integration/test_reverse_text_rl_sft.py +++ b/tests/integration/test_reverse_text_rl_sft.py @@ -59,6 +59,8 @@ def ref_inference(output_dir: Path) -> Generator[subprocess.Popen, None, None]: "inference", "--vllm.model", "PrimeIntellect/Qwen3-0.6B-Reverse-Text-RL", + "--renderer.name", + "prime-qwen3", "--server.port", str(REF_PORT), "--vllm.gpu-memory-utilization", From 56fb11afa696c93cb5a07b7a0b726e23ae87cf26 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Fri, 2 Oct 2026 09:04:25 -0700 Subject: [PATCH 05/13] fix: align CI configs with server-owned rendering --- configs/debug/eval/tb2.toml | 6 ++---- tests/unit/test_configs.py | 8 ++++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/configs/debug/eval/tb2.toml b/configs/debug/eval/tb2.toml index 7d6bde5e82b..644d37f4abd 100644 --- a/configs/debug/eval/tb2.toml +++ b/configs/debug/eval/tb2.toml @@ -9,7 +9,5 @@ group_size = 4 env.taskset.id = "terminal-bench-2" env.agent.harness.id = "bash" env.agent.timeout.rollout = 3600 -env.agent.retries.rules = [ - {type = "ProviderError", max_retries = 2}, - {type = "SandboxError", max_retries = 2}, -] +env.agent.retries.max_retries = 2 +env.agent.retries.include = ["ProviderError", "SandboxError"] diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 53a2930bcf8..373b377eef7 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -1160,18 +1160,18 @@ def test_sft_config_accepts_custom_renderer(tmp_path, custom_renderer_import_pat assert custom_renderer_config(config.renderer).instruction == "Be brief." -def test_orchestrator_config_accepts_custom_renderer(tmp_path, custom_renderer_import_path): - config_file = tmp_path / "orch.toml" +def test_inference_config_accepts_custom_renderer(tmp_path, custom_renderer_import_path): + config_file = tmp_path / "inference.toml" config_file.write_text( tomli_w.dumps( { - "model": {"name": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"}, + "vllm": {"model": "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT"}, "renderer": {"name": "custom", "import_path": custom_renderer_import_path}, } ) ) - config = cli(OrchestratorConfig, args=["@", str(config_file)]) + config = cli(InferenceConfig, args=["@", str(config_file)]) assert config.renderer.name == "custom" assert custom_renderer_config(config.renderer).instruction == "Think step by step." From b4aed52fd17e1c88883573464411b9e4dcce2187 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Fri, 2 Oct 2026 09:15:58 -0700 Subject: [PATCH 06/13] chore: align Verifiers stack with current main --- configs/debug/eval/tb2.toml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/configs/debug/eval/tb2.toml b/configs/debug/eval/tb2.toml index 644d37f4abd..7d6bde5e82b 100644 --- a/configs/debug/eval/tb2.toml +++ b/configs/debug/eval/tb2.toml @@ -9,5 +9,7 @@ group_size = 4 env.taskset.id = "terminal-bench-2" env.agent.harness.id = "bash" env.agent.timeout.rollout = 3600 -env.agent.retries.max_retries = 2 -env.agent.retries.include = ["ProviderError", "SandboxError"] +env.agent.retries.rules = [ + {type = "ProviderError", max_retries = 2}, + {type = "SandboxError", max_retries = 2}, +] From c573018c7a4525065b800c932d8c8d79f51c3f34 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Fri, 2 Oct 2026 10:09:32 -0700 Subject: [PATCH 07/13] fix: preserve all tokenizer vocabulary entries in JSON --- src/prime_rl/inference/vllm/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index e90e0ef9123..647f9a48b85 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -49,7 +49,7 @@ def tokenizer(request: Request): with chat.training_renderer_lock: if chat.training_tokenizer is None: chat.training_tokenizer = load_tokenizer(chat.model_config.tokenizer) - return chat.training_tokenizer.get_vocab() + return JSONResponse(chat.training_tokenizer.get_vocab()) @router.post("/pause") From a48cf6fa388c0d8cb9cd25883b74cace567b70aa Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:09:28 -0700 Subject: [PATCH 08/13] chore: pin the corrected Verifiers delivery fixture --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 00c1534182f..a78f58b13f1 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 00c1534182f5eab7beab1f5230662b960c45e6d3 +Subproject commit a78f58b13f10ffc50d49f06a7dee2a9281e71f0e From d1facb08f2c671504010b8c76d07d71c8c2a955b Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:25:29 -0700 Subject: [PATCH 09/13] refactor(inference): share endpoint rendering and packed training records --- deps/renderers | 2 +- deps/verifiers | 2 +- docs/algorithms.md | 7 +- docs/overview.md | 4 +- docs/training.md | 2 +- .../src/prime_rl/configs/inference.py | 2 +- pyproject.toml | 2 +- src/prime_rl/inference/vllm/server.py | 12 +- src/prime_rl/inference/vllm/serving_chat.py | 187 +++++++----------- src/prime_rl/inference/vllm/serving_tokens.py | 21 +- src/prime_rl/orchestrator/clients.py | 39 ++-- tests/unit/inference/test_serving_tokens.py | 4 +- tests/unit/orchestrator/test_qwen3_vl_e2e.py | 19 +- 13 files changed, 129 insertions(+), 174 deletions(-) diff --git a/deps/renderers b/deps/renderers index 5d870d8eaa6..a3c71b3dfa2 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 5d870d8eaa65114f5bef8a4cffc146fbcb418ec4 +Subproject commit a3c71b3dfa2382293ca60d246ac1265bd59b3338 diff --git a/deps/verifiers b/deps/verifiers index a78f58b13f1..1092382a904 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit a78f58b13f10ffc50d49f06a7dee2a9281e71f0e +Subproject commit 1092382a904f2d7b9db98ede289a658a2466b147 diff --git a/docs/algorithms.md b/docs/algorithms.md index 2ca72503177..fddf336c9d0 100644 --- a/docs/algorithms.md +++ b/docs/algorithms.md @@ -528,11 +528,14 @@ training metadata. Each call is a valid training record in its own right. Verifi records only when their token prefixes match exactly; a re-rendered history starts a new sample with its context masked out. -The inference server uses the `renderers` library to attribute prompt tokens to messages +The inference server uses the configured `renderers` implementation for every Chat +Completions prompt, and to attribute prompt tokens to messages and message bodies. ECHO uses this attribution to select observations, including when tokenization changes across turns. Multimodal records keep inline image URLs and exact placeholder ranges; the trainer builds image features from those URLs. Native vLLM handles -sampling, and the inference renderer parses responses and tool calls. Sampling masks and +sampling, response parsing, tools, reasoning, and streaming for both training and eval. +`return_training_metadata=true` adds exact token IDs and attribution; logprobs remain optional. +Sampling masks and expert routing travel alongside the tokens for replay in the trainer. ```toml diff --git a/docs/overview.md b/docs/overview.md index 7df75dac171..7aa9e7028f4 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -8,8 +8,8 @@ A `prime-rl` RL run is three cooperating processes: ![Architecture](assets/architecture.png) -- **Inference** — vLLM-backed server (or fleet) holding the current policy. The orchestrator drives rollouts through the token-in `/inference/v1/generate` route via the [`renderers`](https://github.com/PrimeIntellect-ai/renderers) package (OpenAI-compatible chat/completions routes are also exposed for external clients). We are trying to stay up-to-date with the latest vLLM features, you can read more about the supported features and deployment options in the dedicated [inference documentation](inference.md). -- **Orchestrator** — Lightweight CPU process that owns the data plane across many [`verifiers`](https://github.com/PrimeIntellect-ai/verifiers) training and eval environments. Each env runs in an isolated subprocess with a variable-size pool of env workers for scalability. The orchestrator drives multi-turn rollouts against the inference fleet (tool use, browsers, sandboxes, long horizons) without re-tokenizing across turns, computes advantages, packs the rollouts into training batches, and relays new weights from trainer to inference. +- **Inference** — vLLM-backed server (or fleet) holding the current policy. The orchestrator drives rollouts through `/v1/chat/completions`, requesting exact tokens and training attribution from the server. The endpoint's [`renderers`](https://github.com/PrimeIntellect-ai/renderers) configuration controls prompt rendering for both training and eval. See the [inference documentation](inference.md) for supported features and deployment options. +- **Orchestrator** — Lightweight CPU process that owns the data plane across many [`verifiers`](https://github.com/PrimeIntellect-ai/verifiers) training and eval environments. Each env runs in an isolated subprocess with a variable-size pool of env workers for scalability. The orchestrator drives multi-turn rollouts against the inference fleet (tool use, browsers, sandboxes, long horizons), merges exact token prefixes across turns, computes advantages, packs the rollouts into training batches, and relays new weights from trainer to inference. - **Trainer** — FSDP2 process group that consumes packed rollouts and steps the optimizer. We ship optimized custom modeling code for many MoE / dense / VLM families that unlocks advanced trainer parallelism — expert parallelism (EP, with DeepEP kernels) and context parallelism (CP) for long-sequence training — plus selective activation checkpointing, FP8 training on Hopper+, and LoRA. You can read more in the dedicated [training documentation](training.md). The three processes communicate through configurable transports — by default the trainer↔orchestrator rollout link uses ZMQ (each trainer rank requests a step before the orchestrator sends it, so the trainer only ever holds the batch it is training on), and weight broadcast uses NCCL for synchronous in-memory transfer (falling back to filesystem when LoRA is enabled or no inference server is configured). Swap the rollout link to the local filesystem (`rollout_transport.type = "filesystem"`) if you want rollouts persisted to disk. See [Scaling](scaling.md) for the deployment options. diff --git a/docs/training.md b/docs/training.md index 9b9fdf4e444..18f473e6977 100644 --- a/docs/training.md +++ b/docs/training.md @@ -193,7 +193,7 @@ import_path = "my_renderers.glm53_no_thinking.NoThinkingGLM53Renderer" clear_thinking = true ``` -The same config selects the renderer for RL under `[orchestrator.renderer]`. The [`renderers` README](https://github.com/PrimeIntellect-ai/renderers#custom-renderers) shows the class behind this example. +The same config selects the renderer for RL under `[inference.renderer]`. The [`renderers` README](https://github.com/PrimeIntellect-ai/renderers#custom-renderers) shows the class behind this example. **Renderer-backed tokenization.** SFT tokenization is renderer-only. The [`renderers`](algorithms.md#renderers) package owns message-to-token conversion and loss attribution end-to-end, so position-dependent chat templates (for example templates that strip past `` blocks across user turns) do not corrupt the loss mask. `[renderer]` defaults to `name = "auto"`; set a typed renderer config only when you need model-specific template controls. Hand-coded renderers ship for Qwen3, Qwen3.5, GLM-5, GLM-4.5, Kimi K2/K2.5, MiniMax M2, DeepSeek V3, Nemotron 3, GPT-OSS, and VLM families such as Qwen3-VL/Qwen3.5. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index dd04dd7ebfd..c8288b8b6ae 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -412,7 +412,7 @@ def num_nodes(self) -> int: class InferenceConfig(BaseConfig): renderer: RendererConfig = AutoRendererConfig() - """Server-side prompt rendering and message/image attribution for training records.""" + """Prompt rendering for every Chat Completions request, including training attribution.""" server: ServerConfig = ServerConfig() diff --git a/pyproject.toml b/pyproject.toml index 88ba2e2de57..cc962dde638 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "setproctitle>=1.3.0", "uvloop>=0.21.0", "verifiers[harbor]>=0.4.0", - "renderers[multimodal]>=0.1.12", + "renderers[multimodal]>=0.1.13.dev6", "pybase64>=1.4.2", "msgspec>=0.18", "fastapi>=0.115", diff --git a/src/prime_rl/inference/vllm/server.py b/src/prime_rl/inference/vllm/server.py index 647f9a48b85..9f2232a166b 100644 --- a/src/prime_rl/inference/vllm/server.py +++ b/src/prime_rl/inference/vllm/server.py @@ -46,10 +46,10 @@ def models(request: Request) -> OpenAIServingModels: def tokenizer(request: Request): """Expose the token-to-ID mapping used for training Chat Completions.""" chat = request.app.state.openai_serving_chat - with chat.training_renderer_lock: - if chat.training_tokenizer is None: - chat.training_tokenizer = load_tokenizer(chat.model_config.tokenizer) - return JSONResponse(chat.training_tokenizer.get_vocab()) + with chat.prompt_renderer_lock: + if chat.prompt_tokenizer is None: + chat.prompt_tokenizer = load_tokenizer(chat.model_config.tokenizer) + return JSONResponse(chat.prompt_tokenizer.get_vocab()) @router.post("/pause") @@ -161,8 +161,8 @@ async def custom_init_app_state( upstream = state.openai_serving_chat chat = object.__new__(PrimeRlServingChat) chat.__dict__.update(upstream.__dict__) - chat.training_renderer_config = args.prime_renderer - chat.training_renderer_lock = Lock() + chat.prompt_renderer_config = args.prime_renderer + chat.prompt_renderer_lock = Lock() state.openai_serving_chat = chat # Under RL load the API server holds millions of live objects, and CPython's diff --git a/src/prime_rl/inference/vllm/serving_chat.py b/src/prime_rl/inference/vllm/serving_chat.py index 6b2766be965..f3e1e17efb3 100644 --- a/src/prime_rl/inference/vllm/serving_chat.py +++ b/src/prime_rl/inference/vllm/serving_chat.py @@ -3,30 +3,41 @@ from __future__ import annotations import asyncio -import json from threading import Lock from renderers import Renderer, RendererConfig, create_renderer -from renderers.base import ToolCallParseStatus, _resolve_renderer_config, is_multimodal, load_tokenizer -from renderers.client import _build_mm_features -from vllm.entrypoints.generate.base.protocol import FunctionCall, ToolCall +from renderers.base import _resolve_renderer_config, is_multimodal, load_tokenizer +from renderers.client import build_mm_features from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, ) from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat -from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item -from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls from vllm.inputs import mm_input, tokens_input -from vllm.multimodal.inputs import MultiModalKwargsItems, PlaceholderRange -from prime_rl.inference.vllm.routed_experts import serialize_routed_experts +from prime_rl.inference.vllm.serving_tokens import _PackedOutputs + + +class PrimeRlChatRequest(ChatCompletionRequest): + return_training_metadata: bool = False + + def to_sampling_params(self, *args, **kwargs): + params = super().to_sampling_params(*args, **kwargs) + # Top-k and prompt-logprob responses retain vLLM's native representation. + params.flat_logprobs = ( + self.return_training_metadata + and params.prompt_logprobs is None + and not self.top_logprobs + and not self.logprob_token_ids + ) + return params class TrainingChoice(ChatCompletionResponseChoice): routed_experts: dict | None = None # type: ignore[assignment] - sampling_mask: list[list[int]] | None = None + completion_logprobs: dict | None = None + sampling_mask: dict | None = None class TrainingCompletion(ChatCompletionResponse): @@ -35,41 +46,53 @@ class TrainingCompletion(ChatCompletionResponse): class PrimeRlServingChat(OpenAIServingChat): - training_renderer_config: RendererConfig - training_renderer_lock: Lock - training_renderer: Renderer | None = None - training_tokenizer = None - training_template_kwargs: dict | None = None + prompt_renderer_config: RendererConfig + prompt_renderer_lock: Lock + prompt_renderer: Renderer | None = None + prompt_tokenizer = None + prompt_template_kwargs: dict | None = None + + async def create_chat_completion(self, request, raw_request=None): + request = PrimeRlChatRequest.model_validate(request.model_dump()) + if request.return_training_metadata: + if request.stream: + return self.create_error_response("Training records require stream=false") + request.return_token_ids = True + request.return_tokens_as_token_ids = True + request.include_reasoning = True + return await super().create_chat_completion(request, raw_request) async def render_chat_request(self, request: ChatCompletionRequest): - if not getattr(request, "return_training_metadata", False): - return await super().render_chat_request(request) - if request.stream or not request.return_token_ids or not request.include_reasoning: - raise ValueError("training metadata requires stream=false, return_token_ids=true, include_reasoning=true") if request.continue_final_message or request.chat_template is not None: - raise ValueError("training records require complete messages and the server's configured renderer") + raise ValueError("Chat requests require complete messages and the server's configured renderer") error = await self._check_model(request) if error is not None: return error self._preflight(request.n or 1) + if request.tool_choice not in (None, "none") and not self.online_renderer.use_harmony: + tool_parser = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None + if tool_parser is None or (request.tool_choice == "auto" and not self.enable_auto_tools): + return self.create_error_response("Tool choice requires an enabled tool-call parser") + if self.parser_cls is not None: + self._make_parser( + request, self.renderer.tokenizer, self._effective_chat_template_kwargs(request) + ).adjust_request(request) # Tokenizer/processor calls are synchronous. One worker-owned renderer is # sufficient; no rollout state or previously sampled tokens are kept here. - return await asyncio.to_thread(self._render_training_prompt, request) + return await asyncio.to_thread(self._render_prompt, request) - def _render_training_prompt(self, request: ChatCompletionRequest): + def _render_prompt(self, request: ChatCompletionRequest): # Hold the lock in the worker thread even if the HTTP request is cancelled. - with self.training_renderer_lock: - if self.training_tokenizer is None: - self.training_tokenizer = load_tokenizer(self.model_config.tokenizer) - self.training_renderer_config = _resolve_renderer_config( - self.training_tokenizer, self.training_renderer_config - ) - template_fields = self.training_renderer_config.template_field_names() + with self.prompt_renderer_lock: + if self.prompt_tokenizer is None: + self.prompt_tokenizer = load_tokenizer(self.model_config.tokenizer) + self.prompt_renderer_config = _resolve_renderer_config(self.prompt_tokenizer, self.prompt_renderer_config) + template_fields = self.prompt_renderer_config.template_field_names() # vLLM defaults may include template options for unrelated models. kwargs = { key: value for key, value in self.default_chat_template_kwargs.items() - if self.training_renderer_config._allow_opaque_template_kwargs or key in template_fields + if self.prompt_renderer_config._allow_opaque_template_kwargs or key in template_fields } kwargs.update(request.chat_template_kwargs or {}) if request.reasoning_effort is not None: @@ -77,12 +100,12 @@ def _render_training_prompt(self, request: ChatCompletionRequest): kwargs.setdefault("reasoning_effort", request.reasoning_effort) elif "enable_thinking" in template_fields: kwargs.setdefault("enable_thinking", request.reasoning_effort != "none") - if self.training_renderer is None or kwargs != self.training_template_kwargs: - self.training_renderer = create_renderer( - self.training_tokenizer, self.training_renderer_config, chat_template_kwargs=kwargs + if self.prompt_renderer is None or kwargs != self.prompt_template_kwargs: + self.prompt_renderer = create_renderer( + self.prompt_tokenizer, self.prompt_renderer_config, chat_template_kwargs=kwargs ) - self.training_template_kwargs = kwargs - renderer = self.training_renderer + self.prompt_template_kwargs = kwargs + renderer = self.prompt_renderer messages = [dict(message) for message in request.messages] for message in messages: if (reasoning := message.pop("reasoning", None)) is not None: @@ -90,6 +113,8 @@ def _render_training_prompt(self, request: ChatCompletionRequest): if message.get("reasoning_content") is None: message.pop("reasoning_content", None) tools = [tool.model_dump(exclude_none=True) for tool in request.tools] if request.tools else None + if request.tool_choice == "none" and self.online_renderer.exclude_tools_when_tool_choice_none: + tools = None rendered = renderer.render(messages, tools=tools, add_generation_prompt=request.add_generation_prompt) if len(rendered.token_ids) >= self.model_config.max_model_len: raise ValueError( @@ -100,100 +125,34 @@ def _render_training_prompt(self, request: ChatCompletionRequest): dict.fromkeys([*(request.stop_token_ids or []), *renderer.get_stop_token_ids()]) ) mmd = rendered.multi_modal_data - # Decode with this request's renderer even if another request replaces the cache. - request._prime_response_renderer = (renderer, tools) - request._prime_training_metadata = { - "message_spans": rendered.message_token_spans(), - "is_content": rendered.is_content or None, - "mm_token_type_id_map": renderer.mm_token_type_id_map if is_multimodal(renderer) else None, - "mm_placeholders": sorted( - (span.offset, span.length) for ranges in mmd.mm_placeholders.values() for span in ranges - ) - if mmd is not None - else None, - } + if request.return_training_metadata: + request._prime_training_metadata = { + "message_spans": rendered.message_token_spans(), + "is_content": rendered.is_content or None, + "mm_token_type_id_map": renderer.mm_token_type_id_map if is_multimodal(renderer) else None, + } engine_input = tokens_input(rendered.token_ids, cache_salt=request.cache_salt) if mmd is not None: - features = _build_mm_features(renderer, mmd) engine_input = mm_input( prompt_token_ids=rendered.token_ids, - mm_kwargs=MultiModalKwargsItems( - { - modality: [decode_mm_kwargs_item(item) for item in items] - for modality, items in features["kwargs_data"].items() - } - ), - mm_hashes=features["mm_hashes"], - mm_placeholders={ - modality: [PlaceholderRange(**span) for span in spans] - for modality, spans in features["mm_placeholders"].items() - }, + **build_mm_features(renderer, mmd), cache_salt=request.cache_salt, ) return messages, [engine_input] async def chat_completion_full_generator(self, request, result_generator, *args, **kwargs): - training = getattr(request, "return_training_metadata", False) - experts = {} - masks = {} - - async def capture(): - async for result in result_generator: - for output in result.outputs: - if training: - routing = serialize_routed_experts(output.routed_experts, request.routed_experts_prompt_start) - if routing is not None: - experts[output.index] = routing - if output.sampling_mask is not None: - masks[output.index] = output.sampling_mask.token_ids - # Replay arrays belong only in explicitly requested training records. - output.routed_experts = None - output.sampling_mask = None - yield result - - response_renderer = getattr(request, "_prime_response_renderer", None) - if response_renderer is not None: - kwargs["parser"] = None - response = await super().chat_completion_full_generator(request, capture(), *args, **kwargs) - if not training or not isinstance(response, ChatCompletionResponse): + packed = _PackedOutputs(result_generator, request, collect=request.return_training_metadata) + response = await super().chat_completion_full_generator(request, packed, *args, **kwargs) + if not request.return_training_metadata or not isinstance(response, ChatCompletionResponse): return response - if response_renderer is not None: - renderer, tools = response_renderer - for choice in response.choices: - parsed = await asyncio.to_thread(renderer.parse_response, choice.token_ids, tools=tools) - choice.message.content = parsed.content or None - choice.message.reasoning = parsed.reasoning_content - choice.message.tool_calls = [ - ToolCall( - **({"id": call.id} if call.id else {}), - function=FunctionCall( - name=call.name, - arguments=call.arguments - if isinstance(call.arguments, str) - else json.dumps(call.arguments or {}), - ), - ) - for call in parsed.tool_calls - if call.status == ToolCallParseStatus.OK - ] - if ( - choice.message.tool_calls - and choice.finish_reason == "stop" - and request.tool_choice in (None, "auto", "required") - ): - choice.finish_reason = "tool_calls" return TrainingCompletion( **response.model_dump(exclude={"choices"}), choices=[ - maybe_filter_parallel_tool_calls( - TrainingChoice( - **choice.model_dump(exclude={"routed_experts"}), - routed_experts=experts.get(choice.index), - sampling_mask=masks.get(choice.index), - ), - request, + TrainingChoice( + **choice.model_dump(exclude={"routed_experts"}), + **packed.fields.get(choice.index, {}), ) for choice in response.choices ], - training_metadata=getattr(request, "_prime_training_metadata", None), + training_metadata=request._prime_training_metadata, ) diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 70474b88ca2..920d6cc6e3a 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -31,6 +31,7 @@ import pybase64 from fastapi import Request from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateRequest, GenerateResponse, @@ -40,6 +41,7 @@ from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logprobs import FlatLogprobs from vllm.outputs import RequestOutput +from vllm.sampling_params import SamplingParams # vLLM's clamp for missing or -inf logprobs; renderers treat it as "no sampling evidence". LOGPROB_SENTINEL = -9999.0 @@ -78,14 +80,25 @@ class _PackedOutputs: """Wraps the result generator: takes the per-token payloads off each final output (so upstream builds no per-token objects) and keeps them packed.""" - def __init__(self, generator: AsyncIterator[RequestOutput], request: GenerateRequest): + def __init__( + self, + generator: AsyncIterator[RequestOutput], + logprob_options: SamplingParams | ChatCompletionRequest, + *, + collect: bool = True, + ): self._generator = generator - self._request = request + self._logprob_options = logprob_options + self._collect = collect self.fields: dict[int, dict[str, Any]] = {} async def __aiter__(self): async for request_output in self._generator: for output in request_output.outputs: + if not self._collect: + output.routed_experts = None + output.sampling_mask = None + continue fields = self.fields.setdefault(output.index, {}) if output.routed_experts is not None: # vLLM's capturer picks uint8 or uint16 from the model's expert count. Rows end at the @@ -100,7 +113,7 @@ async def __aiter__(self): fields["completion_logprobs"] = encode_array(pack_sampled_logprobs(output.logprobs)) output.logprobs = None # Upstream would build ``logprobs.content`` from the logprobs we took. - self._request.sampling_params.logprobs = None + self._logprob_options.logprobs = None mask = output.sampling_mask if mask is not None: fields["sampling_mask"] = {"ids": encode_array(mask.ids), "counts": encode_array(mask.counts)} @@ -142,7 +155,7 @@ async def serve_tokens_full_generator( # type: ignore[override] model_name: str, request_metadata: RequestResponseMetadata, ) -> ErrorResponse | GenerateResponse: - packed = _PackedOutputs(result_generator, request) + packed = _PackedOutputs(result_generator, request.sampling_params) response = await super().serve_tokens_full_generator( request, packed, diff --git a/src/prime_rl/orchestrator/clients.py b/src/prime_rl/orchestrator/clients.py index 1fb69962231..d55b72ec2b2 100644 --- a/src/prime_rl/orchestrator/clients.py +++ b/src/prime_rl/orchestrator/clients.py @@ -18,24 +18,6 @@ from prime_rl.utils.logger import get_logger -class PrefillScorer: - """Prefill-scores token ids against a pool's endpoint, lazily resolving - a single OpenAI client from the pool's endpoint config.""" - - def __init__(self) -> None: - self._client: AsyncOpenAI | None = None - - async def score(self, config: vf.ClientConfig, model: str, token_ids: list[int]) -> list[float]: - if self._client is None: - # Prefill scoring uses the SDK even when generation uses the eval relay. - self._client = build_async_openai(config) - return await prefill_logprobs(self._client, model, token_ids) - - async def aclose(self) -> None: - if self._client is not None: - await self._client.close() - - class InferenceClient: """Data-plane clients for one inference endpoint (the router of the policy deployment, or an external API for frozen models).""" @@ -46,7 +28,7 @@ def __init__( model_name: str, ): self.client = setup_client(client_config) - self._scorer = PrefillScorer() + self._score_client: AsyncOpenAI | None = None # Managed routed deployments set admin_base_url so engine admin traffic # bypasses the client-facing router. External and frozen clients do not. self._session_client = ( @@ -59,10 +41,13 @@ def __init__( async def score(self, token_ids: list[int]) -> list[float]: """Prefill-score ``token_ids`` under this endpoint's model (one logprob per token, 0.0 for the leading token).""" - return await self._scorer.score(self.client, self.model_name, token_ids) + if self._score_client is None: + self._score_client = build_async_openai(self.client) + return await prefill_logprobs(self._score_client, self.model_name, token_ids) async def aclose(self) -> None: - await self._scorer.aclose() + if self._score_client is not None: + await self._score_client.close() if self._session_client is not None: await self._session_client.aclose() @@ -476,21 +461,19 @@ async def prefill_logprobs(openai: AsyncOpenAI, model: str, token_ids: list[int] from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateResponse # `/inference/v1/generate` is mounted at server root, not under `/v1`: pass an - # absolute URL so the SDK skips the base-url merge. vLLM's `GenerateResponse` - # isn't an `openai.BaseModel`, so the SDK parse layer rejects it as `cast_to`; - # `cast_to=httpx.Response` lets the SDK still build the request (auth, retries, - # timeouts) and hand back the raw response for us to validate. + # absolute URL so the SDK skips the base-url merge. Parse JSON first because + # vLLM's `GenerateResponse` isn't an `openai.BaseModel`. base = str(openai.base_url).rstrip("/").removesuffix("/v1") - http_response = await openai.post( + payload = await openai.post( f"{base}/inference/v1/generate", - cast_to=httpx.Response, + cast_to=dict, body={ "model": model, "token_ids": token_ids, "sampling_params": {"max_tokens": 1, "temperature": 1.0, "top_p": 1.0, "prompt_logprobs": 1}, }, ) - response = GenerateResponse.model_validate_json(http_response.content) + response = GenerateResponse.model_validate(payload) # `prompt_logprobs[i]` is a `{token_id: Logprob}` dict, or `None` for the # leading token (no preceding context). Flatten to `list[float]`. flat: list[float] = [] diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 79cb3076982..5be21e15eaa 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -42,9 +42,7 @@ def test_subclass_overrides_serve_tokens_full_generator(): def test_generate_response_post_process_preserves_prompt_metadata(): compact_routed_experts = {"data": "AQID", "shape": [1, 1, 3], "start": 0} - capture = _PackedOutputs( - _empty_request_outputs(), SimpleNamespace(sampling_params=SimpleNamespace(routed_experts_prompt_start=0)) - ) + capture = _PackedOutputs(_empty_request_outputs(), SimpleNamespace(logprobs=None)) capture.fields[0] = {"routed_experts": compact_routed_experts} usage = UsageInfo(prompt_tokens=4, completion_tokens=3, total_tokens=7) response = GenerateResponse( diff --git a/tests/unit/orchestrator/test_qwen3_vl_e2e.py b/tests/unit/orchestrator/test_qwen3_vl_e2e.py index 1c3e8f1d9ee..69ea13a18f5 100644 --- a/tests/unit/orchestrator/test_qwen3_vl_e2e.py +++ b/tests/unit/orchestrator/test_qwen3_vl_e2e.py @@ -27,9 +27,8 @@ def test_generate_qwen3_vl_e2e_preserves_inline_images_and_expanded_prompt_ids() from transformers import AutoProcessor from verifiers.v1.dialects import ChatDialect from verifiers.v1.graph import prepare_turn - from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest - from prime_rl.inference.vllm.serving_chat import PrimeRlServingChat + from prime_rl.inference.vllm.serving_chat import PrimeRlChatRequest, PrimeRlServingChat from prime_rl.orchestrator.trajectories import trace_to_samples tokenizer = load_tokenizer(_MODEL) @@ -48,18 +47,19 @@ def test_generate_qwen3_vl_e2e_preserves_inline_images_and_expanded_prompt_ids() ], } ] - request = ChatCompletionRequest(model=_MODEL, messages=messages, return_token_ids=True) + request = PrimeRlChatRequest(model=_MODEL, messages=messages, return_training_metadata=True) serving = object.__new__(PrimeRlServingChat) serving.model_config = SimpleNamespace(tokenizer=_MODEL, max_model_len=1_000_000) + serving.online_renderer = SimpleNamespace(exclude_tools_when_tool_choice_none=False) serving.chat_template = None serving.chat_template_content_format = "auto" serving.default_chat_template_kwargs = {} - serving.training_renderer_lock = Lock() - serving.training_tokenizer = tokenizer - serving.training_renderer = renderer - serving.training_template_kwargs = {} - serving.training_renderer_config = Qwen3VLRendererConfig() - _, (engine_input,) = serving._render_training_prompt(request) + serving.prompt_renderer_lock = Lock() + serving.prompt_tokenizer = tokenizer + serving.prompt_renderer = renderer + serving.prompt_template_kwargs = {} + serving.prompt_renderer_config = Qwen3VLRendererConfig() + _, (engine_input,) = serving._render_prompt(request) assert engine_input["type"] == "multimodal" assert len(engine_input["mm_hashes"]["image"]) == 1 @@ -95,7 +95,6 @@ def test_generate_qwen3_vl_e2e_preserves_inline_images_and_expanded_prompt_ids() ) assert response.tokens is not None assert "multi_modal_data" not in request._prime_training_metadata - assert response.tokens.mm_placeholders == [(placeholder.offset, placeholder.length)] assert len(response.tokens.is_content) == len(engine_input["prompt_token_ids"]) trace = vf.Trace( agent=vf.AgentInfo(config=vf.AgentConfig()), From 4b65792e21d32fd3eda68b0d58d2902bdfff0579 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:34:18 -0700 Subject: [PATCH 10/13] fix: pin companion revisions for standalone installs --- packages/prime-rl-configs/pyproject.toml | 2 +- pyproject.toml | 5 +++-- uv.lock | 7 ++++--- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/packages/prime-rl-configs/pyproject.toml b/packages/prime-rl-configs/pyproject.toml index 45eb11793ac..9fd0393e39b 100644 --- a/packages/prime-rl-configs/pyproject.toml +++ b/packages/prime-rl-configs/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "renderers>=0.1.11", "tomli>=2.2.1", "tomli-w>=1.2.0", - "verifiers>=0.3.2.dev158", + "verifiers @ git+https://github.com/PrimeIntellect-ai/verifiers.git@1092382", ] [build-system] diff --git a/pyproject.toml b/pyproject.toml index cc962dde638..791b1dc52c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "setproctitle>=1.3.0", "uvloop>=0.21.0", "verifiers[harbor]>=0.4.0", - "renderers[multimodal]>=0.1.13.dev6", + "renderers[multimodal] @ git+https://github.com/PrimeIntellect-ai/renderers.git@a3c71b3", "pybase64>=1.4.2", "msgspec>=0.18", "fastapi>=0.115", @@ -173,6 +173,7 @@ members = [ "deps/verifiers/environments/reverse_text", "deps/verifiers/environments/wiki_search", "deps/verifiers/environments/wordle", + "packages/prime-rl-configs", ] # Tau tasksets require litellm<1.82.7, while Harbor requires litellm>=1.92.0. # Apex pins a verifiers git revision instead of the editable workspace member. @@ -261,7 +262,6 @@ prime-kernels = false flash-mla = false [tool.uv.sources] -prime-rl-configs = { path = "packages/prime-rl-configs", editable = true } # prime-rl consumes `verifiers.v1` (incl. the built-in harnesses/tasksets); env tasksets # live under deps/{verifiers,prime-envs}/environments. verifiers = { path = "deps/verifiers", editable = true } @@ -304,6 +304,7 @@ flash-mla = { url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/downloa flash-linear-attention = { git = "https://github.com/fla-org/flash-linear-attention" } mamba-ssm = { git = "https://github.com/state-spaces/mamba", rev = "f577286" } flash_attn_3 = { index = "pytorch-cu130-test" } +prime-rl-configs = { workspace = true, editable = true } [tool.uv.extra-build-dependencies] flash-attn = [{ requirement = "torch", match-runtime = true }] diff --git a/uv.lock b/uv.lock index d97e33cc599..17d9faed4ac 100644 --- a/uv.lock +++ b/uv.lock @@ -115,6 +115,7 @@ members = [ "pi-terminalbench2-1", "pinchbench", "prime-rl", + "prime-rl-configs", "prolog", "proofnet", "proposer-solver", @@ -5353,12 +5354,12 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "prime-pydantic-config", specifier = ">=0.4.3" }, + { name = "prime-pydantic-config", editable = "deps/pydantic-config" }, { name = "pydantic", specifier = ">=1.10.13" }, - { name = "renderers", specifier = ">=0.1.11" }, + { name = "renderers", editable = "deps/renderers" }, { name = "tomli", specifier = ">=2.2.1" }, { name = "tomli-w", specifier = ">=1.2.0" }, - { name = "verifiers", specifier = ">=0.3.2.dev158" }, + { name = "verifiers", editable = "deps/verifiers" }, ] [[package]] From a954a229dccd602f9c3f1f56a9a633f3f298040f Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:41:29 -0700 Subject: [PATCH 11/13] fix: align scoring test and ablation renderer config --- .../advanced/glm-4.5-air/ipo-ablation/base.toml | 8 ++++---- .../unit/orchestrator/test_prefill_logprobs.py | 17 +++-------------- 2 files changed, 7 insertions(+), 18 deletions(-) diff --git a/examples/advanced/glm-4.5-air/ipo-ablation/base.toml b/examples/advanced/glm-4.5-air/ipo-ablation/base.toml index 40e3c22e641..c94d151dcfd 100644 --- a/examples/advanced/glm-4.5-air/ipo-ablation/base.toml +++ b/examples/advanced/glm-4.5-air/ipo-ablation/base.toml @@ -66,10 +66,6 @@ num_output_tokens_weight = 0.1 num_input_tokens_weight = 0.1 num_turns_weight = 0.1 -[orchestrator.renderer] -name = "glm-4.5" -thinking_retention = "all" - [[orchestrator.train.source]] name = "scaleswe" env.taskset.id = "scaleswe" @@ -89,6 +85,10 @@ env.agent.timeout.rollout = 3600 [inference] +[inference.renderer] +name = "glm-4.5" +thinking_retention = "all" + [inference.vllm] gpu_memory_utilization = 0.85 max_model_len = 131072 diff --git a/tests/unit/orchestrator/test_prefill_logprobs.py b/tests/unit/orchestrator/test_prefill_logprobs.py index 293d621989d..bdab36d46f3 100644 --- a/tests/unit/orchestrator/test_prefill_logprobs.py +++ b/tests/unit/orchestrator/test_prefill_logprobs.py @@ -1,16 +1,10 @@ import asyncio -import json - -import httpx from prime_rl.orchestrator.clients import prefill_logprobs class _FakeOpenAIClient: - """Stand-in for ``AsyncOpenAI`` that captures the sole ``.post()`` call and - returns a synthesized ``httpx.Response`` so ``cast_to=httpx.Response`` is - handed back verbatim, mirroring the real SDK's short-circuit at - ``AsyncAPIClient._process_response``.""" + """Capture the scoring request and return its decoded JSON response.""" def __init__(self, payload: dict): # Match what AsyncOpenAI exposes — prefill_logprobs reads ``str(openai.base_url)``. @@ -20,12 +14,7 @@ def __init__(self, payload: dict): async def post(self, url, *, cast_to, body): self.calls.append({"url": url, "cast_to": cast_to, "body": body}) - request = httpx.Request("POST", url, json=body) - return httpx.Response( - status_code=200, - content=json.dumps(self._payload).encode(), - request=request, - ) + return self._payload def test_prefill_logprobs_uses_inference_generate(): @@ -45,7 +34,7 @@ async def _run(): assert fake_openai.calls == [ { "url": "http://fake-host:8000/inference/v1/generate", - "cast_to": httpx.Response, + "cast_to": dict, "body": { "model": "ref-model", "token_ids": [1, 2, 3], From 317e6f3e4333455afad304c0fa62fff4d77abae5 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:02:20 -0700 Subject: [PATCH 12/13] fix: configure native reasoning parser and pin review fixes --- configs/ci/nightly-fft/hendrycks-sanity.toml | 3 ++- deps/verifiers | 2 +- examples/basic/hendrycks-sanity/rl.toml | 3 ++- packages/prime-rl-configs/pyproject.toml | 2 +- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/configs/ci/nightly-fft/hendrycks-sanity.toml b/configs/ci/nightly-fft/hendrycks-sanity.toml index d1c2a0ee409..98a1acbab4b 100644 --- a/configs/ci/nightly-fft/hendrycks-sanity.toml +++ b/configs/ci/nightly-fft/hendrycks-sanity.toml @@ -50,6 +50,7 @@ attn = "flash_attention_2" [inference.vllm] max_model_len = 8192 data_parallel_size = 4 +reasoning_parser = "qwen3" + [inference.renderer] name = "default" -reasoning_parser = "think" diff --git a/deps/verifiers b/deps/verifiers index 1092382a904..385cf7e7f9b 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 1092382a904f2d7b9db98ede289a658a2466b147 +Subproject commit 385cf7e7f9b924e5de82aac631a9154ffed2b52b diff --git a/examples/basic/hendrycks-sanity/rl.toml b/examples/basic/hendrycks-sanity/rl.toml index ec3e4b95f39..3537789ba9a 100644 --- a/examples/basic/hendrycks-sanity/rl.toml +++ b/examples/basic/hendrycks-sanity/rl.toml @@ -41,6 +41,7 @@ seq_len = 16384 [inference.vllm] max_model_len = 8192 +reasoning_parser = "qwen3" + [inference.renderer] name = "default" -reasoning_parser = "think" diff --git a/packages/prime-rl-configs/pyproject.toml b/packages/prime-rl-configs/pyproject.toml index 9fd0393e39b..00b6fc58d0d 100644 --- a/packages/prime-rl-configs/pyproject.toml +++ b/packages/prime-rl-configs/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "renderers>=0.1.11", "tomli>=2.2.1", "tomli-w>=1.2.0", - "verifiers @ git+https://github.com/PrimeIntellect-ai/verifiers.git@1092382", + "verifiers @ git+https://github.com/PrimeIntellect-ai/verifiers.git@385cf7e", ] [build-system] From 5d1f3f0405b93a974b0c2f08d0948f52f0fde557 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:59:26 -0700 Subject: [PATCH 13/13] fix: reject unsupported prompt truncation and pin rebased Verifiers --- deps/verifiers | 2 +- packages/prime-rl-configs/pyproject.toml | 2 +- src/prime_rl/inference/vllm/serving_chat.py | 2 ++ 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 385cf7e7f9b..37bac61df65 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 385cf7e7f9b924e5de82aac631a9154ffed2b52b +Subproject commit 37bac61df6550e3a19026e0b7b4313634b662d59 diff --git a/packages/prime-rl-configs/pyproject.toml b/packages/prime-rl-configs/pyproject.toml index 00b6fc58d0d..e8d699f8721 100644 --- a/packages/prime-rl-configs/pyproject.toml +++ b/packages/prime-rl-configs/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "renderers>=0.1.11", "tomli>=2.2.1", "tomli-w>=1.2.0", - "verifiers @ git+https://github.com/PrimeIntellect-ai/verifiers.git@385cf7e", + "verifiers @ git+https://github.com/PrimeIntellect-ai/verifiers.git@37bac61", ] [build-system] diff --git a/src/prime_rl/inference/vllm/serving_chat.py b/src/prime_rl/inference/vllm/serving_chat.py index f3e1e17efb3..4359d49737f 100644 --- a/src/prime_rl/inference/vllm/serving_chat.py +++ b/src/prime_rl/inference/vllm/serving_chat.py @@ -63,6 +63,8 @@ async def create_chat_completion(self, request, raw_request=None): return await super().create_chat_completion(request, raw_request) async def render_chat_request(self, request: ChatCompletionRequest): + if request.truncate_prompt_tokens is not None: + return self.create_error_response("Prompt truncation is not supported by the endpoint renderer") if request.continue_final_message or request.chat_template is not None: raise ValueError("Chat requests require complete messages and the server's configured renderer") error = await self._check_model(request)