When multi modal input for trtllm, decode with special token first

verl/experimental/agent_loop/agent_loop.py

@@ -113,7 +113,7 @@ async def generate(
         server = self._choose_server(request_id)
         output = await server.generate.remote(
             request_id=uuid4().hex,  # use new request_id for each turn
-            prompt_ids=prompt_ids,
+            prompt_ids=prompt_ids, # for trtllm, this is the raw prompt
             sampling_params=sampling_params,
             image_data=image_data,
             video_data=video_data,

verl/workers/rollout/trtllm_rollout/trtllm_async_server.py

@@ -184,7 +184,7 @@ async def launch_server(self):
 
     async def generate(
         self,
-        prompt_ids: list[int],
+        prompt_ids: str,
         sampling_params: dict[str, Any],
         request_id: str,
         image_data: Optional[list[Any]] = None,
@@ -201,16 +201,19 @@ async def generate(
 
         trt_llm_sampling_params = SamplingParams(**sampling_params)
         if self.is_vlm_model:
+            org_prompt = self.llm.tokenizer.decode(prompt_ids)
             if image_data or video_data:
+
                 input_dict = {
-                    "prompt_token_ids": prompt_ids,
+                    "prompt": org_prompt,
                     "multi_modal_data": {},
                     "mm_processor_kwargs": {},
                 }
                 if image_data:
                     input_dict["multi_modal_data"]["image"] = image_data
                 if video_data:
                     input_dict["multi_modal_data"]["video"] = video_data
+
                 outputs = await self.llm.generate_async(
                     inputs=input_dict,
                     sampling_params=trt_llm_sampling_params,
@@ -369,7 +372,7 @@ async def launch_servers(self):
                 node_id=node_id,
                 soft=False,
             ),
-            runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}},
+            runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", "TLLM_NUMA_AWARE_WORKER_AFFINITY":"0"}},
             name=name,
         ).remote(
             config=self.config,

verl/workers/rollout/trtllm_rollout/trtllm_rollout.py

@@ -414,7 +414,7 @@ async def update_weights(self, weights: Generator[tuple[str, torch.Tensor], None
         total_available_bytes = int(self.config.checkpoint_engine.update_weights_bucket_megabytes) * 1024 * 1024
 
         try:
-            device_uuid = get_device_uuid(self.gpu_id)
+            device_uuid = get_device_uuid(int(self.gpu_id))
         except Exception as e:
             logger.error(f"Failed to get device UUID in update_weights(): {e}")
             device_uuid = None

verl/workers/rollout/trtllm_rollout/trtllm_worker_extension.py

@@ -62,7 +62,7 @@ def update_weights(self, ipc_handles: Optional[dict] = None):
                     # using restricted unpickler from tensorrt_llm.serialization
                     logger.info("Deserializing base64-encoded weight handles")
                     decoded_data = base64.b64decode(serialized_handles)
-                    # Allow basic builtins and all torch modules
+                    # Allow basic builtins and torch tensor reconstruction classes
                     approved_imports = {
                         "builtins": [
                             "list",
@@ -76,11 +76,47 @@ def update_weights(self, ipc_handles: Optional[dict] = None):
                             "NoneType",
                             "type",
                         ],
+                        "torch": [
+                            "Tensor",
+                            "FloatTensor",
+                            "DoubleTensor",
+                            "HalfTensor",
+                            "BFloat16Tensor",
+                            "IntTensor",
+                            "LongTensor",
+                            "ShortTensor",
+                            "CharTensor",
+                            "ByteTensor",
+                            "BoolTensor",
+                            "Size",
+                            "dtype",
+                            "device",
+                            "float32",
+                            "float16",
+                            "int32",
+                            "int64",
+                            "int16",
+                            "int8",
+                            "uint8",
+                            "bool",
+                        ],
+                        "torch.multiprocessing.reductions": [
+                            "rebuild_cuda_tensor",
+                            "rebuild_tensor",
+                        ],
+                        "torch._utils": [
+                            "_rebuild_tensor_v2",
+                        ],
+                        "torch.storage": [
+                            "_load_from_bytes",
+                            "_TypedStorage",
+                            "UntypedStorage",
+                            "TypedStorage",
+                        ],
                     }
                     all_handles = serialization.loads(
                         decoded_data,
                         approved_imports=approved_imports,
-                        approved_module_patterns=[r"^torch.*"],
                     )
 
                     # Verify the result is a list as expected