Add SGLang responses-API model adaptor for AR models - #1557
Conversation
|
Seems similar to how other training frameworks use sglang
did you compare to vllm and can share?
Also, I wonder if https://github.com/PrimeIntellect-ai/renderers are useful for this |
linnanwang
changed the title
|
we had also captured some thoughts about SGLang support a couple months back - PTAL and let us know your thoughts! Adding a few more backgrounds from Slack to github |
| """Orchestration tests for SGLangModel.chat_completions. | ||
|
|
||
| These exercise the full request->generate->response path with the SGLang HTTP call | ||
| (`ng_request`) and the HF tokenizer mocked out, so no live SGLang server / GPU / model |
There was a problem hiding this comment.
There was a problem hiding this comment.
Yep, exactly that. ng_request is just nemo_gym.server_utils.request aliased on import (in app.py: from nemo_gym.server_utils import request as ng_request) — the pooled aiohttp HTTP helper. We alias it so the native /generate POST reads clearly and so this test can monkeypatch sglang_app.ng_request to run without a live SGLang server / GPU.
There was a problem hiding this comment.
Answered above — ng_request is nemo_gym.server_utils.request aliased on import. One small follow-up on the test harness itself: _patch_http overwrites the module globals without restoring them (
Gym/responses_api_models/sglang_model/tests/test_app.py
Lines 120 to 134 in 35c9773
sglang_app in the same pytest process. Switching to the monkeypatch fixture (as vllm_model/tests/test_app.py does) auto-restores them — left an inline note with details.
There was a problem hiding this comment.
Confirmed, and it's now narrower than it was: ng_request (= nemo_gym.server_utils.request, the pooled aiohttp helper) is only used by the generate transport.
The new default transport: chat goes through the inherited NeMoGymAsyncOpenAI client like vllm_model does, so there's no bespoke HTTP call on that path at all.
cwing-nvidia
left a comment
cwing-nvidia
left a comment
There was a problem hiding this comment.
thanks @linnanwang for your help on this! a few points of feedback:
1 - chat_completions returns text only — no tool-call parsing, so this isn't at parity with vllm_model
The override builds the assistant message as {"role": "assistant", "content": gen_text}. In the vLLM path, tool_calls are populated by vLLM's server-side parser (enable_auto_tool_choice / tool_call_parser); native /generate returns raw text and does none of that, and it isn't re-implemented here. Net effect: any tool-using rollout produces zero tool_calls — on turn 1, not just multi-turn — so an agent harness can't act on them.
Can you run a format-specific function-call parser over gen_text before constructing the message (SGLang ships FunctionCallParser; VeRL does exactly this keyed by format: hermes). Then populate message["tool_calls"] so the inherited converter can emit them.
2 - Please make this multi-turn-correct in this PR rather than scoping to single-turn — the native /generate design actually makes it easy.
Right now chat_completions re-renders the full history with apply_chat_template(messages) and re-tokenizes from scratch on every call. In multi-turn that re-tokenizes the prior assistant spans, which diverges from the token ids the model actually emitted (Problems 1–2 in the on-policy doc)
Note how the ecosystem handles this: NeMo RL's vLLM HTTP path splices real prior token ids (_replace_prefix_tokens), and VeRL's SGLang rollout uses delta-tokenization gated by a tokenization_sanity_check. Both are working around a re-templating boundary. You don't have that boundary — because you POST input_ids to /generate, you can build the prompt from the real token ids directly. The converter already threads prompt_token_ids / generation_token_ids forward on each message; reuse them: concatenate the carried-forward token ids for prior turns and only template+tokenize the new environment-inserted messages (user/tool) between them. That makes the prompt exactly the sequence the policy generated — cleaner than either of the above.
Given how many of our environments are multi-step/multi-turn, I'd like this in before merge rather than a single-turn caveat. If there's a reason it can't be (e.g. a NeMo RL SGLang HTTP worker dependency), let's discuss — but the splice looks doable entirely server-side here.
3 - training run evidence
The sampling_importance_ratio ≈ 1.0 / gen_kl_error ≈ 0 over 89+ steps is the right proof — please attach it so it's reproducible: the SGLang fork commit + launch args, the Gym config + NeMo RL recipe, and the metric curves (not just endpoint values). Also can you clarify whether this has run on any standard HF model or only the one fork.
|
Does SGLang 0.5.13's native TITO change this design? I just saw ProRL-Agent-Server#43 upgraded to SGLang 0.5.13 and removed its TITO patch now that native support exists. |
|
@cmunley1 re "did you compare to vllm and can share?" — so far I've run the Gym setup with NeMo RL and see rewards increasing across multiple agents, but I don't yet have a side-by-side comparison against vLLM. What's your suggested plan to compare? Happy to set up whatever apples-to-apples configuration you'd find most convincing. |
|
@cmunley1 thanks for the pointer — There seem to be two viable paths here:
We'll keep both in consideration and figure out which fits reality best as we make the multi-turn path correct. Appreciate the suggestion! |
|
@cwing-nvidia heads up — on your multi-turn point (#2), @cmunley1 suggested PrimeIntellect-ai/renderers, which targets the same re-tokenization-drift problem: it does token-aware rendering and offers So we have two candidate paths for making the multi-turn path correct:
We'll evaluate both and pick whichever fits best in practice. Wanted to make sure you saw the |
|
@cwing-nvidia good question — I dug in. Short version: 0.5.13's native TITO doesn't change this design; it's the same mechanism we build on. "Native TITO" is the What would actually let us drop the adapter is token-id return on |
Hey @linnanwang SGLang's recent patch https://github.com/sgl-project/sglang/pull/23751/changes attached necessary TITO token & logprob to its v1/chat/completion endpoint. Have you checked it out? |
|
@billxbf @cwing-nvidia thanks both — and @billxbf you're right, I need to correct my earlier note. I cloned SGLang Implementation plan — migrate to the native chat endpointThis is the cleaner path and it resolves both of @cwing-nvidia's blockers in one move:
Net effect:
|
cwing-nvidia
added
the
models
ffrujeri
left a comment
ffrujeri
left a comment
There was a problem hiding this comment.
Background
- SGLang — the serving stack this adapter targets; its native
/generateendpoint acceptsinput_idsand returnsmeta_info.output_token_logprobswhenreturn_logprob=true. - PrimeIntellect renderers — token-aware chat rendering discussed in-thread as one fix for multi-turn re-tokenization drift.
- NeMo Gym on-policy correction doc — why exact token ids/logprobs matter for RL training.
Where this PR sits
This adds a new Model Server backend. It lives on the hot inference path and the training-data serialization boundary: the token ids/logprobs it attaches to assistant messages are exactly what the RL trainer consumes, so silent mistakes here become silent training-signal corruption.
flowchart LR
Agent[Agent Harness] -->|/v1/responses| Conv[responses converter<br/>inherited from VLLMModel]
Conv --> CC[SGLangModel.chat_completions<br/>app.py]
CC -->|"POST /generate (return_logprob)"| SGL[(SGLang server<br/>external)]
SGL -->|token ids + logprobs| CC
CC --> Roll[rollout collection] --> RL[(NeMo RL / GRPO)]
classDef touched fill:#ffe08a,stroke:#d48806;
class CC touched
Summary
Thanks for this contribution — the direction is valuable (exact token-id/logprob recovery for GRPO on SGLang-served models), the pure-logic split with unit tests is well done, and I verified the suite passes end-to-end via ng_test +entrypoint=responses_api_models/sglang_model (24 passed, 3 clean skips; _logic.py at 98% coverage). The logprob_start_len=-1 usage and the (logprob, token_id, text) tuple parse are both correct against SGLang source — the worst-case token-indexing failure modes are avoided, which is great.
The inline comments focus on two themes: (1) places where the override only half-integrates with the inherited VLLMModel contract, so caller-supplied fields are silently dropped (per-request chat_template_kwargs, tools, most sampling params, api_key); and (2) the evidence bar from #976 — the vLLM parity comparison (reward/pass@1 within 1%, greedy logprob parity) is still the one experiment that would make this trustworthy as a training adapter, and the current metrics show self-consistency rather than parity.
A few housekeeping notes:
- The branch is behind
main— a rebase would be good before merge (no conflicts reported). - The three points from the earlier review (tool-call output parsing, multi-turn token splicing, reproducible training evidence) all still apply to the code as of this head; the discussion threads converged on plans, and it would help to state which of those plans land in this PR vs. follow-ups.
- Minor observations, none blocking:
${policy_base_url}must be bare here but end in/v1forvllm_model— worth an explicit warning since migrating users will hit a/v1/generate404;trust_remote_codedefaults toTrue(elsewhere in the repo it's explicit opt-in per config);AutoTokenizer.from_pretrained(config.model)will hit the network at startup if given a hub id rather than the documented local path;transformersis unpinned althoughnormalize_token_idsbranches on 5.x return shapes; thetry/except ImportError+sys.path.insertaround the cross-server import is unique to this server (all others use a plain import); and the inherited preprocess injectslogprobs=True/return_tokens_as_token_ids=Truewhich the override ignores — harmless, but a hint that the preprocess contract and this override have drifted apart.
| messages = body_dict["messages"] | ||
|
|
||
| # 1) prompt token ids via the model's own chat template (local tokenizer). | ||
| ct_kwargs = self.config.chat_template_kwargs or {} |
There was a problem hiding this comment.
Was it intentional that per-request chat_template_kwargs are not used here? The inherited _preprocess_chat_completion_create_params (called at L101) merges config + per-request metadata.chat_template_kwargs into body_dict["chat_template_kwargs"] (see vllm_model/app.py L299-310, whose comment describes per-sample reasoning on/off), but this line reads only the static config value — so per-sample template overrides are silently dropped and the prompt is tokenized with the wrong template, which corrupts prompt_token_ids for training.
| ct_kwargs = self.config.chat_template_kwargs or {} | |
| ct_kwargs = body_dict.get("chat_template_kwargs") or self.config.chat_template_kwargs or {} |
There was a problem hiding this comment.
Fixed, and on the default path the class of bug is gone rather than patched.
transport: chat (new default) sends messages to SGLang and lets the server apply the chat template, so there is no local templating that could read the wrong kwargs. The merged value from the inherited preprocess is forwarded as-is.
Where local templating still happens (transport: generate), it now reads the merged value exactly as you suggested:
ct_kwargs = body_dict.get("chat_template_kwargs") or {}And your root-cause diagnosis was right: the tests stubbed the real preprocess, which is why this slipped through. That's fixed too — see the test_app.py:108 thread.
|
|
||
| # 1) prompt token ids via the model's own chat template (local tokenizer). | ||
| ct_kwargs = self.config.chat_template_kwargs or {} | ||
| rendered = self._tokenizer.apply_chat_template( |
There was a problem hiding this comment.
Adding to the earlier tool-call feedback: beyond the missing output parsing, apply_chat_template(...) is never given tools=, so tool schemas never enter the prompt at all — the model isn't told the tools exist even before the parsing question. If a request carries tools, could they be forwarded to the chat template (HF templates accept a tools= kwarg) so the input half is covered too?
There was a problem hiding this comment.
Good catch — the input half was indeed missing.
On transport: chat this resolves structurally: tools go to SGLang with the request and the server renders them into the prompt and parses tool calls out of the output, so both halves are covered by the same mechanism that serves vllm_model.
On transport: generate, tools is now forwarded to the template:
rendered = self._tokenizer.apply_chat_template(..., tools=body_dict.get("tools"), ...)The output half still isn't parsed there — that's inherent to /generate, so it's documented as a limitation of that transport rather than papered over.
| def _post_init(self) -> None: | ||
| super()._post_init() | ||
| # The model name is a local path with tokenizer + chat_template.jinja. | ||
| self._tokenizer = AutoTokenizer.from_pretrained( |
There was a problem hiding this comment.
The prompt is tokenized locally while the SGLang server has its own copy of the tokenizer/chat template — if the local config.model path differs from what the server was launched with (different snapshot, edited chat_template.jinja, added tokens), generation is silently conditioned on ids from the wrong template, and nothing reconciles the two (the vLLM path recovers prompt ids from the server's authoritative /tokenize). Since #976 explicitly calls for token-id round-trip consistency: would a startup assertion work here — e.g. render a fixed probe prompt and compare against the server's /tokenize (or check /get_model_info) — plus a doc note that model must be the exact path/revision the server was launched with?
There was a problem hiding this comment.
This was the strongest argument for changing the design, and it's why the default transport changed.
transport: chat does no local tokenization at all. return_prompt_token_ids makes SGLang report the prompt ids it actually tokenized, which is authoritative in a way a local tokenizer can't be — so there is nothing left to reconcile and no /tokenize round-trip needed:
prompt_token_ids = choice_dict.get("prompt_token_ids")If the server predates chat-side TITO and returns nothing, we now fail loudly with a message naming the 0.5.13 requirement instead of emitting empty ids.
The drift risk only remains on transport: generate, where it's inherent; the README now states that model must be the exact path/revision the server was launched with.
| # 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) |
There was a problem hiding this comment.
api_key is required by the config (and the training YAML wires api_key: ${policy_api_key}), but the /generate POST sends no Authorization header. SGLang applies its API-key middleware to /generate when launched with --api-key (http_server.py L715-716 @ v0.4.6), so such deployments would get a guaranteed 401. Could the header be added (Authorization: Bearer {self.config.api_key}) to match how the inherited OpenAI-client path behaves?
There was a problem hiding this comment.
Correct, and thank you for the citation — that would have been a guaranteed 401 on any --api-key deployment.
transport: chat inherits the NeMoGymAsyncOpenAI client, so auth is handled the same way as vllm_model. For transport: generate the header is now set explicitly:
if self.config.api_key:
extra_request_kwargs["headers"] = {"Authorization": f"Bearer {self.config.api_key}"}Worth noting for anyone reading later: it has to be conditional. server_utils.request() does kwargs.setdefault("headers", ...), which would keep an explicit headers=None and then crash setting Content-Type on it.
| # 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) |
There was a problem hiding this comment.
apply_chat_template(..., tokenize=True) (L106) and this decode run synchronously inside the async endpoint — CPU-bound tokenization of up to ~ctx-length prompts will block the event loop under high rollout concurrency (the vLLM path does no local tokenization, so this is a new hazard on the hot path). Consider wrapping both in await asyncio.to_thread(...); there's repo precedent in e.g. resources_servers/newton_bench and resources_servers/evalplus.
There was a problem hiding this comment.
Agreed, and this one also disappears on the default path rather than being mitigated.
transport: chat does no local tokenization or decoding, so there's no CPU-bound work on the event loop — the hot path is the same pooled-aiohttp call vllm_model makes.
It remains a real consideration for transport: generate. Since that path is now a fallback for older builds rather than the primary, I left it synchronous rather than add asyncio.to_thread — happy to add it if you'd rather it be safe under concurrency there too.
| ) | ||
|
|
||
|
|
||
| MODEL_PATH = os.environ.get("SGLANG_MODEL_PATH", "/linnanw/justGRPO/asset/Nemotron-Labs-Diffusion-3B") |
There was a problem hiding this comment.
The default parity-test checkpoint is a diffusion model in ar_mode on a custom fork, while the PR title claims support for AR models generally. Two asks: (1) could at least one stock-SGLang + standard AR model run back the general claim? (2) More importantly — for this fork, are output_token_logprobs the true sampling logprobs of the tokens as they were drawn, or a teacher-forced re-score after generation? If a re-score, the on-policy guarantee doesn't hold even single-turn, which would matter more than any other issue in this review.
There was a problem hiding this comment.
This was the most important question in the review and it deserved a direct answer. Both parts:
(1) The parity-test default is now Qwen/Qwen2.5-1.5B-Instruct rather than the diffusion fork, so the general AR claim is backed by a stock model on stock SGLang.
(2) For AR decoding they are true sampling logprobs, not a teacher-forced re-score. Every path in sampler.py:127-192 derives the logprob from the same logits tensor that produced the token, in the same forward pass — there is no second scoring pass. Two caveats worth recording:
- With
top_p<1/top_k, the returned logprob is over the temperature-scaled full distribution while the token was drawn from the truncated one. Our recipe pinstemperature/top_p/top_k = 1.0/1.0/null(SGLang'ssimple_sampling_case), so it doesn't bite — but it would silently start to if someone settop_p: 0.95. - SGLang also has an explicit RL on-policy path (
rl_on_policy_target+enable_deterministic+ simple sampling) that samples directly from the bf16log_softmax(logits/T)it returns, for exact trainer parity.
But for block diffusion the answer is different, and you were right to probe here: tokens are unmasked in reveal order, not drawn left-to-right, so the per-token logprob is an estimator. The on-policy guarantee in this PR is therefore scoped to single-turn AR — which is what the title meant by "for AR models", now made explicit in the README rather than left to inference.
| _sglang_urls=["http://sglang-host:30000"], | ||
| ) | ||
| # _preprocess_chat_completion_create_params is inherited from VLLMModel; stub as passthrough. | ||
| me._preprocess_chat_completion_create_params = lambda request, body_dict: body_dict |
There was a problem hiding this comment.
Stubbing _preprocess_chat_completion_create_params as identity means the real preprocess contract is never exercised — which is exactly how the per-request chat_template_kwargs issue (see app.py L105) slipped past these tests. The suite also only ever passes SimpleNamespace() requests (so only the sid="" branch of URL selection runs) and has no cases for abort finish reasons, tools, response_format, or truncation. Worth one test that runs the real preprocess against a fake config to lock in that contract?
There was a problem hiding this comment.
You were right about the mechanism, not just the symptom — stubbing the preprocess is precisely how the app.py:105 bug survived the suite.
The tests are rebuilt on the vllm_model convention: a real SGLangModel with a mocked ServerClient, so the real inherited preprocess runs and the contract is locked in:
def _make_server(monkeypatch, **overrides) -> SGLangModel:
...
return SGLangModel(config=config, server_client=MagicMock(spec=ServerClient, global_config_dict={}))test_preprocess_honors_per_request_chat_template_kwargs asserts the metadata merge specifically, so the original bug is now caught by construction.
Coverage added for the other gaps you listed: abort finish reasons (both transports), tools, and truncation. response_format is covered indirectly via the unsupported-params reporting rather than a dedicated case.
| return dict(self._fields) | ||
|
|
||
|
|
||
| def _patch_http(monkeypatch_like, *, result=None, resp=None): |
There was a problem hiding this comment.
_patch_http overwrites the module globals sglang_app.ng_request / sglang_app.get_response_json without restoring them, so the fakes stay installed for the rest of the pytest process for any test importing this module. The parent vllm_model/tests/test_app.py (and the other model-server suites) use the monkeypatch fixture, which auto-restores — could this harness take monkeypatch and use monkeypatch.setattr(sglang_app, "ng_request", ...)?
There was a problem hiding this comment.
Fixed exactly as suggested — the harness takes monkeypatch and the globals are restored:
def _patch_http(monkeypatch: MonkeyPatch, *, result=None, resp=None):
monkeypatch.setattr(sglang_app, "ng_request", fake_ng_request)
monkeypatch.setattr(sglang_app, "get_response_json", fake_get_response_json)The whole file now uses the monkeypatch fixture rather than assigning module globals, matching vllm_model/tests/test_app.py, so nothing can leak into later tests in the same process.
| assert len(ids) + sp["max_new_tokens"] == 4095 < 4096 | ||
|
|
||
|
|
||
| def test_cap_does_not_mutate_input(): |
There was a problem hiding this comment.
Two small additions would close the remaining _logic.py coverage gaps (the falsy-ctx early return at L105, and the output_ids fallback branch at L54); both verified passing against this branch:
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_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]There was a problem hiding this comment.
Both added verbatim — thank you for writing them out and verifying them against the branch.
They're joined by coverage for the behaviour that changed since: the ctx<2 guard, would_truncate, max_tokens=0, the expanded passthrough params, the unsupported-param reporting, and a case asserting the same parser handles a chat choice (the meta_info shape is identical on both transports, which is what lets the two paths share one kernel).
Pure-logic suite is 27 passing, up from 19.
| @@ -0,0 +1,155 @@ | |||
| # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
There was a problem hiding this comment.
diagnostic_vllm_vs_sglang.py:1
The script is nicely env-gated (no hardcoded endpoints, non-zero exit when unconfigured), but the name promises a vLLM-vs-SGLang comparison while it only probes whether a live SGLang server satisfies three vllm_model assumptions — it never runs vLLM. It's also the only standalone diagnostic shipped inside a server directory in the repo, and it isn't collected by gym env test. Would it make sense to capture its findings table durably in the README (the durable value is the results), and either rename it to reflect the one-sided probe (e.g. probe_sglang_vllm_assumptions.py) and wire it as a skippable test, or drop the script from the shipped tree?
There was a problem hiding this comment.
Dropped from the shipped tree — the last of your three options.
Your critique of the name was fair, but the deeper issue is that the script's premise expired: it existed to argue "why /generate instead of vllm_model", and with the chat transport on stock 0.5.13 that argument no longer holds. Keeping a renamed probe would have preserved a conclusion that's now wrong.
The durable value went to the README as you suggested — the version boundary and each transport's limitations are stated there directly.
35c9773 to
ac66157
Compare
Update: migrated to SGLang's OpenAI chat endpoint on stock 0.5.13 — and a correction to my own earlier commentPushed as three commits on top of First, a correction@billxbf was right, and I was more wrong than I acknowledged. On Jun 25 I said chat-side TITO was And ProRL's What changed
That pushes templating, tool-call parsing, sampling params, auth and overflow handling back to the server, which removes rather than patches a whole class of review comments:
Directly fixed: On sampling-vs-rescore (@ffrujeri) — the question you flagged as mattering mostIt's a sampling-time capture, never a re-score. In every path the logprobs come from the same logits tensor that produced the token (
On multi-turn (@cwing-nvidia #2) — I'd like to push back, with evidenceI designed the splice and had it adversarially reviewed before writing it, and I don't think it's safe to ship yet. The blocker isn't SGLang — A splice that derives the prefix from carried Two further blockers found: for Nemotron-style templates that emit glue between content and the turn marker, the trim degenerates and every multi-turn turn either raises or double-emits the marker; and nothing validates that the carried prefix still corresponds to the caller's current messages. The splice needs a prefix-validation layer that doesn't exist yet. I'd like to do that as a focused follow-up rather than land it unsound here — happy to discuss if you'd rather block on it. Bonus: a real bug in core
Still outstandingThe vLLM parity run (@ffrujeri's protocol: stock SGLang + a standard AR model, greedy, per-token max/mean |Δlogprob| + pass@1) and the 89-step artifacts (@cwing-nvidia #3). Both need cluster time and are in progress. One caveat in the interest of honesty: local verification here covered the pure-logic suite only — CI is the first execution of the app-level and converter tests. |
Brought over from the core working copy: a responses_api_models/sglang_model package (app, logic, config, requirements, tests). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
Self-contained, server-free probe showing why a dedicated sglang_model server is needed instead of reusing vllm_model. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
…nt message
`ResponsesConverterState.flush_assistant()` reset `content_buffer` and
`tool_calls_buffer` but left `token_information` in place. An assistant message
that carried no token ids of its own was therefore stamped with the PREVIOUS
assistant turn's `prompt_token_ids` / `generation_token_ids` /
`generation_log_probs`, attributing one turn's generated tokens to another
turn's text in the training data.
Reachable whenever a harness injects or rewrites an assistant message without
token ids, e.g.
assistant(ids) -> tool_output(flush) -> assistant(no ids) -> tool_output(flush)
where the second flush reuses the first turn's ids.
`token_information` describes exactly one assistant turn, so consume it into a
local and clear it on every path -- including the empty-buffer early return,
since an item can carry ids while producing no content and no tool calls (a
zero-token generation), and that message is dropped.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: linnan wang <linnanw@nvidia.com>
… hook Move the body of the `return_token_id_information` block in `chat_completions` into `_attach_token_id_information(choice_dict, body_dict, client)`. Pure code move: the 26 executable lines are byte-identical and the call site is unchanged, so behavior is identical for every existing caller. Splitting it out lets a backend whose server returns token ids natively -- e.g. SGLang, via `meta_info` -- override this one step instead of reimplementing the whole endpoint and drifting from the inherited preprocess contract. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
Chat-side TITO is native in the 0.5.13 release, not main-only: sgl-project/sglang#23751 merged 2026-06-04 and v0.5.13 was cut 2026-06-13, so the tag contains it. Stock 0.5.13 exposes `return_prompt_token_ids` / `return_meta_info` on the request and `prompt_token_ids` / `meta_info` on each choice. No patched build or fork is required -- ProRL's `patch_sglang_0513_token_metadata.sh` only makes `logprobs=true` imply those flags, which this server sets explicitly. `transport: chat` (new default) therefore inherits everything from VLLMModel and overrides only `_attach_token_id_information`. Prompt templating, tool-call parsing, sampling params, auth and context-overflow handling all move server-side, which removes -- rather than patches -- the local-tokenizer drift, the missing Authorization header, the client-side `cap_to_context` truncation, the dropped sampling params, and the event-loop blocking of local tokenization. `transport: generate` keeps the native `/generate` path for builds that predate chat-side TITO (e.g. forks serving diffusion LLMs), with fixes: merged `chat_template_kwargs` so per-request overrides are honored, `tools` rendered into the prompt, an Authorization header, expanded sampling params plus a loud warning for ones it cannot honor, and an overflowing prompt returning a filterable empty `finish_reason="length"` instead of a head-truncated prompt. In both transports a generation SGLang reports as `finish_reason="abort"` now raises, so a server-cancelled partial rollout cannot enter a training batch looking like a normal completion. Also: `trust_remote_code` defaults to false, `transformers` is pinned and imported lazily, the cross-server import uses a plain import (matching local_vllm_model_proxy), and `cap_to_context` rejects a context too small to generate instead of POSTing a negative `max_new_tokens`. Tests are rebuilt on the vllm_model convention -- a real SGLangModel with a mocked ServerClient -- so the real inherited preprocess runs rather than being stubbed, and all patching goes through `monkeypatch` so module globals cannot leak between tests. 27 pure-logic tests pass (up from 19). The on-policy guarantee is documented as single-turn and AR-mode: under block diffusion the per-token logprob is an estimator, not the logprob of an AR draw. The obsolete vllm-vs-sglang diagnostic is dropped; its finding (the version boundary) is now stated in the README. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
Blank line before the module constant, and collapse a call that fits within the 119-column limit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: linnan wang <linnanw@nvidia.com>
ac66157 to
8528f86
Compare
5 participants
Add
sglang_model: a Responses-API model server backed by SGLang native/generateSummary
Adds a new model server,
responses_api_models/sglang_model/, that drives an SGLang server's native/generateendpoint withreturn_logprob=trueinstead of the OpenAI-compatible/v1/chat/completionspath.This lets Gym recover the exact generated token ids and their logprobs for use as the policy model server in RL training, such as GRPO.
Motivation
The existing
vllm_modelserver recovers training token ids by parsingtoken_id:NNNlogprob tokens out of/v1/chat/completions.Some SGLang-served models — notably custom forks — do not expose those tokens through the OpenAI-compatible path. As a result, the exact token sequence and logprobs emitted by the policy cannot be recovered, forcing the trainer to re-tokenize a decoded string, which makes training off-policy.
SGLang's native
/generatereturns the generated token ids and logprobs directly, eliminating that gap.What it does
SGLangModelsubclassesvllm_model'sVLLMModeland overrides onlychat_completions.Per request, it:
Renders the prompt to token ids via the model's own HF chat template using the local tokenizer.
Caps the prompt to the context window and shrinks
max_new_tokenssoinput + gen < context, avoiding SGLang 400s.POSTs to
{base_url}/generatewith thoseinput_ids, then parsesmeta_info.output_token_logprobsinto:generation_token_idsgeneration_log_probsAttaches:
prompt_token_idsgeneration_token_idsgeneration_log_probsto the assistant message when
return_token_id_information: true, exactly as the vLLM path does.The graded
contentis decoded withskip_special_tokens=True; raw token ids are preserved for training.Everything else — Responses ↔ ChatCompletions conversion,
responses(), and the assistant-message training-class upgrade — is inherited unchanged fromvllm_model.All HTTP goes through
nemo_gym.server_utils.request, the pooledaiohttpclient.Files
app.py—SGLangModel/SGLangModelConfig; overrideschat_completions._logic.py— pure, framework-free transforms:extract_generated_tokens_and_logprobsnormalize_token_idsbuild_sampling_paramscap_to_contextSplit out for unit testing.
configs/sglang_model_for_training.yaml— registers the server;base_url/api_key/modelvia${policy_*}interpolation; enablesreturn_token_id_information: true.requirements.txt—-e nemo-gym[dev] @ ../../plustransformers.tests/test_logic.py— unit and tokenizer-parity tests for_logic.py.README.mdTesting
tests/test_logic.pycovers the pure logic:ng_test +entrypoint=responses_api_models/sglang_modelbuilds the server venv and runs the suite.Exercised end-to-end as the policy server in a GRPO run:
Testing & verification
Reviewer checklist → what we verified
vllm_worker_async_replace_prefix_tokens)nemo_rlworker concern (multi-turn retokenization). We're single-turn (max_rollout_turns=1) and passinput_ids/keep raw output ids, so it doesn't apply. Not part of this Gym PR.vllm_modelbreaks on the real server (returns token text, ignoresreturn_tokens_as_token_ids→ silent id corruption).sglang_modelfixes it via native/generate(real integer ids).sglang_modelsubclassesVLLMModel, overrides onlychat_completions; converter reused verbatim.cap_to_contextvllm_model's overflow detection works on this fork.sglang_modelprevents overflow viacap_to_context(a design choice). Found + fixed an off-by-one there.The "run SGLang via
vllm_model, see what breaks" experimentWe implemented exactly this suggestion, in two iterations:
diagnostic_vllm_vs_sglang.py) — flagged as "loose" (we supplied the responses), so it over-predicted breakage.ar_mode, same args).return_tokens_as_token_ids=True)/tokenize→ Only 1 of 3 actually breaks. Running it for real corrected the mock's overclaims and pinned the single genuine incompatibility: the token-id logprob format.
Other sanity checks
test_logic.py(22 tests: 19 L1 + 3 L6)test_app.py(5 tests)chat_completionsorchestration (payload, id/logprob attach,skip_special_tokens, error path)nemogym_smoke*)raise_for_statusmiddleware, chat-template dict, context 400, `<gen_kl_error ≈ 0,sampling_importance_ratio ≈ 1.0over 89+ GRPO stepscap_to_contextoff-by-oneSummary
/tokenizeand overflow work on this fork, socap_to_context/ native/generateremain clean design choices, not failure workarounds.cap_to_contextbug (fixed in code).Net:
sglang_modelis justified by the single, real, live-confirmed token-id gap; everything else is verified or scoped out.Closes #976