feat(trtllm): support agentic training for SWE rollout

docs/guides/swe-rl-qwen3.md

@@ -176,6 +176,20 @@ uv run --frozen ./examples/nemo_gym/run_grpo_nemo_gym.py \
 
 Stage 2 uses a smaller global batch (`num_prompts_per_step=8`, `train_global_batch_size=64`) because each rollout is an expensive multi-turn trajectory (`agent_max_turns=200`, `swebench_agent_timeout=1800s`), and bumps Megatron TP to 4. `policy.model_name` points at the HF-converted SWE1 checkpoint.
 
+## [Alternative] Using TensorRT-LLM generation backend in Stage 2
+
+The recipe also supports running Stage 2 generation with [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) instead of vLLM. This can be useful when TensorRT-LLM's in-flight weight update path offers better throughput for a given cluster configuration.
+
+Use [`examples/swe_bench/grpo_qwen3_30b_async_swe_trtllm.yaml`](../../examples/swe_bench/grpo_qwen3_30b_async_swe_trtllm.yaml), which inherits the SWE2 training and environment settings and switches the generation backend. Set `container_formatter` in the YAML to your `.sif` image template (same as Stage 2 vLLM), then launch with your Stage 1 checkpoint and data:
+
+```bash
+uv run --frozen ./examples/nemo_gym/run_grpo_nemo_gym.py \
+  --config examples/swe_bench/grpo_qwen3_30b_async_swe_trtllm.yaml \
+  policy.model_name=/path/to/swe1_checkpoint_hf \
+  data.train.data_path=/path/to/data/swe2/train-split.jsonl \
+  data.validation.data_path=/path/to/data/swe2/val-split.jsonl
+```
+
 ## What to monitor
 
 - **`train/total_reward/mean`** — primary signal and the checkpointing metric for both stages.

docs/nsys-profiling.md

@@ -61,6 +61,7 @@ payload raises at startup so misconfiguration surfaces immediately.
 The supported worker types are:
 - **DTensorPolicyWorker**: Pattern matched against `"dtensor_policy_worker"`
 - **VllmGenerationWorker**: Pattern matched against `"vllm_generation_worker"`
+- **TrtllmAsyncGenerationWorker**: Pattern matched against `"trtllm_async_generation_worker"`
 
 ## Example Usage
 
@@ -91,6 +92,14 @@ LD_LIBRARY_PATH="/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:
 NRL_NSYS_PROFILE_STEP_RANGE=2:3 NRL_NSYS_WORKER_PATTERNS="megatron_policy_worker,vllm_generation_worker" uv run examples/run_grpo.py --config examples/configs/grpo_math_1B_megatron.yaml grpo.max_num_steps=5
 ```
 
+### Profile TRT-LLM Generation Workers
+
+```bash
+NRL_NSYS_PROFILE_STEP_RANGE=2:3 NRL_NSYS_WORKER_PATTERNS="trtllm_async_generation_worker" uv run examples/run_grpo.py --config examples/configs/grpo_math_1B_trtllm.yaml grpo.max_num_steps=5
+```
+
+The outer `TrtllmAsyncGenerationWorker` actor is CPU-only, so nsys wraps TRT-LLM's **internal** RayExecutor GPU workers instead. This is done by injecting `ray_worker_nsight_options` (with `capture-range=cudaProfilerApi`, deferred capture) into the `AsyncLLM` constructor. When `start_gpu_profiling()` is called, it broadcasts `collective_rpc("start_gpu_profiling")` to the internal GPU workers, each of which calls `torch.cuda.profiler.start()` to trigger the capture. Traces are one `.nsys-rep` per internal GPU worker (replicas × TP).
+
 ## Profile Output
 
 When profiling is enabled, it generates the following logs and files:
@@ -104,9 +113,10 @@ When profiling is enabled, it generates the following logs and files:
    ```
    dtensor_policy_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep
    vllm_generation_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep
+   trtllm_async_generation_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep
    worker_process_<PID>.nsys-rep
    ```
-If you are not using model parallelism in Vllm, you should directly refer to `vllm_generation_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep` for nsight reports; If you are using model parallelism, nsight is NOT applied to the outer `VllmGenerationWorker` to avoid interfering with Ray's compiled DAG. Instead, `ray_workers_use_nsight` is enabled and vLLM's default nsight config is monkey-patched to use `capture-range=cudaProfilerApi` (deferred capture). This means the internal TP workers run under nsys with near-zero overhead until `start_gpu_profiling()` triggers `cudaProfilerStart()` on each worker via `collective_rpc`. The `vllm_tp_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep` files are the nsight profiles from the internal TP workers. (refer to https://github.com/vllm-project/vllm/blob/7e3a8dc90670fd312ce1e0d4eba9bf11c571e3ad/vllm/executor/ray_distributed_executor.py#L136 for more information).
+For TRT-LLM, the meaningful generation profiles are the per-internal-GPU-worker files (`trtllm_async_generation_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep`), one per GPU (replicas × TP). If you are not using model parallelism in Vllm, you should directly refer to `vllm_generation_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep` for nsight reports; If you are using model parallelism, nsight is NOT applied to the outer `VllmGenerationWorker` to avoid interfering with Ray's compiled DAG. Instead, `ray_workers_use_nsight` is enabled and vLLM's default nsight config is monkey-patched to use `capture-range=cudaProfilerApi` (deferred capture). This means the internal TP workers run under nsys with near-zero overhead until `start_gpu_profiling()` triggers `cudaProfilerStart()` on each worker via `collective_rpc`. The `vllm_tp_worker_<NRL_NSYS_PROFILE_STEP_RANGE>_<PID>.nsys-rep` files are the nsight profiles from the internal TP workers. (refer to https://github.com/vllm-project/vllm/blob/7e3a8dc90670fd312ce1e0d4eba9bf11c571e3ad/vllm/executor/ray_distributed_executor.py#L136 for more information).
 
 3. **File Location**: Profile files are saved in `/tmp/ray/session*/logs/nsight/` directory on each worker node. Ensure you check both `ls /tmp/ray/session_[0-9]*/logs/nsight` and `ls /tmp/ray/session_latest/logs/nsight` for the profiles, since the "latest" pointer may be stale.
 

examples/swe_bench/grpo_qwen3_30b_async_swe_trtllm.yaml

@@ -0,0 +1,31 @@
+# Async GRPO SWE2 end-to-end agentic recipe for Qwen3-30B-A3B-Thinking with the TRT-LLM generation backend.
+# For the vLLM variant see examples/nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml.
+
+defaults: ../nemo_gym/grpo_qwen3_30ba3b_thinking_swe2.yaml
+
+policy:
+  generation:
+    backend: trtllm
+    stop_token_ids: [151643, 151645]
+    trtllm_cfg:
+      tensor_parallel_size: 2
+      precision: ${policy.precision}
+      max_model_len: ${policy.max_total_sequence_length}
+      max_batch_size: 256
+      max_num_tokens: ${policy.max_total_sequence_length}
+      gpu_memory_utilization: 0.7
+      async_engine: true
+      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
+      reasoning_parser: "deepseek-r1"
+      default_chat_template_kwargs:
+        enable_thinking: true
+        truncate_history_thinking: false
+
+logger:
+  wandb:
+    name: qwen3-30b-thinking-swe2-agentic-trtllm
+  mlflow:
+    experiment_name: qwen3-30b-thinking-swe2-agentic-trtllm
+    run_name: qwen3-30b-thinking-swe2-agentic-trtllm

nemo_rl/algorithms/grpo.py

@@ -864,6 +864,11 @@ def _spinup_nemo_gym(base_urls, model_name):
                     gpus_per_instance = vllm_cfg["tensor_parallel_size"] * vllm_cfg.get(
                         "pipeline_parallel_size", 1
                     )
+                elif generation_config["backend"] == "trtllm":
+                    trtllm_cfg = generation_config.get("trtllm_cfg", {})
+                    gpus_per_instance = trtllm_cfg[
+                        "tensor_parallel_size"
+                    ] * trtllm_cfg.get("pipeline_parallel_size", 1)
                 else:
                     sglang_cfg = generation_config.get("sglang_cfg", {})
                     gpus_per_instance = sglang_cfg.get("gpus_per_server", 1)
@@ -932,8 +937,12 @@ def _spinup_nemo_gym(base_urls, model_name):
             node_resource_constraints=inference_node_resource_constraints,
         )
         if inference_node_resource_constraints is not None:
-            VllmGeneration.init_cluster_placement_groups(
-                inference_cluster, generation_config
+            {
+                "vllm": VllmGeneration,
+                "trtllm": TrtllmGeneration,
+            }[generation_config["backend"]].init_cluster_placement_groups(
+                inference_cluster,
+                generation_config,
             )
         print(
             f"  ✓ Ray inference cluster initialized with {inference_nodes} nodes with {inference_gpus_per_node} GPUs per node",
@@ -1317,6 +1326,13 @@ def init_trtllm():
             flush=True,
         )
 
+        if enable_nemo_gym:
+            nemo_gym_actor, nemo_gym_time = _spinup_nemo_gym(
+                policy_generation.dp_openai_server_base_urls,
+                generation_config["model_name"],
+            )
+            worker_init_timing_metrics["nemo_gym_init_time_s"] = nemo_gym_time
+
     # Record when worker initialization completes (for calculating other setup time)
     worker_init_complete_time = time.perf_counter() - setup_start_time
 
@@ -2011,9 +2027,8 @@ def _should_use_nemo_gym(master_config: MasterConfig) -> bool:
             "expose_http_server"
         )
     elif generation_config["backend"] == "trtllm":
-        raise NotImplementedError(
-            "NeMo-Gym is not supported with the TRT-LLM generation backend "
-            "(the TRT-LLM OpenAI-compatible HTTP server was removed)."
+        should_expose_http_server = generation_config["trtllm_cfg"].get(
+            "expose_http_server"
         )
     else:
         should_expose_http_server = False

nemo_rl/data/llm_message_utils.py

@@ -374,7 +374,11 @@ def batched_message_log_to_flat_message(
         # Filter out None values and validate consistency
         values: list[Tensor | None] = cast(list[Tensor | None], values)
         tensors = cast(list[Tensor], [t for t in values if t is not None])
-        _validate_tensor_consistency(tensors)
+        try:
+            _validate_tensor_consistency(tensors)
+        except RuntimeError as e:
+            e.add_note(f"[key={key!r}]")
+            raise
 
         # Create zero tensors for None values
         filled_values: list[Tensor] = [

nemo_rl/distributed/virtual_cluster.py

@@ -259,9 +259,9 @@ def init_ray(log_dir: Optional[str] = None) -> None:
         if _k.startswith(("PMIX_", "PMI_", "MPI_", "OMPI_", "SLURM_")):
             os.environ.pop(_k, None)
 
-    # Set up runtime environment
     env_vars = dict(os.environ)
     env_vars.pop("RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", None)
+
     runtime_env = {
         "env_vars": env_vars,  # Pass thru all user environment variables
     }

nemo_rl/environments/nemo_gym.py

@@ -361,7 +361,11 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
             # Eventually we can maybe be smarter about this, but this is functional for now.
 
             # Note that NeMo-Gym will only return token ids on "assistant" messages and not other message types.
-            if "generation_token_ids" not in output_item_dict:
+            # Also skip if generation_token_ids is present but empty, e.g. all-EOS generation stripped to [] — torch.tensor([]) defaults to float32 and breaks batch dtype consistency.
+            if (
+                "generation_token_ids" not in output_item_dict
+                or not output_item_dict["generation_token_ids"]
+            ):
                 continue
 
             assert (

nemo_rl/models/generation/openai_server_utils.py

@@ -0,0 +1,125 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+"""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.
+"""
+
+from typing import Any
+
+
+def 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
+    (<think>) 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]
+    """
+    if not model_prefix_token_ids:
+        return template_token_ids
+
+    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
+
+    count_needed = template_prefix_token_ids.count(eos_token_id)
+    count_seen = 0
+    template_cut_start = -1
+    for pos, tid in enumerate(template_token_ids):
+        if tid == eos_token_id:
+            count_seen += 1
+            if count_seen == count_needed:
+                template_cut_start = pos
+                break
+
+    assert template_cut_start >= 0, (
+        f"EOS token #{count_needed} not found in template_token_ids "
+        f"(only found {count_seen} EOS tokens total)!\n"
+        f"Template prefix token IDs (everything before the final assistant message): {template_prefix_token_ids}\n\n"
+        f"Template token IDs (everything that was sent to the model endpoint): {template_token_ids}\n\n"
+        f"Template prefix repr (detokenized): {repr(tokenizer.decode(template_prefix_token_ids))}\n\n"
+        f"Template repr (detokenized): {repr(tokenizer.decode(template_token_ids))}"
+    )
+
+    return (
+        model_prefix_token_ids[:model_cut_end] + template_token_ids[template_cut_start:]
+    )

nemo_rl/models/generation/trtllm/config.py

@@ -19,11 +19,13 @@
 
 class TrtllmSpecificArgs(TypedDict):
     tensor_parallel_size: int
+    model_name: NotRequired[str]
     gpu_memory_utilization: NotRequired[float]
     max_model_len: int
     precision: str
     max_batch_size: int
     max_num_tokens: int
+    expose_http_server: NotRequired[bool]
     async_engine: NotRequired[bool]
     # MoE expert parallelism. TRT-LLM splits the TP dimension on MoE layers
     # into moe_tp × moe_ep, so the constraint is
@@ -41,6 +43,12 @@ class TrtllmSpecificArgs(TypedDict):
     # grpo.async_grpo so they cannot diverge).
     in_flight_weight_updates: NotRequired[bool]
     recompute_kv_cache_after_weight_updates: NotRequired[bool]
+    default_chat_template_kwargs: NotRequired[dict[str, Any]]
+    # TRT-LLM's registered parser names:
+    #   "qwen3"       -> Qwen3ToolParser      (JSON format: {"name":..., "arguments":{...}})
+    #   "qwen3_coder" -> Qwen3CoderToolParser  (XML format: <function=...>)
+    tool_parser: NotRequired[str]
+    reasoning_parser: NotRequired[str]
 
 
 class TrtllmConfig(GenerationConfig):