fix: Add chat completion parity tests for TRTLLM v.s. vLLM and fix disparities

examples/configs/recipes/llm/grpo-qwen3-0.6b-1n4g-megatron-trtllm-tp2-noncolocated-gym.yaml

@@ -0,0 +1,81 @@
+defaults: ../../../nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
+grpo:
+  num_prompts_per_step: 4
+  num_generations_per_prompt: 2
+  adv_estimator:
+    name: reinforce_plus_plus
+    minus_baseline: false
+  async_grpo:
+    in_flight_weight_updates: true
+  val_period: 1000
+  val_at_start: false
+checkpointing:
+  checkpoint_dir: results/grpo-qwen3-0.6b-1n4g-megatron-trtllm-tp2-noncolocated-gym
+policy:
+  model_name: Qwen/Qwen3-0.6B
+  tokenizer:
+    name: Qwen/Qwen3-0.6B
+    chat_template_kwargs:
+      enable_thinking: false
+  train_global_batch_size: 8
+  max_total_sequence_length: 2048
+  make_sequence_length_divisible_by: 2
+  megatron_cfg:
+    sequence_parallel: true
+    activation_checkpointing: false
+    apply_rope_fusion: false
+    defer_fp32_logits: true
+    moe_per_layer_logging: false
+  optimizer:
+    name: torch.optim.AdamW
+    kwargs:
+      lr: 5.0e-06
+      weight_decay: 0.01
+      betas:
+      - 0.9
+      - 0.999
+      eps: 1.0e-08
+  generation:
+    backend: trtllm
+    max_new_tokens: 512
+    stop_token_ids:
+    - 151643
+    - 151645
+    trtllm_cfg:
+      tensor_parallel_size: 2
+      moe_tensor_parallel_size: null
+      moe_expert_parallel_size: null
+      max_model_len: ${policy.max_total_sequence_length}
+      precision: ${policy.precision}
+      async_engine: true
+      gpu_memory_utilization: 0.6
+      max_batch_size: 64
+      max_num_tokens: 4096
+      in_flight_weight_updates: ${grpo.async_grpo.in_flight_weight_updates}
+      recompute_kv_cache_after_weight_updates: ${grpo.async_grpo.recompute_kv_cache_after_weight_updates}
+      expose_http_server: true
+      default_chat_template_kwargs:
+        enable_thinking: false
+      tool_parser: qwen3_coder
+      reasoning_parser: qwen3
+    trtllm_kwargs:
+      batch_wait_timeout_iters: 32
+      batch_wait_max_tokens_ratio: 0.5
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 2
+        num_nodes: 1
+env:
+  nemo_gym:
+    policy_model:
+      responses_api_models:
+        vllm_model:
+          uses_reasoning_parser: true
+logger:
+  tensorboard_enabled: true
+  wandb:
+    project: nemo-rl
+    name: grpo-qwen3-0.6b-1n4g-megatron-trtllm-tp2-noncolocated-gym
+cluster:
+  gpus_per_node: 4

nemo_rl/models/generation/openai_server_utils.py

@@ -22,6 +22,7 @@
 so it has no retokenization drift to correct.
 """
 
+from collections.abc import Collection
 from typing import Any
 
 
@@ -30,6 +31,7 @@ def replace_prefix_tokens(
     model_prefix_token_ids: list[int],
     template_prefix_token_ids: list[int],
     template_token_ids: list[int],
+    model_stop_token_ids: Collection[int] | None = None,
 ) -> list[int]:
     """This is a subroutine used inside the OpenAI-compatible Chat Completion server.
 
@@ -95,11 +97,9 @@ def replace_prefix_tokens(
     eos_token_id = tokenizer.eos_token_id
     assert eos_token_id is not None, "Tokenizer must have an EOS token ID"
 
-    # The model isn't guaranteed to end on EOS (e.g. it hit max_tokens); chat
-    # templates always add one, so cut the model input to just before its EOS.
-    model_cut_end = len(model_prefix_token_ids)
-    if model_prefix_token_ids[-1] == eos_token_id:
-        model_cut_end -= 1
+    effective_stop_token_ids = set(model_stop_token_ids or ())
+    effective_stop_token_ids.add(eos_token_id)
+    model_ended_with_stop = model_prefix_token_ids[-1] in effective_stop_token_ids
 
     # Locate the turn boundary by EOS count rather than token position. Qwen3
     # templates may strip prior reasoning blocks when re-rendering history;
@@ -112,7 +112,8 @@ def replace_prefix_tokens(
         if tid == eos_token_id:
             count_seen += 1
             if count_seen == count_needed:
-                template_cut_start = pos
+                # Keep sampled stops; otherwise let the template add the missing EOS.
+                template_cut_start = pos + int(model_ended_with_stop)
                 break
 
     assert template_cut_start >= 0, (
@@ -124,6 +125,4 @@ def replace_prefix_tokens(
         f"Template repr (detokenized): {repr(tokenizer.decode(template_token_ids))}"
     )
 
-    return (
-        model_prefix_token_ids[:model_cut_end] + template_token_ids[template_cut_start:]
-    )
+    return model_prefix_token_ids + template_token_ids[template_cut_start:]

nemo_rl/models/generation/trtllm/trtllm_http_server.py

@@ -33,18 +33,31 @@
 logger = logging.getLogger(__name__)
 
 
-def _build_reasoning_parser(name: str, chat_template_kwargs: dict[str, Any]) -> Any:
-    from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory
+def _tokens_for_response_text(
+    token_ids: list[int], stop_token_ids: set[int]
+) -> list[int]:
+    text_token_end = len(token_ids)
+    while text_token_end and token_ids[text_token_end - 1] in stop_token_ids:
+        text_token_end -= 1
+    return token_ids if text_token_end == len(token_ids) else token_ids[:text_token_end]
 
-    if name == "deepseek-r1" and "enable_thinking" in chat_template_kwargs:
-        from tensorrt_llm.llmapi.reasoning_parser import DeepSeekR1Parser
 
-        return DeepSeekR1Parser(
-            reasoning_at_start=bool(chat_template_kwargs["enable_thinking"]),
-            chat_template_kwargs=chat_template_kwargs,
-        )
+def _build_reasoning_parser(
+    name: str,
+    chat_template_kwargs: dict[str, Any],
+    *,
+    reasoning_at_start: bool = False,
+) -> Any:
+    from tensorrt_llm.llmapi.reasoning_parser import (
+        DeepSeekR1Parser,
+        ReasoningParserFactory,
+    )
 
-    return ReasoningParserFactory.create_reasoning_parser(name, chat_template_kwargs)
+    parser = ReasoningParserFactory.create_reasoning_parser(name, chat_template_kwargs)
+    if isinstance(parser, DeepSeekR1Parser):
+        parser.reasoning_at_start = reasoning_at_start
+        parser.in_reasoning = reasoning_at_start
+    return parser
 
 
 def _build_sampling_params(
@@ -107,6 +120,10 @@ def create_app(
         **(default_chat_template_kwargs or {}),
     }
 
+    def _prompt_opens_reasoning(token_ids: list[int]) -> bool:
+        tail = tokenizer.decode(token_ids[-16:], skip_special_tokens=False)
+        return tail.rstrip().endswith("<think>")
+
     # Use the configured parser or infer one from the model config.
     _tool_parser_name = _resolve_tool_parser_name(tool_parser, model_name)
     _tool_parser_instance = _build_tool_parser(_tool_parser_name)
@@ -123,8 +140,7 @@ def create_app(
     if reasoning_parser is not None:
         _build_reasoning_parser(reasoning_parser, _server_template_kwargs)
 
-    # Match TRT-LLM's effective EOS set so this adapter can trim every returned
-    # stop token and its logprob together, preserving multi-turn continuity.
+    # Retain stop IDs in metadata but exclude them from parsed text.
     _eos_token_ids: set[int] = set(stop_token_ids or [])
 
     def _add_eos_token_ids(token_ids: Any) -> None:
@@ -170,14 +186,8 @@ async def chat_completions(request: Request):
         per_request_kwargs: dict[str, Any] = body.get("chat_template_kwargs") or {}
         effective_template_kwargs = {**_server_template_kwargs, **per_request_kwargs}
 
-        _active_reasoning_parser = (
-            _build_reasoning_parser(reasoning_parser, effective_template_kwargs)
-            if reasoning_parser is not None
-            else None
-        )
-
         try:
-            conversation, mm_coroutine, _ = parse_chat_messages_coroutines(
+            conversation, mm_coroutine, *_ = parse_chat_messages_coroutines(
                 messages, model_config
             )
             mm_data, mm_embeddings = await mm_coroutine
@@ -216,6 +226,21 @@ async def chat_completions(request: Request):
             model_prefix_token_ids=required_prefix_ids,
             template_prefix_token_ids=template_prefix_ids,
             template_token_ids=prompt_token_ids,
+            model_stop_token_ids=_eos_token_ids,
+        )
+
+        # Infer parser state from the exact engine prompt.
+        reasoning_at_start = reasoning_parser is not None and _prompt_opens_reasoning(
+            adj_prompt
+        )
+        _active_reasoning_parser = (
+            _build_reasoning_parser(
+                reasoning_parser,
+                effective_template_kwargs,
+                reasoning_at_start=reasoning_at_start,
+            )
+            if reasoning_parser is not None
+            else None
         )
 
         max_tokens_requested = (
@@ -272,14 +297,9 @@ async def chat_completions(request: Request):
                 else:
                     raise TypeError(f"Unsupported TRT-LLM logprob type: {type(lp)}")
 
-        # Strip trailing stop tokens TRT-LLM appends — apply_chat_template doesn't reproduce
-        # <|endoftext|>, so they'd break seen_token_ids contiguity. Trim logprobs in lockstep.
-        while gen_token_ids and gen_token_ids[-1] in _eos_token_ids:
-            gen_token_ids.pop()
-            if gen_logprobs:
-                gen_logprobs.pop()
-
-        gen_text = tokenizer.decode(gen_token_ids, skip_special_tokens=False)
+        # Parse without trailing stops; retain full response metadata.
+        text_token_ids = _tokens_for_response_text(gen_token_ids, _eos_token_ids)
+        gen_text = tokenizer.decode(text_token_ids, skip_special_tokens=False)
 
         finish_reason = "stop"
         if gen.finish_reason is not None:

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -397,6 +397,14 @@ def _setup_vllm_openai_api_server(self, app: FastAPI) -> FastAPI:
 
         engine_client = self.llm
         model_config = self.llm_async_engine_args.create_model_config()
+        model_stop_token_ids = set(self.cfg.get("stop_token_ids") or ())
+        generation_eos_token_ids = model_config.try_get_generation_config().get(
+            "eos_token_id"
+        )
+        if isinstance(generation_eos_token_ids, int):
+            model_stop_token_ids.add(generation_eos_token_ids)
+        elif generation_eos_token_ids is not None:
+            model_stop_token_ids.update(generation_eos_token_ids)
         base_model_paths = [
             BaseModelPath(
                 name=model_config.served_model_name, model_path=model_config.model
@@ -563,10 +571,10 @@ async def preprocess_chat(
                     model_prefix_token_ids=request.required_prefix_token_ids,
                     template_prefix_token_ids=actual_corresponding_token_ids,
                     template_token_ids=engine_prompt["prompt_token_ids"],
+                    model_stop_token_ids=model_stop_token_ids,
                 )
 
                 engine_prompt["prompt_token_ids"] = final_prompt_token_ids
-
                 # Clamp after prefix replacement since the prompt length may have changed.
                 if actual_request_max_tokens is not None:
                     self._clamp_max_tokens(
@@ -914,6 +922,10 @@ def _setup_vllm_server(self) -> "tuple[threading.Thread, str, uvicorn.Server]":
         # e.g. last-run middleware.
         app = FastAPI()
 
+        @app.get("/health")
+        async def health() -> dict[str, str]:
+            return {"status": "ok"}
+
         app = self._setup_vllm_openai_api_server(app)
         if self._sparse_refit_receiver is not None:
             self._sparse_refit_receiver.setup_api_server(app)

tests/test_suites/llm/grpo-qwen3-0.6b-1n4g-megatron-trtllm-tp2-noncolocated-gym.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source "$SCRIPT_DIR/common.env"
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=1
+GPUS_PER_NODE=4
+STEPS_PER_RUN=6
+MAX_STEPS=6
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))
+NUM_MINUTES=90
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+cd "$PROJECT_ROOT"
+
+DATA_DIR="$EXP_DIR/data"
+mkdir -p "$DATA_DIR"
+cd 3rdparty/Gym-workspace/Gym
+if [[ ! -f env.yaml ]]; then
+    if [[ -z "${HF_TOKEN:-}" ]]; then
+        echo "[ERROR] HF_TOKEN is not set"
+        exit 1
+    fi
+    echo "hf_token: $HF_TOKEN" >> env.yaml
+fi
+uv run ng_prepare_data "+config_paths=[resources_servers/workplace_assistant/configs/workplace_assistant.yaml]" \
+    +output_dirpath=data/workplace_assistant \
+    +mode=train_preparation \
+    +should_download=true \
+    +data_source=huggingface
+cd -
+
+TRAIN_PATH="$DATA_DIR/workplace_assistant_train.jsonl"
+VALIDATION_PATH="$DATA_DIR/workplace_assistant_validation.jsonl"
+jq -c '.responses_create_params.tools |= (.[0:1])' \
+    3rdparty/Gym-workspace/Gym/data/workplace_assistant/train.jsonl > "$TRAIN_PATH"
+jq -c '.responses_create_params.tools |= (.[0:1])' \
+    3rdparty/Gym-workspace/Gym/data/workplace_assistant/validation.jsonl > "$VALIDATION_PATH"
+
+uv run examples/nemo_gym/run_grpo_nemo_gym.py \
+    --config "$CONFIG_PATH" \
+    grpo.max_num_steps="$MAX_STEPS" \
+    logger.log_dir="$LOG_DIR" \
+    logger.wandb_enabled=true \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name="$EXP_NAME" \
+    logger.monitor_gpus=true \
+    logger.tensorboard_enabled=true \
+    checkpointing.enabled=true \
+    checkpointing.checkpoint_dir="$CKPT_DIR" \
+    data.train.data_path="$TRAIN_PATH" \
+    data.validation.data_path="$VALIDATION_PATH" \
+    "$@" \
+    2>&1 | tee "$RUN_LOG"
+
+uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS"
+
+LAST_STEP=$(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' "$JSON_METRICS")
+if [[ "$LAST_STEP" -lt "$MAX_STEPS" ]]; then
+    echo "[ERROR] Expected step $MAX_STEPS, reached $LAST_STEP"
+    exit 1
+fi
+
+uv run tests/check_metrics.py "$JSON_METRICS" \
+    'mean(data["train/reward"]) > 0.05' \
+    'median(data["train/token_mult_prob_error"]) < 1.1' \
+    'mean(data["train/gen_kl_error"]) < 0.02' \
+    'mean(data["train/grad_norm"], 2, 0) > 0.1' \
+    'mean(data["train/grad_norm"], 2, 0) < 30.0'
+
+rm -rf "$CKPT_DIR"

tests/test_suites/nightly_gb200.txt

@@ -20,6 +20,7 @@ tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_generation-noncolocated-m
 # TRT-LLM generation backend
 tests/test_suites/llm/grpo-qwen3-1.7b-2n4g-fsdp2-trtllm.sh
 tests/test_suites/llm/grpo-qwen2.5-0.5b-1n4g-megatron-trtllm-noncolocated-async.sh
+tests/test_suites/llm/grpo-qwen3-0.6b-1n4g-megatron-trtllm-tp2-noncolocated-gym.sh
 
 # Functional moonlight run
 tests/test_suites/llm/grpo-moonlight-16ba3b-4n4g-megatron.sh