[trtllm] fix: Fixes for TRTLLM rollout (#5032)

### What does this PR do?

Several fixes/improvements for trtllm rollout:
- Remove dependency on `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES` for
before Ray cluster starts.
- Fix memory issue during IPC handle creation in `update_weights()`.
- Fix reward model path (Note: these codes are for future use when
TRTLLM rollout switches from `TRTLLMSampler` to `TorchSampler`.
Currently, TRTLLM rollout still doesn't support training involving
reward models. )

### Checklist Before Starting

- [ ] Search for similar PRs. Paste at least one query link here: ...
- [ ] Format the PR title as `[{modules}] {type}: {description}` (This
will be checked by the CI)
- `{modules}` include `fsdp`, `megatron`, `veomni`, `sglang`, `vllm`,
`rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`,
`deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`,
`model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`
- If this PR involves multiple modules, separate them with `,` like
`[megatron, fsdp, doc]`
  - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test`
- If this PR breaks any API (CLI arguments, config, function signature,
etc.), add `[BREAKING]` to the beginning of the title.
  - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching`

### Test

> For changes that can not be tested by CI (e.g., algorithm
implementation, new model support), validate by experiment(s) and show
results like training curve plots, evaluation results, etc.

### API and Usage Example

> Demonstrate how the API changes if any, and provide usage example(s)
if possible.

```python
# Add code snippet or script demonstrating how to use this
```

### Design & Code Changes

> Demonstrate the high-level design if this PR is complex, and list the
specific changes.

### Checklist Before Submitting

> [!IMPORTANT]
> Please check all the following items before requesting a review,
otherwise the reviewer might deprioritize this PR for review.

- [x] Read the [Contribute
Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).
- [x] Apply [pre-commit
checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting):
`pre-commit install && pre-commit run --all-files --show-diff-on-failure
--color=always`
- [x] Add / Update [the
documentation](https://github.com/volcengine/verl/tree/main/docs).
- [x] Add unit or end-to-end test(s) to [the CI
workflow](https://github.com/volcengine/verl/tree/main/.github/workflows)
to cover all the code. If not feasible, explain why: ...
- [x] Once your PR is ready for CI, send a message in [the `ci-request`
channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the
`verl` Slack
workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ).
(If not accessible, please try [the Feishu group
(飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)
- [ ] If your PR is related to the `recipe` submodule, please also
update the reference to the submodule commit via `git submodule update
--remote` or `cd recipe && git pull origin main`.

---------

Co-authored-by: Liwei Ma <liweim@nvidia.com>

Author: hchings

.github/workflows/e2e_ppo_grpo_trainer_trtllm.yml

@@ -116,7 +116,6 @@ jobs:
       NO_PROXY: "localhost,127.0.0.1,hf-mirror.com"
       HF_ENDPOINT: "https://hf-mirror.com"
       HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable
-      RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES: "1"
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
@@ -152,7 +151,6 @@ jobs:
       NO_PROXY: "localhost,127.0.0.1,hf-mirror.com"
       HF_ENDPOINT: "https://hf-mirror.com"
       HF_HUB_ENABLE_HF_TRANSFER: "0" # This is more stable
-      RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES: "1"
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:

docs/workers/trtllm_worker.rst

@@ -36,9 +36,6 @@ Install verl with TensorRT-LLM:
             unset "$v"
         done
 
-        # Required for IPC UUID detection
-        export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1
-
 Using TensorRT-LLM as the Rollout Engine for GRPO
 -------------------------------------------------
 

examples/grpo_trainer/run_qwen2-7b_math_megatron_trtllm.sh

@@ -7,7 +7,6 @@ for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do
     unset "$v"
 done
 
-export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1
 export RAY_DEDUP_LOGS=0
 
 # -----

examples/grpo_trainer/run_qwen2-7b_math_trtllm.sh

@@ -5,7 +5,6 @@ for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do
     unset "$v"
 done
 
-export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1
 export RAY_DEDUP_LOGS=0
 
 # -----

tests/special_sanity/check_license.py

@@ -17,6 +17,7 @@
 
 license_head_bytedance = "Copyright 2024 Bytedance Ltd. and/or its affiliates"
 license_head_bytedance_25 = "Copyright 2025 Bytedance Ltd. and/or its affiliates"
+license_head_bytedance_26 = "Copyright 2026 Bytedance Ltd. and/or its affiliates"
 # Add custom license headers below
 license_head_prime = "Copyright 2024 PRIME team and/or its affiliates"
 license_head_individual = "Copyright 2025 Individual Contributor:"
@@ -29,6 +30,7 @@
 license_headers = [
     license_head_bytedance,
     license_head_bytedance_25,
+    license_head_bytedance_26,
     license_head_prime,
     license_head_individual,
     license_head_sglang,

verl/experimental/reward_loop/reward_loop.py

@@ -218,6 +218,23 @@ async def compute_score_disrm(self, data: DataProto) -> dict:
             }
             output = await self._post_request(payloads, "v1/embeddings")
             rm_score = output["data"][-1]["embedding"][-1]
+        elif engine_name == "trtllm":
+            # TODO: remove this once TRT-LLM switches to TorchSampler
+            raise ValueError("TensorRT-LLM backend does not support reward models currently.")
+
+            payloads = {
+                "model": model_name,
+                "prompt": disrm_prompt,
+                "return_context_logits": True,
+            }
+            output = await self._post_request(payloads, "v1/completions")
+            rm_score = output["choices"][0]["context_logits"]
+            assert isinstance(rm_score, list) and len(rm_score) > 0, (
+                "TensorRT-LLM OpenAI server response for reward score is not in the expected format."
+            )
+
+            rm_score = float(rm_score[0][0])
+            logger.debug(f"rm score: {rm_score}")
         else:
             raise NotImplementedError(f"RewardLoopManager does not support {engine_name}")
 

verl/workers/rollout/trtllm_rollout/trtllm_async_rollout.md

@@ -23,9 +23,6 @@ Note that using the TRT-LLM rollout requires setting the following environment v
 for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do
     unset "$v"
 done
-
-# Required for IPC UUID detection
-export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1
 ```
 
 ## 2. Architecture Design

verl/workers/rollout/trtllm_rollout/trtllm_async_server.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -17,14 +17,14 @@
 from typing import Any, Optional
 
 import ray
+import torch
 from omegaconf import DictConfig
 from ray.actor import ActorHandle
 from ray.util import placement_group_table
 from ray.util.placement_group import PlacementGroup
 
 from verl.single_controller.ray import RayClassWithInitArgs, SubRayResourcePool
 from verl.utils.config import omega_conf_to_dataclass
-from verl.utils.device import is_cuda_available
 from verl.utils.net_utils import is_valid_ipv6_address
 from verl.workers.config import HFModelConfig, RolloutConfig
 from verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput
@@ -42,6 +42,7 @@ class TRTLLMHttpServer:
     Args:
         config (DictConfig): full config.
         model_config (HFModelConfig): model config.
+        is_reward_model (bool): whether this is a reward model.
         rollout_mode (RolloutMode): rollout mode.
         workers (list[ActorHandle]): list of rollout workers.
         replica_rank (int): replica rank, a replica may contain multiple nodes.
@@ -54,6 +55,7 @@ def __init__(
         self,
         config: RolloutConfig,
         model_config: HFModelConfig,
+        is_reward_model: bool,
         rollout_mode: RolloutMode,
         workers: list[ActorHandle],
         replica_rank: int,
@@ -62,10 +64,11 @@ def __init__(
         bundle_indices: list[list[int]] = None,
     ):
         os.environ["TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL"] = "1"
-        assert is_cuda_available, "TRTLLM http server should run on GPU node"
+        assert torch.cuda.is_available(), "TRTLLM http server should run on GPU node"
 
         self.config: RolloutConfig = omega_conf_to_dataclass(config)
         self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)
+        self.is_reward_model = is_reward_model
         self.config.max_model_len = self.config.prompt_length + self.config.response_length
         self.rollout_mode = rollout_mode
         self.workers = workers
@@ -82,7 +85,7 @@ def __init__(
         self._server_address = ray.util.get_node_ip_address().strip("[]")
         self._server_port = None
 
-        logger.info(f"TRTLLMHttpServer, replica_rank: {self.replica_rank}, ")
+        logger.info(f"TRTLLMHttpServer, replica_rank: {self.replica_rank}")
 
         self.sampling_args = {
             "detokenize": False,
@@ -107,11 +110,6 @@ async def launch_server(self):
             enable_block_reuse=True,
             free_gpu_memory_fraction=self.config.gpu_memory_utilization,
         )
-        cuda_graph_config = CudaGraphConfig(
-            enable_padding=True,
-            batch_sizes=self.config.cudagraph_capture_sizes,
-            max_batch_size=0 if self.config.cudagraph_capture_sizes else self.config.max_num_seqs,
-        )
 
         per_worker_gpu_share = 1.0 / self.max_colocate_count
 
@@ -121,7 +119,6 @@ async def launch_server(self):
             "orchestrator_type": "ray",
             "ray_worker_extension_cls": "tensorrt_llm.llmapi.rlhf_utils.WorkerExtension",
             "kv_cache_config": kv_cache_config,
-            "cuda_graph_config": cuda_graph_config,
             "max_seq_len": self.config.max_model_len,
             "max_batch_size": self.config.max_num_seqs,
             "max_num_tokens": self.config.max_num_batched_tokens,
@@ -136,6 +133,24 @@ async def launch_server(self):
             **engine_kwargs,
         }
 
+        if self.is_reward_model:
+            llm_kwargs.update(
+                {
+                    "cuda_graph_config": None,
+                    "disable_overlap_scheduler": True,
+                }
+            )
+        else:
+            llm_kwargs.update(
+                {
+                    "cuda_graph_config": CudaGraphConfig(
+                        enable_padding=True,
+                        batch_sizes=self.config.cudagraph_capture_sizes,
+                        max_batch_size=0 if self.config.cudagraph_capture_sizes else self.config.max_num_seqs,
+                    )
+                }
+            )
+
         self.llm = await AsyncLLM(**llm_kwargs)
 
         trtllm_server = OpenAIServer(
@@ -313,6 +328,7 @@ async def launch_servers(self):
         ).remote(
             config=self.config,
             model_config=self.model_config,
+            is_reward_model=self.is_reward_model,
             rollout_mode=self.rollout_mode,
             workers=self.workers,
             replica_rank=self.replica_rank,

verl/workers/rollout/trtllm_rollout/trtllm_rollout.py

@@ -1,4 +1,4 @@
-# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -31,6 +31,7 @@
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.multiprocessing.reductions import reduce_tensor
 
+from verl.utils.memory_utils import aggressive_empty_cache
 from verl.utils.net_utils import is_valid_ipv6_address
 from verl.workers.config import HFModelConfig, RolloutConfig
 from verl.workers.rollout.base import BaseRollout
@@ -46,21 +47,6 @@
 DEFAULT_MAX_WAIT_TIME = 300.0
 
 
-def device_id_to_physical_device_id(id: int) -> int:
-    """Convert a logical device ID to a physical device ID considering CUDA_VISIBLE_DEVICES."""
-    if "CUDA_VISIBLE_DEVICES" in os.environ:
-        device_ids = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
-        try:
-            physical_device_id = int(device_ids[id])
-            return physical_device_id
-        except (ValueError, IndexError) as err:
-            raise RuntimeError(
-                f"Failed to convert logical device ID {id} to physical device ID. Available devices are: {device_ids}."
-            ) from err
-    else:
-        return id
-
-
 @contextlib.contextmanager
 def nvml_context():
     """Context manager for NVML initialization and shutdown.
@@ -95,24 +81,19 @@ def get_device_uuid(id: int) -> str:
             except pynvml.NVMLError as e:
                 raise RuntimeError(f"Failed to initialize NVML: {e}") from e
 
-    # The process has visibility to all GPUs within the TP group
-    global_device_idx = device_id_to_physical_device_id(id)
-
     # Get the device handle and UUID
     try:
-        handle = pynvml.nvmlDeviceGetHandleByIndex(global_device_idx)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(id)
         uuid = pynvml.nvmlDeviceGetUUID(handle)
         # Ensure the UUID is returned as a string, not bytes
         if isinstance(uuid, bytes):
             return uuid.decode("utf-8")
         elif isinstance(uuid, str):
             return uuid
         else:
-            raise RuntimeError(
-                f"Unexpected UUID type: {type(uuid)} for device {id} (global index: {global_device_idx})"
-            )
+            raise RuntimeError(f"Unexpected UUID type: {type(uuid)} for device {id} (global index: {id})")
     except pynvml.NVMLError as e:
-        raise RuntimeError(f"Failed to get device UUID for device {id} (global index: {global_device_idx}): {e}") from e
+        raise RuntimeError(f"Failed to get device UUID for device {id} (global index: {id}): {e}") from e
 
 
 async def _read_async_response(resp: aiohttp.ClientResponse) -> dict[str, Any]:
@@ -330,8 +311,6 @@ def __init__(
         assert self.replica_rank >= 0, "replica_rank is not set"
         assert self.is_leader_rank is not None, "is_leader_rank is not set"
 
-        print(f"ServerAdapter, replica_rank: {self.replica_rank}, is_leader_rank: {self.is_leader_rank}")
-
         self.node_ip = ray.util.get_node_ip_address().strip("[]")
 
     async def _init_server_adapter(self):
@@ -412,8 +391,7 @@ async def update_weights(self, weights: Generator[tuple[str, torch.Tensor], None
         try:
             device_uuid = get_device_uuid(self.gpu_id)
         except Exception as e:
-            logger.error(f"Failed to get device UUID: {e}")
-            logger.error("Did you miss to set RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1 before ray start?")
+            logger.error(f"Failed to get device UUID in update_weights(): {e}")
             device_uuid = None
             raise e
 
@@ -447,3 +425,4 @@ async def flush():
             # Finalize update weights
             await self._adapter.update_weights(None)
         await asyncio.to_thread(dist.barrier, group=self.hybrid_device_mesh["exclude_dp"].get_group())
+        aggressive_empty_cache(force_sync=False)