nemo_rl.models.generation.openai_server_utils

Shared helpers for the OpenAI-compatible HTTP generation servers.

These utilities are backend-agnostic: they operate on token-ID lists plus a tokenizer, with no engine calls. They are shared by the vLLM async worker (vllm_worker_async.py) and the TRT-LLM HTTP server (trtllm_http_server.py), which both put a message-based /v1/chat/completions layer in front of a token engine for the agentic NeMo-Gym path. SGLang does not use these — it is driven token-in/token-out via generate(input_ids) and never re-templates messages, so it has no retokenization drift to correct.

Module Contents

Functions

replace_prefix_tokens

This is a subroutine used inside the OpenAI-compatible Chat Completion server.

API

nemo_rl.models.generation.openai_server_utils.replace_prefix_tokens(

tokenizer: Any,

model_prefix_token_ids: list[int],

template_prefix_token_ids: list[int],

template_token_ids: list[int],

) → list[int]

This is a subroutine used inside the OpenAI-compatible Chat Completion server.

This function is for fixing up the chat template-tokenized messages history to match the model output tokenization up to the last assistant turn, in order to preserve the monotonic tokens property for optimized multi-turn training.

Some environments (namely NeMo-Gym) require an OpenAI compatible server endpoint rather than an inference engine handle. This is fine for the most part, but it may cause issues when the environment is used as a part of training.

RL training frameworks train models on token IDs, but the OpenAI compatible server communicates in what is basically de-tokenized text. When multiple model calls are made to the OpenAI compatible server in a single trajectory, model generations in previous model calls may be re-tokenized to something that is different than what was generated. This is not too big of an issue (that we know of) at inference time, but the log probs the model produces are different enough for the differently re-tokenized generation result that it causes the training to be off policy. Off policy isn’t necessarily a bad thing in isolation, but this source of off-policyness may cause unexpected issues if not properly accounted for. It also mis-aligns the token ID sequences across model calls, which feels very strange during training.

There are real cases where the model output string does not match the chat template tokenization of the parsed model output. A concrete example is inconsistent whitespace tokens around tool call special tokens.

TODO When NeMo RL supports training image generation models, we want to revisit and possibly update this function. This issue occurs when the model generates tokens that are de-tokenized into text or images, and then re-tokenized into tokens. So if there is a situation like that with images and image tokenization is non-unique, then we will need to uppdate this function.

The splice boundary is located by EOS count, not position: count the EOS tokens in template_prefix_token_ids and cut at the N-th EOS in template_token_ids. This is robust to chat templates that strip reasoning () blocks from history when the last message is a user turn – that shifts token positions but not the per-message EOS count, so counting still finds the same boundary (and reduces to the last EOS of the prefix when nothing is stripped).

Example (turn-by-turn, concise; eos_token_id = 2): Turn 1: - prefill_T1 (template prefill) = [11,12,13,40,41] - model output = [220,17,2] # decodes to “ 4” + EOS - model_prefix_token_ids = prefill_T1 + model output => [11,12,13,40,41,220,17,2]

Turn 2 (template retokenizes prior assistant text differently):
    - template_prefix_token_ids = [11,12,13,40,41,1001,2]  # 1001 decodes to " 4"
    - template_token_ids = [11,12,13,40,41,1001,2,21,22,40,41]

replace_prefix_tokens keeps the exact prior model tokens up to EOS and
resumes from the template after that EOS:
    output => [11,12,13,40,41,220,17,2,21,22,40,41]