feat(sglang): add token-exact SGLang model server for RL training

[Kh4L]

Problem

With the SGLang generation backend, multi-turn agentic (SWE-bench-style) rollouts hit a fatal prefix-stability assert in nemo_gym.py on ~every tool-using turn (48/48 turns failed in our runs). The proxy must guarantee that each turn's freshly-built prompt has the prior accumulated tokens as an exact prefix (seen == prompt[:len(seen)]). Two root causes broke this:

1. Proxy parse drift: re-rendering a prior assistant turn from parsed text dropped multi-line tool-call blocks and mangled </think>.
2. Retokenization: byte-identical text re-tokenizes to a different BPE split, so prior tokens were no longer a prefix.

Separately, exact sampled integer token ids are needed for token-level RL. On the SGLang version NeMo-RL main currently pins (0.5.12.post1), /v1/chat/completions does not expose them (logprobs.token is a decoded string; there is no return_tokens_as_token_ids), and /tokenize only accepts a raw prompt string, so the proxy recovers ids another way.

Fix

Adds a standalone responses_api_models/sglang_model/ server. SGLangModel subclasses the vLLM VLLMModel and overrides only the generation path; SGLangModelConfig extends VLLMModelConfig with the SGLang-specific fields (a required context_length, an optional inline/loaded chat template, and a sglang_tool_format). The vLLM model server is left byte-identical and carries zero SGLang references, so the two backends are fully decoupled and selected by which server a recipe points at.

• Token-splice contiguity fix (the durable contribution). Build each turn's prompt as prompt_{K-1} + gen_{K-1}(verbatim) + delta_K, splicing the prior assistant turn's exact sampled generation_token_ids instead of re-tokenizing them (_build_sglang_prompt_ids, _update_sglang_session_seq, _sglang_followup_fragment_ids, keyed by session id). Prefix-stable by construction; cache-miss falls back to a full chat-template tokenize. A session whose tools or chat-template inputs change mid-trajectory fails loudly rather than silently splicing a stale prefix.
• Exact token ids via /generate. Generate through SGLang's native /generate (return_logprob=True) and read ids+logprobs from the returned meta_info. New NeMoGymAsyncOpenAI.create_generate. The extractor tolerates the field shapes SGLang has shipped (mapping, tuple, and split value/index arrays), cross-checks them against output_ids, and treats missing or mismatched token data as a hard error rather than silently emitting an empty completion that would corrupt the training loss mask.
• </think> preservation + tool parsing. Decode the sampled ids with skip_special_tokens=False so the reasoning close-tag survives, then re-parse into reasoning + tool_calls so the returned object is shaped exactly like the vLLM /v1/chat/completions response; every downstream Responses-API conversion is identical. Supports both hermes-JSON and qwen3_coder-style XML tool blocks.
• Over-context guard. A prompt that already fills context_length returns a terminal finish_reason="length" turn instead of letting /generate hard-reject the request and corrupt the trajectory.

Local tokenization uses transformers; the leaf package pins the version that matches NeMo-RL's SGLang worker environment.

Relationship to native token-in / token-out (SGLang 0.5.13+)

SGLang 0.5.13 added native token-in / token-out on the chat endpoint (/v1/chat/completions returns meta_info.output_token_logprobs). On 0.5.13+ the /generate id-recovery mechanism becomes unnecessary and the server can read ids straight from the chat response. Two reasons this PR is still the right shape today:

1. NeMo-RL main still pins 0.5.12.post1, which has no native TITO, so the id-recovery path is required on the current stack. It is cleanly isolated and can be swapped for the chat-endpoint read on a version bump.
2. Native TITO does not address multi-turn prefix stability. The splice/contiguity fix and the over-context guard are independent of how ids are obtained and remain necessary on every SGLang version. That is the server's durable value. Tool-call parsing is a separate matter: SGLang's native chat endpoint already parses tools server-side (including a qwen3_coder parser via --tool-call-parser), so the custom parser is a compatibility fallback for the /generate path (which returns raw text) and for older or affected server-parser versions, not the primary value of the native chat design.

Relationship to #1557

This supersedes the earlier engine-flag prototype that lived inside the vLLM server and converges with the standalone-adapter structure proposed in #1557: a separate sglang_model package rather than an engine switch on the vLLM config. It layers the multi-turn token-splice, the over-context guard, reasoning/tool parsing, and the fail-loud metadata validation on top of that structure.

Result

In our runs: multi-turn contiguity failures 48 to 0 (8/8 rollouts complete), throughput ~ the vLLM path, and the engine emits training-grade per-token logprobs (validated to within the model's own bf16/MoE numerical noise vs vLLM). The focused CPU test suite for the package (splice, over-context, tool parsing, metadata extraction) passes 36/36.

Status: draft, ported onto current main

This ports a fix developed against an older Gym revision onto the refactored main proxy. Logic was validated end-to-end on the source base; the ported version needs a functional run before merge. Companion NeMo-RL work: the Megatron->SGLang refit transport lands in NVIDIA-NeMo/RL#3190 (part of the SGLang stack #3187 -> #3190 -> #3188 -> #3189), and the async-GRPO loop + gym enablement composes on top of it. Research-only logprob-parity instrumentation from the source branch is intentionally excluded.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[cwing-nvidia]

@Kh4L can you also take a look at #1557

[Kh4L]

Pushed 261615e: added an over-context guard to the SGLang /generate path.

SGLang hard-rejects input >= context_length with a 500 (vLLM tolerates it). In a multi-turn rollout the spliced prompt can already fill the window; the 500 then corrupts the trajectory and trips the downstream contiguity assert, stalling the async step. The guard pre-checks the prompt against sglang_max_total_sequence_length and defensively catches the ClientResponseError, ending the turn cleanly with finish_reason="length" — mirroring the existing vLLM context-length handling.

The splice fix + this guard have now been exercised end-to-end by a 16-node R=32 SGLang async-GRPO run: 0 chat-500s, 15+ training steps with graceful trajectory truncation and no stall. Still draft pending a full convergence-vs-baseline comparison.

[Kh4L]

Convergence validation complete. The splice fix + over-context guard in this PR (as carried on the fork branch) have now been validated by a full-scale end-to-end run: Qwen3-30B-A3B-Thinking SWE-bench async-GRPO on SGLang, 16 nodes (8 train + 8 gen, 32 generation replicas), 63 training steps across four clean checkpoint+resume segments.

• train/total_reward/mean rose 0.10 → 0.19 → 0.22 and then held ~0.22 for two independent segments — the model learns and saturates; no stalls.
• Zero over-context 500s (the guard's pre-check + error catch both exercised).
• Rare non-contiguous turns were handled gracefully throughout (handled on the NeMo-RL side in the fork; the assert-side resilience is a separate NeMo-RL change).

Remaining known caveat: this port of the fix onto refactored Gym main is still not GPU-tested in this exact form (the validated code is the fork's pre-refactor equivalent) — keeping the PR draft until that's run.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.