diff --git a/nemo_gym/responses_converter.py b/nemo_gym/responses_converter.py index bfd96fc53c..c9b854a8e6 100644 --- a/nemo_gym/responses_converter.py +++ b/nemo_gym/responses_converter.py @@ -78,6 +78,19 @@ class ResponsesConverterState(BaseModel): token_information: Optional[TokenIDLogProbMixin] = None def flush_assistant(self) -> None: + # `token_information` describes exactly one assistant turn, so consume it here rather than + # letting it persist. It is cleared on EVERY path, including the empty-buffer early return + # (an assistant item can carry ids while producing no content and no tool calls, e.g. a + # zero-token generation, and that message is dropped below). + # + # Without this, a later assistant message that carries no ids of its own silently inherits + # this turn's -- attributing one turn's generated tokens to another turn's text in the + # training data. Harnesses hit this whenever they inject or rewrite an assistant message + # without token ids (canned turns, compacted history, agents that stamp ids onto only the + # most recent assistant message). + token_information = self.token_information + self.token_information = None + if not (self.content_buffer or self.tool_calls_buffer): return @@ -87,10 +100,10 @@ def flush_assistant(self) -> None: tool_calls=self.tool_calls_buffer, ) - if self.return_token_id_information and self.token_information: + if self.return_token_id_information and token_information: message = NeMoGymChatCompletionAssistantMessageForTrainingParam( **shared_params, - **self.token_information.model_dump(exclude_none=True), + **token_information.model_dump(exclude_none=True), ) else: message = NeMoGymChatCompletionAssistantMessageParam(**shared_params) diff --git a/responses_api_models/sglang_model/README.md b/responses_api_models/sglang_model/README.md new file mode 100644 index 0000000000..303dc767cc --- /dev/null +++ b/responses_api_models/sglang_model/README.md @@ -0,0 +1,86 @@ +# Description + +A Responses-API **model server** for policies served by [SGLang](https://github.com/sgl-project/sglang). + +RL trainers (GRPO etc.) need the *exact* token ids the policy emitted and their logprobs, not a +re-tokenized decode of the text. `vllm_model` recovers those by parsing `token_id:NNN` logprob +tokens out of `/v1/chat/completions` — a vLLM-specific encoding SGLang does not produce. This +server closes that gap. + +It subclasses `vllm_model`'s `VLLMModel`. Everything else — Responses<->ChatCompletions +conversion, `responses()`, and the assistant-message training-class upgrade — is inherited +unchanged. + +## Transports + +### `transport: chat` (default) — requires **sglang >= 0.5.13** + +Drives SGLang's OpenAI-compatible `/v1/chat/completions`, requesting the training metadata via +SGLang's native `return_meta_info` / `return_prompt_token_ids` extensions. Token ids and logprobs +come back on each choice as `meta_info.output_token_logprobs` and `prompt_token_ids`. + +These extensions landed with the sglang-miles TITO sync series and are in the tree as of the +0.5.13 release — **no patched build or fork is required**. (ProRL ships a +`patch_sglang_0513_token_metadata.sh`, but that patch only makes `logprobs=true` *imply* those +flags for clients that don't set them; this server sets them explicitly, so the patch is +unnecessary here.) + +This is the path to use whenever you can. Only the token extraction is overridden, so prompt +templating, **tool-call parsing**, sampling parameters, auth, and context-overflow handling are +all done server-side by SGLang — there is no local tokenizer that can drift from the server's. + +### `transport: generate` — fallback for older builds + +Drives SGLang's native `/generate` with `return_logprob=true`, tokenizing the prompt locally. +Use only for SGLang builds/forks predating chat-side TITO. Inherent limitations of this path: + +- **Tool calls are not parsed** out of the generated text (tool *schemas* are rendered into the + prompt only if the local chat template does so). +- The prompt is tokenized by a **local** copy of the chat template, so `model` must be the exact + path/revision the SGLang server was launched with — otherwise generation is silently + conditioned on ids from the wrong template. +- Sampling params `/generate` does not accept (`n`, `seed`, `response_format`, `logit_bias`) are + logged and dropped rather than honored. +- A prompt that overflows `context_length` yields an empty completion with + `finish_reason="length"` (matching the inherited vLLM behavior) so it is filterable, rather + than a head-truncated prompt that would generate from a malformed context. + +In both transports, a generation SGLang reports as `finish_reason="abort"` raises instead of +being emitted, so a server-cancelled partial rollout cannot enter a training batch looking like +a normal completion. + +## On-policy scope + +Single-turn rollouts are exactly on-policy: the ids attached to the assistant message are the +ids the policy emitted. + +**Multi-turn is not yet on-policy.** Turn N re-renders the whole history through the chat +template, which re-tokenizes prior assistant spans; those generally differ from the +`generation_token_ids` the policy actually produced. The fix is to splice the carried-forward +token ids and tokenize only the newly inserted environment/user messages — tracked as follow-up +work, not yet implemented here. + +## Config + +- `transport`: `chat` (default) or `generate`. +- `base_url`: **must end in `/v1` for `transport: chat`** (like `vllm_model`), and must be the + **bare** server URL for `transport: generate` (the server appends `/generate`). Migrating a + bare URL to the chat transport 404s. +- `context_length` (generate only): SGLang context window; keep in sync with the trainer's max + sequence length. +- `default_max_new_tokens` (generate only): used only when a request carries no + `max_(completion_)tokens`. +- `add_generation_prompt`, `trust_remote_code` (generate only): forwarded to the local + tokenizer / chat template. `trust_remote_code` defaults to `false`. + +See `configs/sglang_model_for_training.yaml` and +`configs/sglang_model_for_training_generate.yaml`. + +## Licensing information + +Code: Apache 2.0 +Data: N/A + +Dependencies +- nemo_gym: Apache 2.0 +- transformers: Apache 2.0 diff --git a/responses_api_models/sglang_model/__init__.py b/responses_api_models/sglang_model/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/responses_api_models/sglang_model/_logic.py b/responses_api_models/sglang_model/_logic.py new file mode 100644 index 0000000000..2869f0df4b --- /dev/null +++ b/responses_api_models/sglang_model/_logic.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Pure, framework-free logic for the sglang_model Gym server. + +Split out from app.py so it can be unit-tested without importing nemo_gym / the server +framework (which is only available inside the per-server venv). No third-party imports. +""" + +from typing import Any, Dict, List, Tuple + + +def extract_generated_tokens_and_logprobs( + result: Dict[str, Any], +) -> Tuple[List[int], List[float]]: + """Parse an SGLang `meta_info` carrier into (generated_token_ids, logprobs). + + Works for both transports, because SGLang emits the same `meta_info` shape either way: + the `/generate` response body, and (since 0.5.13, with `return_meta_info=true`) each + `/v1/chat/completions` *choice*. + + Mirrors nemo_rl's sglang_worker: handles the dict-form and tuple-form + `meta_info.output_token_logprobs`, plus the `output_token_logprobs_val/idx` fallback. + Raises RuntimeError on malformed / missing logprobs. + """ + meta = result.get("meta_info", {}) or {} + otl = meta.get("output_token_logprobs", []) + if otl: + toks: List[int] = [] + lps: List[float] = [] + for item in otl: + if isinstance(item, dict): + tid = item.get("token_id", item.get("id")) + lp = item.get("logprob") + else: # [logprob, token_id, (optional) text] + lp, tid = item[0], item[1] + if tid is None or lp is None: + raise RuntimeError(f"Malformed SGLang output_token_logprobs entry: {item!r}") + toks.append(int(tid)) + lps.append(float(lp)) + return toks, lps + + val = meta.get("output_token_logprobs_val", result.get("output_token_logprobs_val", [])) + idx = meta.get("output_token_logprobs_idx", result.get("output_token_logprobs_idx", [])) + ids = result.get("output_ids", meta.get("output_ids", [])) + if val: + new = idx if idx else ids + if len(new) != len(val): + raise RuntimeError(f"SGLang mismatched gen logprob fields: {len(new)} ids vs {len(val)} logprobs") + return [int(x) for x in new], [float(x) for x in val] + + raise RuntimeError( + f"SGLang /generate returned no generation logprobs (keys={sorted(result)}, meta={sorted(meta)}). " + "Ensure the request set return_logprob=true." + ) + + +def normalize_token_ids(rendered: Any) -> List[int]: + """Normalize a chat-template tokenization result to a flat list[int]. + + transformers 5.x `apply_chat_template(tokenize=True)` can return a dict / BatchEncoding + (in which case `list(...)` would grab the KEYS), or a nested `[[...]]` for a single + conversation. This collapses all of those to a flat list of python ints. + """ + if isinstance(rendered, dict) or hasattr(rendered, "input_ids"): + rendered = rendered["input_ids"] + rendered = list(rendered) + if rendered and isinstance(rendered[0], (list, tuple)): + rendered = list(rendered[0]) + return [int(t) for t in rendered] + + +# OpenAI chat-completion param -> SGLang sampling_params key, for params that pass through +# unchanged when set. `n` is deliberately absent: /generate would return a single choice +# regardless, so a caller asking for n>1 must be told rather than silently under-served. +_PASSTHROUGH_SAMPLING_PARAMS = { + "top_k": "top_k", + "stop": "stop", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", + "repetition_penalty": "repetition_penalty", + "min_p": "min_p", +} + +# Params this transport cannot honor. Reported so a caller never silently gets a different +# sampling distribution than the one their recipe configured. +_UNSUPPORTED_SAMPLING_PARAMS = ("n", "seed", "response_format", "logit_bias") + + +def build_sampling_params(body_dict: Dict[str, Any], default_max_new_tokens: int) -> Dict[str, Any]: + """Map OpenAI chat-completion params to SGLang sampling_params (``/generate`` transport). + + Only used by the legacy ``/generate`` transport; the ``chat`` transport forwards the + request to SGLang's OpenAI-compatible endpoint verbatim, so it needs no mapping. + """ + max_new = body_dict.get("max_completion_tokens") + if max_new is None: + max_new = body_dict.get("max_tokens") + if max_new is None: + max_new = default_max_new_tokens + sp: Dict[str, Any] = { + "temperature": body_dict.get("temperature", 1.0), + "top_p": body_dict.get("top_p", 1.0), + "max_new_tokens": int(max_new), + } + for openai_key, sglang_key in _PASSTHROUGH_SAMPLING_PARAMS.items(): + value = body_dict.get(openai_key) + if value is not None and value != []: + sp[sglang_key] = value + return sp + + +def unsupported_sampling_params(body_dict: Dict[str, Any]) -> List[str]: + """Names of request params the ``/generate`` transport cannot honor, for a loud warning. + + Silently ignoring these would make the realized rollout distribution diverge from the + configured recipe -- which is invisible in the training data. + """ + return [key for key in _UNSUPPORTED_SAMPLING_PARAMS if body_dict.get(key) is not None] + + +def cap_to_context( + prompt_token_ids: List[int], sampling_params: Dict[str, Any], ctx: int +) -> Tuple[List[int], Dict[str, Any]]: + """Keep the request within the context window: guarantee input_len + max_new_tokens < ctx + while always leaving room for at least one generated token (SGLang /generate errors when a + request exceeds the context). If the prompt alone is too long it is truncated. Mirrors + nemo_rl's SGLang worker. Returns the (possibly truncated) ids and (possibly adjusted) params; + does not mutate the inputs. + """ + if not ctx: + return prompt_token_ids, sampling_params + if ctx < 2: + # Below this there is no room for both a prompt and a generated token. Raising beats + # silently POSTing a negative `max_new_tokens`, which SGLang rejects with an opaque 400. + raise ValueError(f"context_length={ctx} is too small; need >= 2 to leave room for one generated token") + # Cap the prompt at ctx-2 so that input + (>=1 generated token) <= ctx-1 < ctx. (A ctx-1 + # truncation combined with a max(1, ...) floor on `room` could yield input+gen == ctx, which + # violates the bound and can overflow the context.) + max_prompt_len = ctx - 2 + truncated = len(prompt_token_ids) > max_prompt_len + if truncated: + prompt_token_ids = prompt_token_ids[:max_prompt_len] + room = ctx - 1 - len(prompt_token_ids) # >= 1 (since len <= ctx-2); input + room == ctx-1 < ctx + if sampling_params["max_new_tokens"] > room: + sampling_params = {**sampling_params, "max_new_tokens": room} + return prompt_token_ids, sampling_params + + +def would_truncate(prompt_token_ids: List[int], ctx: int) -> bool: + """Whether `cap_to_context` would drop prompt tokens. + + Truncation keeps the prompt *head*, so it discards the newest user turn and the + `add_generation_prompt` cue -- the model then generates from a malformed context. Callers + use this to surface the rollout as degenerate instead of letting it enter the training set + looking valid. + """ + return bool(ctx) and ctx >= 2 and len(prompt_token_ids) > ctx - 2 diff --git a/responses_api_models/sglang_model/app.py b/responses_api_models/sglang_model/app.py new file mode 100644 index 0000000000..6840372e8f --- /dev/null +++ b/responses_api_models/sglang_model/app.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""NeMo-Gym Responses-API model server for SGLang-served policies. + +Why this exists +--------------- +RL trainers (e.g. GRPO) need the *exact* token ids the policy emitted and their logprobs. +`vllm_model` recovers those by parsing `token_id:NNN` logprob tokens out of +`/v1/chat/completions`, which is a vLLM-specific encoding that SGLang does not produce. + +Two transports +-------------- +``transport: chat`` (default, **requires sglang >= 0.5.13**) + Drives SGLang's OpenAI-compatible ``/v1/chat/completions`` and asks for the training + metadata via the native ``return_meta_info`` / ``return_prompt_token_ids`` request + extensions (added by the TITO series, in the tree as of 0.5.13). Ids and logprobs come + back on each choice as ``meta_info.output_token_logprobs`` / ``prompt_token_ids``. + + This is the preferred path: everything except the token extraction is the inherited + ``VLLMModel`` behavior, so tool-call parsing, prompt templating, sampling params, auth + and context-overflow handling are all done *server-side* by SGLang -- there is no local + tokenizer to drift from the server's, and no client-side reimplementation to keep in sync. + +``transport: generate`` (fallback) + Drives SGLang's native ``/generate`` with ``return_logprob=true``, tokenizing the prompt + locally. Needed only for SGLang builds/forks predating chat-side TITO. It cannot parse + tool calls and it tokenizes with a *local* copy of the chat template, so it is off by + construction if that copy differs from the server's. Prefer ``chat`` whenever possible. + +The pure request/response transforms live in `_logic.py` (unit-tested in tests/). +""" + +from time import time +from typing import Any, Dict, List, Literal, Optional +from uuid import uuid4 + +from fastapi import Body, Request + +from nemo_gym.openai_utils import ( + NeMoGymAsyncOpenAI, + NeMoGymChatCompletion, + NeMoGymChatCompletionCreateParamsNonStreaming, +) +from nemo_gym.server_utils import ( + SESSION_ID_KEY, + get_response_json, + is_nemo_gym_fastapi_entrypoint, +) +from nemo_gym.server_utils import ( + request as ng_request, +) +from responses_api_models.sglang_model._logic import ( + build_sampling_params, + cap_to_context, + extract_generated_tokens_and_logprobs, + normalize_token_ids, + unsupported_sampling_params, + would_truncate, +) +from responses_api_models.vllm_model.app import VLLMModel, VLLMModelConfig + + +# Chat-side TITO (`return_meta_info` / `return_prompt_token_ids` on /v1/chat/completions) +# landed in the sglang-miles sync series and is present in the 0.5.13 release tree. +MIN_SGLANG_VERSION_FOR_CHAT_TRANSPORT = "0.5.13" + + +class SGLangModelConfig(VLLMModelConfig): + # `chat` requires sglang >= 0.5.13; `generate` is the fallback for older builds/forks. + transport: Literal["chat", "generate"] = "chat" + + # --- `generate` transport only (ignored when transport == "chat") --- + # Used only when the request carries no max_(completion_)tokens. + default_max_new_tokens: int = 1024 + # Opt-in, matching the rest of the repo: executing model-repo code is the caller's decision. + trust_remote_code: bool = False + add_generation_prompt: bool = True + # SGLang context window. Keep in sync with policy.max_total_sequence_length: + # the prompt is truncated and max_new_tokens shrunk so input_len + max_new_tokens < ctx, + # else SGLang /generate returns 400. + context_length: int = 4096 + + +class SGLangModel(VLLMModel): + config: SGLangModelConfig + + def _post_init(self) -> None: + super()._post_init() + if self.config.transport != "generate": + return + # Local tokenization is only needed by the /generate fallback. Imported lazily so the + # default `chat` transport does not pay for (or depend on) transformers at startup. + from transformers import AutoTokenizer + + # The model name must be the exact path/revision the SGLang server was launched with; + # any divergence silently conditions generation on ids from the wrong template. + self._tokenizer = AutoTokenizer.from_pretrained( + self.config.model, trust_remote_code=self.config.trust_remote_code + ) + # Bare SGLang server base url(s); we hit `{base}/generate`, not `/v1/...`. + self._sglang_urls: List[str] = [u.rstrip("/") for u in self.config.base_url] + + # ------------------------------------------------------------------ + # `chat` transport: inherit everything, override only token extraction + # ------------------------------------------------------------------ + + def _preprocess_chat_completion_create_params(self, request: Request, body_dict: Dict[str, Any]) -> Dict[str, Any]: + body_dict = super()._preprocess_chat_completion_create_params(request, body_dict) + if self.config.transport == "chat" and self.config.return_token_id_information: + # vLLM-only knob: SGLang has no `token_id:NNN` logprob encoding, and leaving it set + # would be a silent no-op that misrepresents where the ids come from. + body_dict.pop("return_tokens_as_token_ids", None) + # SGLang's native request extensions for the training metadata. + body_dict["return_meta_info"] = True + body_dict["return_prompt_token_ids"] = True + return body_dict + + async def _attach_token_id_information( + self, choice_dict: Dict[str, Any], body_dict: Dict[str, Any], client: NeMoGymAsyncOpenAI + ) -> None: + """Read the ids/logprobs SGLang already returned, instead of vLLM's `token_id:` parse. + + No `/tokenize` round-trip is needed: `return_prompt_token_ids` makes SGLang report the + prompt ids it actually tokenized, which is authoritative in a way a local tokenizer + cannot be. + """ + # An aborted generation is a truncated fragment, not a completion. Letting it through + # would put a poisoned rollout into the GRPO batch looking like a normal `stop`. + if choice_dict.get("finish_reason") == "abort": + raise RuntimeError( + f"`{self.config.name}`: SGLang reported finish_reason='abort' (generation was " + "cancelled server-side, e.g. preemption or a shutdown). Refusing to emit a " + "partial rollout as a completed one." + ) + + generation_token_ids, generation_log_probs = extract_generated_tokens_and_logprobs(choice_dict) + + prompt_token_ids: Optional[List[int]] = choice_dict.get("prompt_token_ids") + if prompt_token_ids is None: + raise RuntimeError( + f"`{self.config.name}` requested prompt token ids from SGLang " + "(return_token_id_information=True, so return_prompt_token_ids=True was sent), " + f"but the response carried none (choice keys={sorted(choice_dict)}). This server " + f"requires sglang >= {MIN_SGLANG_VERSION_FOR_CHAT_TRANSPORT} for " + "transport='chat'; set transport='generate' for older builds." + ) + + choice_dict["message"].update( + dict( + prompt_token_ids=[int(t) for t in prompt_token_ids], + generation_token_ids=generation_token_ids, + generation_log_probs=generation_log_probs, + ) + ) + + # Clean the duplicated / non-OpenAI information so the response validates. + choice_dict.pop("logprobs", None) + choice_dict.pop("prompt_token_ids", None) + choice_dict.pop("meta_info", None) + + # ------------------------------------------------------------------ + # `generate` transport: fallback for builds without chat-side TITO + # ------------------------------------------------------------------ + + async def chat_completions( + self, request: Request, body: NeMoGymChatCompletionCreateParamsNonStreaming = Body() + ) -> NeMoGymChatCompletion: + if self.config.transport == "chat": + return await super().chat_completions(request, body) + return await self._chat_completions_via_generate(request, body) + + async def _chat_completions_via_generate( + self, request: Request, body: NeMoGymChatCompletionCreateParamsNonStreaming + ) -> NeMoGymChatCompletion: + """Drive SGLang's native /generate, tokenizing the prompt locally. + + Only for SGLang builds without chat-side TITO. Known gaps vs. the `chat` transport: + tool calls are not parsed (and tool schemas are rendered only if the local chat + template does so), and the prompt is tokenized locally rather than by the server. + """ + body_dict = body.model_dump(exclude_unset=True) + body_dict = self._preprocess_chat_completion_create_params(request, body_dict) + messages = body_dict["messages"] + + ignored = unsupported_sampling_params(body_dict) + if ignored: + print( + f"[sglang_model] transport='generate' cannot honor {ignored}; the realized " + "sampling distribution will differ from the configured recipe. " + "Use transport='chat' (sglang >= 0.5.13) for full parameter support.", + flush=True, + ) + + # 1) prompt token ids via the model's own chat template (local tokenizer). Use the + # *merged* kwargs the inherited preprocess produced, so per-request overrides in + # `metadata.chat_template_kwargs` (e.g. per-sample reasoning on/off) are honored -- + # reading self.config here would tokenize with the wrong template. + ct_kwargs = body_dict.get("chat_template_kwargs") or {} + rendered = self._tokenizer.apply_chat_template( + messages, + add_generation_prompt=self.config.add_generation_prompt, + tokenize=True, + return_dict=False, + # Let the template render tool schemas into the prompt when the caller sent them. + tools=body_dict.get("tools"), + **ct_kwargs, + ) + prompt_token_ids = normalize_token_ids(rendered) + + # 2) cap to the context window, then generate via SGLang native /generate. + sampling_params = build_sampling_params(body_dict, self.config.default_max_new_tokens) + truncated = would_truncate(prompt_token_ids, self.config.context_length) + prompt_token_ids, sampling_params = cap_to_context( + prompt_token_ids, sampling_params, self.config.context_length + ) + if truncated: + # Truncation keeps the prompt head, dropping the newest turn and the generation + # cue. Match the inherited overflow behavior instead of emitting a rollout that + # looks valid: an empty completion with finish_reason="length" is filterable. + print( + f"[sglang_model] prompt exceeded context_length={self.config.context_length}; " + "returning an empty completion with finish_reason='length'.", + flush=True, + ) + res = self._create_empty_chat_completion() + res.choices[0].finish_reason = "length" + return res + + payload = { + "input_ids": prompt_token_ids, + "sampling_params": sampling_params, + "return_logprob": True, + "logprob_start_len": -1, + } + sid = request.session.get(SESSION_ID_KEY, "") if hasattr(request, "session") else "" + url = f"{self._sglang_urls[hash(sid) % len(self._sglang_urls)]}/generate" + # SGLang applies its --api-key middleware to /generate as well as /v1/*, so an + # authenticated deployment 401s without this. Only set when non-empty: `request()` + # does `kwargs.setdefault("headers", ...)`, which would keep an explicit None. + extra_request_kwargs: Dict[str, Any] = {} + if self.config.api_key: + extra_request_kwargs["headers"] = {"Authorization": f"Bearer {self.config.api_key}"} + # Use NeMo-Gym's pooled aiohttp client. Raw aiohttp + native raise_for_status + # trips the framework's exception_handling_middleware (it requires the escaping + # exception to carry `response_content`). + resp = await ng_request("POST", url, json=payload, **extra_request_kwargs) + if not resp.ok: + content = await resp.read() + print( + f"[sglang_model] SGLang /generate -> {resp.status}; " + f"input_len={len(prompt_token_ids)} max_new_tokens={sampling_params['max_new_tokens']}; " + f"body={content[:800]!r}", + flush=True, + ) + try: + resp.raise_for_status() + except Exception as e: + e.response_content = content # satisfy nemo_gym exception middleware + raise + result = await get_response_json(resp) + + meta = result.get("meta_info", {}) or {} + finish = meta.get("finish_reason") + if isinstance(finish, dict): + finish = finish.get("type") + if finish == "abort": + raise RuntimeError( + f"`{self.config.name}`: SGLang reported finish_reason='abort' (generation was " + "cancelled server-side). Refusing to emit a partial rollout as a completed one." + ) + finish_reason = "length" if finish == "length" else "stop" + + gen_token_ids, gen_log_probs = extract_generated_tokens_and_logprobs(result) + # generation_token_ids stay RAW (incl. EOS/special — the policy generated them and + # we train on them), but the assistant *content* the verifier grades must be clean, + # matching vLLM's server-side decode. A trailing special token otherwise breaks + # strict parsers (e.g. structured_outputs json.loads). + gen_text = self._tokenizer.decode(gen_token_ids, skip_special_tokens=True) + + # 3) OpenAI chat.completion dict + the training token fields on the assistant message. + chat_completion_dict: Dict[str, Any] = { + "id": f"chtcmpl-{uuid4().hex}", + "object": "chat.completion", + "created": int(time()), + "model": self.config.model, + "choices": [ + { + "index": 0, + "finish_reason": finish_reason, + "message": {"role": "assistant", "content": gen_text}, + } + ], + "usage": { + "prompt_tokens": len(prompt_token_ids), + "completion_tokens": len(gen_token_ids), + "total_tokens": len(prompt_token_ids) + len(gen_token_ids), + }, + } + if self.config.return_token_id_information: + chat_completion_dict["choices"][0]["message"].update( + prompt_token_ids=prompt_token_ids, + generation_token_ids=gen_token_ids, + generation_log_probs=gen_log_probs, + ) + return NeMoGymChatCompletion.model_validate(chat_completion_dict) + + +if __name__ == "__main__": + SGLangModel.run_webserver() +elif is_nemo_gym_fastapi_entrypoint(__file__): + app = SGLangModel.run_webserver() # noqa: F401 diff --git a/responses_api_models/sglang_model/configs/sglang_model_for_training.yaml b/responses_api_models/sglang_model/configs/sglang_model_for_training.yaml new file mode 100644 index 0000000000..a2afa04a3e --- /dev/null +++ b/responses_api_models/sglang_model/configs/sglang_model_for_training.yaml @@ -0,0 +1,23 @@ +# SGLang-backed policy model server for RL training. +# Mirrors responses_api_models/vllm_model/configs/vllm_model_for_training.yaml. +# +# transport: chat (the default) REQUIRES sglang >= 0.5.13. It drives SGLang's +# OpenAI-compatible /v1/chat/completions and reads the training token ids and logprobs from +# the native `meta_info` / `prompt_token_ids` response extensions. +# +# IMPORTANT: `base_url` must END IN /v1 here, exactly like vllm_model. This differs from +# transport=generate, where it must be the BARE server url -- pointing a bare url at the +# chat transport 404s. +# +# For SGLang builds/forks that predate chat-side TITO, use +# configs/sglang_model_for_training_generate.yaml instead. +policy_model: + responses_api_models: + sglang_model: + entrypoint: app.py + transport: chat + base_url: ${policy_base_url} + api_key: ${policy_api_key} + model: ${policy_model_name} + return_token_id_information: true + uses_reasoning_parser: false diff --git a/responses_api_models/sglang_model/configs/sglang_model_for_training_generate.yaml b/responses_api_models/sglang_model/configs/sglang_model_for_training_generate.yaml new file mode 100644 index 0000000000..a458932736 --- /dev/null +++ b/responses_api_models/sglang_model/configs/sglang_model_for_training_generate.yaml @@ -0,0 +1,25 @@ +# Fallback policy model server for SGLang builds/forks WITHOUT chat-side TITO +# (i.e. older than 0.5.13). Prefer configs/sglang_model_for_training.yaml when you can. +# +# Drives SGLang's native /generate and tokenizes the prompt locally. Known gaps vs. the +# `chat` transport, all inherent to /generate rather than to this implementation: +# - tool calls are NOT parsed out of the generated text +# - the prompt is tokenized by a LOCAL copy of the chat template, so `model` must be the +# exact path/revision the SGLang server was launched with, or generation is silently +# conditioned on ids from the wrong template +# - sampling params SGLang's /generate does not accept (n, seed, response_format, +# logit_bias) are logged and dropped rather than honored +# +# base_url here is the BARE SGLang server url (NO /v1); the adapter appends /generate. +# Keep context_length in sync with policy.max_total_sequence_length. +policy_model: + responses_api_models: + sglang_model: + entrypoint: app.py + transport: generate + base_url: ${policy_base_url} + api_key: ${policy_api_key} + model: ${policy_model_name} + return_token_id_information: true + uses_reasoning_parser: false + context_length: 4096 diff --git a/responses_api_models/sglang_model/requirements.txt b/responses_api_models/sglang_model/requirements.txt new file mode 100644 index 0000000000..6d4b3e8663 --- /dev/null +++ b/responses_api_models/sglang_model/requirements.txt @@ -0,0 +1,5 @@ +-e nemo-gym[dev] @ ../../ +# Only used by transport="generate" (local prompt tokenization); imported lazily. +# Lower bound: `normalize_token_ids` handles the dict/BatchEncoding return shape that +# `apply_chat_template(tokenize=True)` grew in 4.44. Upper bound: 6.x may change it again. +transformers>=4.44,<6 diff --git a/responses_api_models/sglang_model/tests/__init__.py b/responses_api_models/sglang_model/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/responses_api_models/sglang_model/tests/test_app.py b/responses_api_models/sglang_model/tests/test_app.py new file mode 100644 index 0000000000..8e3de385b2 --- /dev/null +++ b/responses_api_models/sglang_model/tests/test_app.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Orchestration tests for SGLangModel, for both transports. + +Servers are built the same way `vllm_model/tests/test_app.py` builds them -- a real +`SGLangModel` with a mocked `ServerClient` -- so the *inherited* `VLLMModel` behavior is +genuinely exercised rather than stubbed. In particular the real +`_preprocess_chat_completion_create_params` runs, which is what locks in the contract between +it and this subclass. + +All patching goes through pytest's `monkeypatch` fixture so module globals are restored +afterwards and cannot leak into later tests in the same process. + +Pure-logic coverage lives in test_logic.py. +""" + +from typing import Any, Dict +from unittest.mock import MagicMock + +from pytest import MonkeyPatch, raises + +import nemo_gym.server_utils +from nemo_gym.server_utils import ServerClient +from responses_api_models.sglang_model.app import SGLangModel, SGLangModelConfig + + +def _make_server(monkeypatch: MonkeyPatch, **overrides: Any) -> SGLangModel: + config = SGLangModelConfig( + host="0.0.0.0", + port=8081, + base_url=overrides.pop("base_url", "http://sglang-host:30000/v1"), + api_key="dummy_key", # pragma: allowlist secret + model="dummy_model", + entrypoint="", + name="sglang_model", + return_token_id_information=True, + uses_reasoning_parser=False, + **overrides, + ) + get_global_config_dict_mock = MagicMock() + get_global_config_dict_mock.return_value = dict() + monkeypatch.setattr(nemo_gym.server_utils, "get_global_config_dict", get_global_config_dict_mock) + return SGLangModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={})) + + +def _choice(**overrides: Any) -> Dict[str, Any]: + """An SGLang >= 0.5.13 chat choice with the native TITO extensions populated.""" + choice: Dict[str, Any] = { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "the answer"}, + "prompt_token_ids": [1, 2, 3], + "meta_info": {"output_token_logprobs": [[-0.5, 10, "a"], [-0.25, 11, "b"]]}, + "logprobs": {"content": [{"token": "a", "logprob": -0.5}]}, + } + choice.update(overrides) + return choice + + +# --------------------------- chat transport --------------------------- + + +class TestChatTransport: + def test_preprocess_requests_sglang_tito_extensions(self, monkeypatch: MonkeyPatch) -> None: + """The real inherited preprocess runs, then we swap vLLM's knob for SGLang's.""" + server = _make_server(monkeypatch) + body: Dict[str, Any] = {"messages": [{"role": "user", "content": "hi"}]} + + out = server._preprocess_chat_completion_create_params(MagicMock(), body) + + assert out["return_meta_info"] is True + assert out["return_prompt_token_ids"] is True + # vLLM's `token_id:NNN` encoding does not exist on SGLang; leaving it set would be a + # silent no-op that misrepresents where the ids come from. + assert "return_tokens_as_token_ids" not in out + # ...and the inherited behavior is still in force. + assert out["logprobs"] is True + assert out["model"] == "dummy_model" + + def test_preprocess_honors_per_request_chat_template_kwargs(self, monkeypatch: MonkeyPatch) -> None: + """Per-sample overrides must survive; dropping them tokenizes with the wrong template.""" + server = _make_server(monkeypatch, chat_template_kwargs={"enable_thinking": False}) + body: Dict[str, Any] = { + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"chat_template_kwargs": '{"enable_thinking": true}'}, + } + + out = server._preprocess_chat_completion_create_params(MagicMock(), body) + + assert out["chat_template_kwargs"] == {"enable_thinking": True} + + async def test_attach_reads_native_ids_and_logprobs(self, monkeypatch: MonkeyPatch) -> None: + server = _make_server(monkeypatch) + choice = _choice() + + await server._attach_token_id_information(choice, {}, MagicMock()) + + assert choice["message"]["prompt_token_ids"] == [1, 2, 3] + assert choice["message"]["generation_token_ids"] == [10, 11] + assert choice["message"]["generation_log_probs"] == [-0.5, -0.25] + # Non-OpenAI / duplicated fields are stripped so the response validates. + for key in ("logprobs", "prompt_token_ids", "meta_info"): + assert key not in choice + + async def test_attach_rejects_aborted_generation(self, monkeypatch: MonkeyPatch) -> None: + """An abort is a truncated fragment; it must not enter a training batch as a `stop`.""" + server = _make_server(monkeypatch) + + with raises(RuntimeError, match="abort"): + await server._attach_token_id_information(_choice(finish_reason="abort"), {}, MagicMock()) + + async def test_attach_errors_when_server_predates_chat_tito(self, monkeypatch: MonkeyPatch) -> None: + """Older SGLang ignores the extensions; say so instead of emitting empty token ids.""" + server = _make_server(monkeypatch) + choice = _choice() + choice.pop("prompt_token_ids") + + with raises(RuntimeError, match="0.5.13"): + await server._attach_token_id_information(choice, {}, MagicMock()) + + +# ------------------------- generate transport ------------------------- + + +class _FakeTokenizer: + def __init__(self, prompt_ids=(1, 2, 3, 4, 5), decoded="the answer"): + self._prompt_ids = list(prompt_ids) + self._decoded = decoded + self.decode_calls: list = [] + self.template_calls: list = [] + + def apply_chat_template(self, messages, add_generation_prompt, tokenize, return_dict, **kw): + assert tokenize is True and return_dict is False + self.template_calls.append(kw) + return list(self._prompt_ids) + + def decode(self, token_ids, skip_special_tokens=False): + self.decode_calls.append({"token_ids": list(token_ids), "skip_special_tokens": skip_special_tokens}) + return self._decoded + + +class _FakeResp: + def __init__(self, ok=True, status=200, body=b""): + self.ok = ok + self.status = status + self._body = body + + async def read(self): + return self._body + + def raise_for_status(self): + if not self.ok: + raise RuntimeError(f"HTTP {self.status}") + + +def _make_generate_server(monkeypatch: MonkeyPatch, tokenizer=None, **overrides: Any) -> SGLangModel: + server = _make_server(monkeypatch, transport="generate", base_url="http://sglang-host:30000", **overrides) + server._tokenizer = tokenizer or _FakeTokenizer() + server._sglang_urls = ["http://sglang-host:30000"] + return server + + +def _patch_http(monkeypatch: MonkeyPatch, *, result=None, resp=None) -> Dict[str, Any]: + """Patch the module globals via monkeypatch so they are restored after the test.""" + import responses_api_models.sglang_model.app as sglang_app + + rec: Dict[str, Any] = {} + + async def fake_ng_request(method, url, json=None, **kw): + rec.update(payload=json, url=url, headers=kw.get("headers")) + return resp if resp is not None else _FakeResp(ok=True) + + async def fake_get_response_json(_resp): + return result + + monkeypatch.setattr(sglang_app, "ng_request", fake_ng_request) + monkeypatch.setattr(sglang_app, "get_response_json", fake_get_response_json) + return rec + + +class _Body: + def __init__(self, **fields: Any): + self._fields = {"messages": [{"role": "user", "content": "hi"}], **fields} + + def model_dump(self, exclude_unset=True): + return dict(self._fields) + + +_GENERATE_RESULT = { + "meta_info": { + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.5, 10, "a"], [-0.25, 11, "b"]], + } +} + + +class TestGenerateTransport: + async def test_posts_input_ids_with_auth_header(self, monkeypatch: MonkeyPatch) -> None: + server = _make_generate_server(monkeypatch) + rec = _patch_http(monkeypatch, result=_GENERATE_RESULT) + + res = await server.chat_completions(MagicMock(spec=[]), _Body()) + + assert rec["url"] == "http://sglang-host:30000/generate" + assert rec["payload"]["input_ids"] == [1, 2, 3, 4, 5] + assert rec["payload"]["return_logprob"] is True + # SGLang's --api-key middleware guards /generate too. + assert rec["headers"] == {"Authorization": "Bearer dummy_key"} + + message = res.choices[0].message + assert message.generation_token_ids == [10, 11] + assert message.generation_log_probs == [-0.5, -0.25] + + async def test_graded_content_drops_special_tokens(self, monkeypatch: MonkeyPatch) -> None: + tokenizer = _FakeTokenizer() + server = _make_generate_server(monkeypatch, tokenizer=tokenizer) + _patch_http(monkeypatch, result=_GENERATE_RESULT) + + await server.chat_completions(MagicMock(spec=[]), _Body()) + + # Raw ids are kept for training, but the graded text must not carry a trailing + # special token (it breaks strict parsers like structured_outputs json.loads). + assert tokenizer.decode_calls[0]["skip_special_tokens"] is True + assert tokenizer.decode_calls[0]["token_ids"] == [10, 11] + + async def test_overflowing_prompt_is_filterable_not_head_truncated(self, monkeypatch: MonkeyPatch) -> None: + """Head-truncating would drop the newest turn and the generation cue silently.""" + server = _make_generate_server(monkeypatch, tokenizer=_FakeTokenizer(prompt_ids=range(50)), context_length=8) + rec = _patch_http(monkeypatch, result=_GENERATE_RESULT) + + res = await server.chat_completions(MagicMock(spec=[]), _Body()) + + assert res.choices[0].finish_reason == "length" + assert not res.choices[0].message.content + assert rec == {}, "no /generate call should be made for an overflowing prompt" + + async def test_tools_are_rendered_into_the_prompt(self, monkeypatch: MonkeyPatch) -> None: + tokenizer = _FakeTokenizer() + server = _make_generate_server(monkeypatch, tokenizer=tokenizer) + _patch_http(monkeypatch, result=_GENERATE_RESULT) + tools = [{"type": "function", "function": {"name": "f", "parameters": {}}}] + + await server.chat_completions(MagicMock(spec=[]), _Body(tools=tools)) + + assert tokenizer.template_calls[0]["tools"] == tools + + async def test_abort_is_rejected(self, monkeypatch: MonkeyPatch) -> None: + server = _make_generate_server(monkeypatch) + _patch_http(monkeypatch, result={"meta_info": {"finish_reason": {"type": "abort"}}}) + + with raises(RuntimeError, match="abort"): + await server.chat_completions(MagicMock(spec=[]), _Body()) + + async def test_http_error_carries_response_content(self, monkeypatch: MonkeyPatch) -> None: + """nemo_gym's exception middleware asserts on `response_content` being present.""" + server = _make_generate_server(monkeypatch) + _patch_http(monkeypatch, resp=_FakeResp(ok=False, status=400, body=b"boom")) + + with raises(Exception) as excinfo: + await server.chat_completions(MagicMock(spec=[]), _Body()) + + assert getattr(excinfo.value, "response_content", None) == b"boom" diff --git a/responses_api_models/sglang_model/tests/test_logic.py b/responses_api_models/sglang_model/tests/test_logic.py new file mode 100644 index 0000000000..338073aeef --- /dev/null +++ b/responses_api_models/sglang_model/tests/test_logic.py @@ -0,0 +1,372 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the sglang_model adapter. + +L1: pure unit tests of _logic.py (no framework / no model needed). +L6: tokenization-parity tests against the real diffusion-model tokenizer (needs transformers). + +Run standalone: python tests/test_logic.py +Or via pytest: pytest tests/test_logic.py +""" + +import os +import sys + + +HERE = os.path.dirname(os.path.abspath(__file__)) +GYM_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(HERE))) +sys.path.insert(0, GYM_ROOT) + +from responses_api_models.sglang_model._logic import ( # noqa: E402 + build_sampling_params, + cap_to_context, + extract_generated_tokens_and_logprobs, + normalize_token_ids, + unsupported_sampling_params, + would_truncate, +) + + +# A stock, widely-available AR instruct model, so the parity tests back the general claim +# rather than one custom fork. Override with SGLANG_MODEL_PATH to check a specific checkpoint. +MODEL_PATH = os.environ.get("SGLANG_MODEL_PATH", "Qwen/Qwen2.5-1.5B-Instruct") + +# L6 tokenization-parity tests need the real diffusion-model tokenizer (a local path) + transformers, +# and the chat-template ones additionally need jinja2. None of those are guaranteed in CI, so L6 SKIPS +# cleanly when unavailable (per Gym's test-skip-guard convention). L1 is pure and always runs. Point at +# a model with SGLANG_MODEL_PATH to exercise L6 locally. +_HAS_MODEL = os.path.isdir(MODEL_PATH) +try: + import transformers # noqa: F401 + + _HAS_TRANSFORMERS = True +except Exception: + _HAS_TRANSFORMERS = False +try: + import jinja2 # noqa: F401 (transformers.apply_chat_template requires it) + + _HAS_JINJA2 = True +except Exception: + _HAS_JINJA2 = False + +try: + import pytest +except ImportError: # allow standalone `python tests/test_logic.py` + pytest = None + + +def _skip_unless(cond, reason): + """Skip a test when cond is False — works under both pytest and the standalone runner.""" + + def deco(fn): + if not cond: + fn._skip_reason = reason + return pytest.mark.skipif(not cond, reason=reason)(fn) if pytest is not None else fn + + return deco + + +requires_real_model = _skip_unless( + _HAS_MODEL and _HAS_TRANSFORMERS, + f"L6 needs real model (present={_HAS_MODEL}, SGLANG_MODEL_PATH) + transformers={_HAS_TRANSFORMERS}", +) +requires_chat_template = _skip_unless( + _HAS_MODEL and _HAS_TRANSFORMERS and _HAS_JINJA2, + f"L6 chat-template needs real model + transformers + jinja2 (jinja2={_HAS_JINJA2})", +) + + +# ----------------------------- L1: extract_generated_tokens_and_logprobs ----------------------------- +def test_extract_dict_form(): + r = {"meta_info": {"output_token_logprobs": [{"token_id": 5, "logprob": -0.1}, {"id": 7, "logprob": -0.2}]}} + toks, lps = extract_generated_tokens_and_logprobs(r) + assert toks == [5, 7] and lps == [-0.1, -0.2] + + +def test_extract_tuple_form(): + r = {"meta_info": {"output_token_logprobs": [[-0.5, 11, "a"], [-0.6, 12, "b"]]}} + toks, lps = extract_generated_tokens_and_logprobs(r) + assert toks == [11, 12] and lps == [-0.5, -0.6] + + +def test_extract_val_idx_fallback(): + r = {"meta_info": {"output_token_logprobs_val": [-0.1, -0.2], "output_token_logprobs_idx": [3, 4]}} + toks, lps = extract_generated_tokens_and_logprobs(r) + assert toks == [3, 4] and lps == [-0.1, -0.2] + + +def test_extract_empty_raises(): + for r in ({}, {"meta_info": {}}, {"meta_info": {"output_token_logprobs": []}}): + try: + extract_generated_tokens_and_logprobs(r) + assert False, "expected RuntimeError" + except RuntimeError: + pass + + +def test_extract_mismatch_raises(): + r = {"meta_info": {"output_token_logprobs_val": [-0.1, -0.2], "output_token_logprobs_idx": [3]}} + try: + extract_generated_tokens_and_logprobs(r) + assert False, "expected RuntimeError" + except RuntimeError: + pass + + +def test_extract_malformed_entry_raises(): + r = {"meta_info": {"output_token_logprobs": [{"token_id": None, "logprob": -0.1}]}} + try: + extract_generated_tokens_and_logprobs(r) + assert False, "expected RuntimeError" + except RuntimeError: + pass + + +# ----------------------------- L1: normalize_token_ids ----------------------------- +def test_normalize_flat_list(): + assert normalize_token_ids([1, 2, 3]) == [1, 2, 3] + + +def test_normalize_dict(): + # transformers 5.x bug: list(dict) would grab the KEYS -> this must extract input_ids + assert normalize_token_ids({"input_ids": [9, 8, 7], "attention_mask": [1, 1, 1]}) == [9, 8, 7] + + +def test_normalize_batchencoding_like(): + class FakeBE: + def __init__(self, ids): + self._d = {"input_ids": ids} + + def __getitem__(self, k): + return self._d[k] + + @property + def input_ids(self): + return self._d["input_ids"] + + assert normalize_token_ids(FakeBE([4, 5, 6])) == [4, 5, 6] + + +def test_normalize_nested(): + assert normalize_token_ids([[1, 2, 3]]) == [1, 2, 3] + + +def test_normalize_casts_to_int(): + class IntLike(int): + pass + + out = normalize_token_ids([IntLike(1), IntLike(2)]) + assert out == [1, 2] and all(type(t) is int for t in out) + + +# ----------------------------- L1: build_sampling_params ----------------------------- +def test_sp_default_when_absent(): + sp = build_sampling_params({}, 777) + assert sp == {"temperature": 1.0, "top_p": 1.0, "max_new_tokens": 777} + + +def test_sp_precedence(): + # max_completion_tokens wins over max_tokens which wins over default + assert build_sampling_params({"max_completion_tokens": 10, "max_tokens": 20}, 99)["max_new_tokens"] == 10 + assert build_sampling_params({"max_tokens": 20}, 99)["max_new_tokens"] == 20 + + +def test_sp_optional_topk_stop(): + sp = build_sampling_params({"temperature": 0.7, "top_p": 0.9, "top_k": 40, "stop": [""]}, 50) + assert sp["temperature"] == 0.7 and sp["top_p"] == 0.9 and sp["top_k"] == 40 and sp["stop"] == [""] + # top_k None and falsy stop are omitted + sp2 = build_sampling_params({"top_k": None, "stop": []}, 50) + assert "top_k" not in sp2 and "stop" not in sp2 + + +# ----------------------------- L1: cap_to_context ----------------------------- +def test_cap_no_change_when_short(): + ids, sp = cap_to_context([1, 2, 3], {"max_new_tokens": 100}, 4096) + assert ids == [1, 2, 3] and sp["max_new_tokens"] == 100 + + +def test_cap_truncates_long_prompt(): + ids, sp = cap_to_context(list(range(5000)), {"max_new_tokens": 100}, 4096) + # prompt alone exceeds ctx -> truncate to ctx-2, leaving room for >=1 gen token, total < ctx + assert len(ids) == 4094 and ids == list(range(4094)) + assert len(ids) + sp["max_new_tokens"] < 4096 + + +def test_cap_shrinks_max_new(): + # prompt 4000, ctx 4096 -> room = 4096-4000-1 = 95 + ids, sp = cap_to_context(list(range(4000)), {"max_new_tokens": 2048}, 4096) + assert len(ids) == 4000 and sp["max_new_tokens"] == 95 + + +def test_cap_room_floor_at_one(): + # prompt == ctx: truncate to ctx-2 and floor generation at 1 token, with input+gen still < ctx + ids, sp = cap_to_context(list(range(4096)), {"max_new_tokens": 2048}, 4096) + assert len(ids) == 4094 and sp["max_new_tokens"] == 1 + assert len(ids) + sp["max_new_tokens"] == 4095 < 4096 + + +def test_cap_does_not_mutate_input(): + sp_in = {"max_new_tokens": 2048} + _, sp_out = cap_to_context(list(range(4000)), sp_in, 4096) + assert sp_in["max_new_tokens"] == 2048 and sp_out["max_new_tokens"] == 95 + + +def test_cap_ctx_zero_is_noop(): + # ctx falsy (0 or None) -> passthrough, no truncation/shrink + ids, sp = cap_to_context([1, 2, 3], {"max_new_tokens": 10}, 0) + assert ids == [1, 2, 3] and sp["max_new_tokens"] == 10 + ids2, sp2 = cap_to_context(list(range(9999)), {"max_new_tokens": 10}, None) + assert len(ids2) == 9999 and sp2["max_new_tokens"] == 10 + + +def test_cap_rejects_context_too_small_to_generate(): + # ctx == 1 previously produced prompt[:-1] and a NEGATIVE max_new_tokens, which SGLang + # rejects with an opaque 400. Fail loudly instead. + for ctx in (1,): + try: + cap_to_context([1, 2, 3], {"max_new_tokens": 10}, ctx) + except ValueError as e: + assert "too small" in str(e) + else: + raise AssertionError(f"expected ValueError for ctx={ctx}") + # ctx == 2 is the smallest workable window: empty prompt, exactly one generated token. + ids, sp = cap_to_context([1, 2, 3], {"max_new_tokens": 10}, 2) + assert ids == [] and sp["max_new_tokens"] == 1 + + +def test_would_truncate_flags_only_real_truncation(): + assert would_truncate(list(range(4095)), 4096) is True + assert would_truncate(list(range(10)), 4096) is False + # disabled / degenerate windows are not "truncation" + assert would_truncate(list(range(10)), 0) is False + assert would_truncate(list(range(10)), 1) is False + + +# ----------------------------- L1: sampling param coverage ----------------------------- +def test_sp_zero_max_tokens_is_honored_not_defaulted(): + # `or`-chaining treated max_tokens=0 as absent and silently substituted the default. + assert build_sampling_params({"max_tokens": 0}, 1024)["max_new_tokens"] == 0 + assert build_sampling_params({"max_completion_tokens": 0}, 1024)["max_new_tokens"] == 0 + + +def test_sp_forwards_penalties_and_min_p(): + sp = build_sampling_params( + {"frequency_penalty": 0.5, "presence_penalty": 0.25, "repetition_penalty": 1.1, "min_p": 0.05}, 50 + ) + assert sp["frequency_penalty"] == 0.5 and sp["presence_penalty"] == 0.25 + assert sp["repetition_penalty"] == 1.1 and sp["min_p"] == 0.05 + + +def test_unsupported_sampling_params_are_reported(): + # These would otherwise be dropped silently, so the realized rollout distribution would + # differ from the configured recipe with nothing in the data to show it. + assert unsupported_sampling_params({"n": 4, "seed": 7}) == ["n", "seed"] + assert unsupported_sampling_params({"temperature": 0.7}) == [] + + +def test_extract_val_fallback_uses_output_ids_when_idx_empty(): + # idx empty -> fall back to output_ids + r = { + "meta_info": {"output_token_logprobs_val": [-0.1, -0.2], "output_token_logprobs_idx": []}, + "output_ids": [21, 22], + } + toks, lps = extract_generated_tokens_and_logprobs(r) + assert toks == [21, 22] and lps == [-0.1, -0.2] + + +def test_extract_reads_a_chat_choice_not_just_a_generate_body(): + # The chat transport hands the *choice* to the same parser; same meta_info shape. + choice = {"meta_info": {"output_token_logprobs": [[-0.5, 10, "a"], [-0.25, 11, "b"]]}} + toks, lps = extract_generated_tokens_and_logprobs(choice) + assert toks == [10, 11] and lps == [-0.5, -0.25] + + +# ----------------------------- L6: tokenization parity (real tokenizer) ----------------------------- +def _load_tokenizer(): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + + +@requires_chat_template +def test_l6_chat_template_parity(): + tok = _load_tokenizer() + msgs = [{"role": "user", "content": "What is 2 + 2? Put the answer in \\boxed{}."}] + ids_tok = normalize_token_ids( + tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=True, return_dict=False) + ) + assert len(ids_tok) > 0 and all(isinstance(t, int) for t in ids_tok) + text = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False) + ids_txt = tok.encode(text, add_special_tokens=False) + assert ids_tok == ids_txt, f"template tokenize != encode(template text): {len(ids_tok)} vs {len(ids_txt)}" + + +@requires_chat_template +def test_l6_dict_return_normalization_real(): + tok = _load_tokenizer() + msgs = [{"role": "user", "content": "hello"}] + flat = normalize_token_ids( + tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=True, return_dict=False) + ) + as_dict = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=True, return_dict=True) + assert normalize_token_ids(as_dict) == flat # normalizer handles the real dict/BatchEncoding + + +@requires_real_model +def test_l6_skip_special_tokens_strips(): + tok = _load_tokenizer() + special_id = tok.eos_token_id + if special_id is None and tok.all_special_ids: + special_id = tok.all_special_ids[0] + assert special_id is not None + body = tok.encode("hello world", add_special_tokens=False) + ids = body + [special_id] + clean = tok.decode(ids, skip_special_tokens=True) + raw = tok.decode(ids, skip_special_tokens=False) + special_str = tok.convert_ids_to_tokens(special_id) + # the special-token string is present raw but stripped from the clean (verifier) content + assert clean != raw + assert special_str not in clean + + +# ----------------------------- runner ----------------------------- +def _run(): + import traceback + + items = sorted(globals().items()) + l1 = [k for k, v in items if k.startswith("test_") and not k.startswith("test_l6") and callable(v)] + l6 = [k for k, v in items if k.startswith("test_l6") and callable(v)] + npass = nfail = nskip = 0 + for name in l1 + l6: + fn = globals()[name] + reason = getattr(fn, "_skip_reason", None) + if reason: + print(f" SKIP {name}: {reason}") + nskip += 1 + continue + try: + fn() + print(f" PASS {name}") + npass += 1 + except Exception as e: + print(f" FAIL {name}: {type(e).__name__}: {e}") + traceback.print_exc() + nfail += 1 + print(f"\n{npass} passed, {nfail} failed, {nskip} skipped (L1={len(l1)}, L6={len(l6)})") + return nfail + + +if __name__ == "__main__": + sys.exit(1 if _run() else 0) diff --git a/responses_api_models/vllm_model/app.py b/responses_api_models/vllm_model/app.py index ce388e6203..c6abdbdec5 100644 --- a/responses_api_models/vllm_model/app.py +++ b/responses_api_models/vllm_model/app.py @@ -516,64 +516,76 @@ async def chat_completions( ) if self.config.return_token_id_information and "prompt_token_ids" not in choice_dict["message"]: - # Check vLLM honored the logprobs request. - # It returns choice.logprobs=None when it computed none. - # That happens when a null top_logprobs reached it, or the contract changed across versions. - # Without this check the code below raises a TypeError or emits empty token ids that zero the loss mask. - # An empty content list is a valid zero-token generation and passes through. - logprobs_block = choice_dict.get("logprobs") - if not logprobs_block or logprobs_block.get("content") is None: - raise RuntimeError( - f"`{self.config.name}` requested per-token logprobs from vLLM " - f"(return_token_id_information=True, logprobs=True, top_logprobs=0), but the response " - f"had none (choice.logprobs={logprobs_block!r}). Cannot extract token ids or logprobs." - ) - log_probs = logprobs_block["content"] - generation_log_probs = [log_prob["logprob"] for log_prob in log_probs] + await self._attach_token_id_information(choice_dict, body_dict, client) - """ - START TODO remove this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids - """ - # Looks like `"token_id:151667"` - generation_token_ids = [log_prob["token"].removeprefix("token_id:") for log_prob in log_probs] - - # The tokenize endpoint doesn't accept any sampling parameters - # The only relevant params are model, messages, and tools. - # - # IMPORTANT: pass through chat-template knobs (e.g. enable_thinking) - # when tokenizing, otherwise `prompt_token_ids` (and therefore logged - # `prompt_str`) can be built with different chat template settings than - # the actual generation request. - tokenize_body_dict = dict() - for key in ("model", "messages", "tools", "chat_template_kwargs"): - if key in body_dict: - tokenize_body_dict[key] = body_dict[key] - - # The base url has /v1 at the end but vLLM's tokenize endpoint does not have v1, hence the .. - tokenize_response = await client.create_tokenize(**tokenize_body_dict) - """ - END - """ + return NeMoGymChatCompletion.model_validate(chat_completion_dict) - message_dict = choice_dict["message"] - message_dict.update( - dict( - # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids - # prompt_token_ids=chat_completion_dict["prompt_token_ids"], - prompt_token_ids=tokenize_response["tokens"], - # generation_token_ids=choice_dict["token_ids"], - generation_token_ids=generation_token_ids, - generation_log_probs=generation_log_probs, - ) + async def _attach_token_id_information( + self, choice_dict: Dict[str, Any], body_dict: Dict[str, Any], client: NeMoGymAsyncOpenAI + ) -> None: + """Attach the training token fields to ``choice_dict["message"]`` in place. + + Sets ``prompt_token_ids``, ``generation_token_ids`` and ``generation_log_probs``. + Split out from ``chat_completions`` so a backend whose server returns this + information natively (e.g. SGLang's ``meta_info``) can override this one step + rather than reimplementing the whole endpoint. + """ + # Check vLLM honored the logprobs request. + # It returns choice.logprobs=None when it computed none. + # That happens when a null top_logprobs reached it, or the contract changed across versions. + # Without this check the code below raises a TypeError or emits empty token ids that zero the loss mask. + # An empty content list is a valid zero-token generation and passes through. + logprobs_block = choice_dict.get("logprobs") + if not logprobs_block or logprobs_block.get("content") is None: + raise RuntimeError( + f"`{self.config.name}` requested per-token logprobs from vLLM " + f"(return_token_id_information=True, logprobs=True, top_logprobs=0), but the response " + f"had none (choice.logprobs={logprobs_block!r}). Cannot extract token ids or logprobs." ) + log_probs = logprobs_block["content"] + generation_log_probs = [log_prob["logprob"] for log_prob in log_probs] - # Clean the duplicated information - choice_dict.pop("logprobs") - # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids - # chat_completion_dict.pop("prompt_token_ids") - # choice_dict.pop("token_ids") + """ + START TODO remove this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids + """ + # Looks like `"token_id:151667"` + generation_token_ids = [log_prob["token"].removeprefix("token_id:") for log_prob in log_probs] - return NeMoGymChatCompletion.model_validate(chat_completion_dict) + # The tokenize endpoint doesn't accept any sampling parameters + # The only relevant params are model, messages, and tools. + # + # IMPORTANT: pass through chat-template knobs (e.g. enable_thinking) + # when tokenizing, otherwise `prompt_token_ids` (and therefore logged + # `prompt_str`) can be built with different chat template settings than + # the actual generation request. + tokenize_body_dict = dict() + for key in ("model", "messages", "tools", "chat_template_kwargs"): + if key in body_dict: + tokenize_body_dict[key] = body_dict[key] + + # The base url has /v1 at the end but vLLM's tokenize endpoint does not have v1, hence the .. + tokenize_response = await client.create_tokenize(**tokenize_body_dict) + """ + END + """ + + message_dict = choice_dict["message"] + message_dict.update( + dict( + # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids + # prompt_token_ids=chat_completion_dict["prompt_token_ids"], + prompt_token_ids=tokenize_response["tokens"], + # generation_token_ids=choice_dict["token_ids"], + generation_token_ids=generation_token_ids, + generation_log_probs=generation_log_probs, + ) + ) + + # Clean the duplicated information + choice_dict.pop("logprobs") + # TODO add this when NeMo RL upgrades to vLLM 0.10.2 support for prompt token ids + # chat_completion_dict.pop("prompt_token_ids") + # choice_dict.pop("token_ids") async def _chat_completions_via_completions_api( self, request: Request, body: NeMoGymChatCompletionCreateParamsNonStreaming diff --git a/tests/unit_tests/test_responses_converter.py b/tests/unit_tests/test_responses_converter.py index 400b1a4ba0..9023ec913d 100644 --- a/tests/unit_tests/test_responses_converter.py +++ b/tests/unit_tests/test_responses_converter.py @@ -134,6 +134,57 @@ def test_flush_assistant_emits_training_message_when_token_info_present(): assert state.messages[0]["generation_token_ids"] == [3] +def test_flush_assistant_does_not_leak_token_info_to_a_later_message(): + """Token ids describe one turn; a later assistant turn must not inherit them. + + Regression: `token_information` used to survive the flush, so an assistant message that + carried no ids of its own was stamped with the PREVIOUS turn's ids -- attributing one turn's + generated tokens to another turn's text in the training data. + """ + from nemo_gym.openai_utils import TokenIDLogProbMixin + + state = ResponsesConverterState(return_token_id_information=True) + + # Turn 1 carries ids. + state.content_buffer = "turn one" + state.token_information = TokenIDLogProbMixin( + prompt_token_ids=[1, 2], + generation_token_ids=[3], + generation_log_probs=[-0.1], + ) + state.flush_assistant() + + # Turn 2 carries none (e.g. a harness-injected or rewritten assistant message). + state.content_buffer = "turn two" + state.flush_assistant() + + assert state.messages[0]["content"] == "turn one" + assert state.messages[0]["prompt_token_ids"] == [1, 2] + + assert state.messages[1]["content"] == "turn two" + for field in ("prompt_token_ids", "generation_token_ids", "generation_log_probs"): + assert field not in state.messages[1], f"turn two inherited turn one's {field}" + + +def test_flush_assistant_clears_token_info_even_when_buffers_are_empty(): + """An ids-carrying item that produces no content/tool_calls is dropped -- its ids must go too.""" + from nemo_gym.openai_utils import TokenIDLogProbMixin + + state = ResponsesConverterState(return_token_id_information=True) + state.token_information = TokenIDLogProbMixin( + prompt_token_ids=[1, 2], + generation_token_ids=[], + generation_log_probs=[], + ) + state.flush_assistant() # early-returns: nothing buffered + assert state.messages == [] + assert state.token_information is None + + state.content_buffer = "later turn" + state.flush_assistant() + assert "prompt_token_ids" not in state.messages[0] + + # =========================================================================== # responses_to_chat_completion_create_params # ===========================================================================