Add SGLang responses-API model adaptor for AR models #1557

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft

linnanwang wants to merge 10 commits into NVIDIA-NeMo:main from linnanwang:linnan_sglang_adaptor

nemo_gym/responses_converter.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -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
@@ Expand All / @@ -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)
@@ Expand Down @@

responses_api_models/sglang_model/README.md

-Original file line number
+Diff line change
@@ -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
+.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

responses_api_models/sglang_model/__init__.py

Empty file.

responses_api_models/sglang_model/_logic.py

-Original file line number
+Diff line change
@@ -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

Navigation Menu

Uh oh!

Add SGLang responses-API model adaptor for AR models #1557

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Add SGLang responses-API model adaptor for AR models #1557

Are you sure you want to change the base?

Uh oh!

Add SGLang responses-API model adaptor for AR models #1557

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!