Fix partial load problem, Add vlm support for trtllm rollout

Author: SchumiDing

verl/workers/rollout/trtllm_rollout/trtllm_async_server.py

@@ -125,7 +125,7 @@ async def launch_server(self):
             "model": self.model_config.local_path,
             "backend": "pytorch",
             "orchestrator_type": "ray",
-            "ray_worker_extension_cls": "tensorrt_llm.llmapi.rlhf_utils.WorkerExtension",
+            "ray_worker_extension_cls": "verl.workers.rollout.trtllm_rollout.trtllm_worker_extension.WorkerExtension",
             "kv_cache_config": kv_cache_config,
             "max_seq_len": self.config.max_model_len,
             "max_batch_size": self.config.max_num_seqs,
@@ -159,18 +159,45 @@ async def launch_server(self):
                 }
             )
 
-        self.llm = await AsyncLLM(**llm_kwargs)
-
-        trtllm_server = OpenAIServer(
-            llm=self.llm,
-            model=self.model_config.local_path,
-            tool_parser=None,
-            server_role=None,
-            metadata_server_cfg=None,
-        )
+        if self.is_vlm_model:
+            from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
+            multimodal_config = MultimodalServerConfig(
+                media_io_kwargs={
+                    "image": {
+                        "format": "pil",
+                        "device": "cpu",
+                    },
+                    "video": {
+                        "num_frames": 8,
+                        "fps": 30,
+                        "format": "pil",
+                        "device": "cpu",
+                    },
+                }
+            )
+            self.llm = await AsyncLLM(**llm_kwargs)
+            trtllm_server = OpenAIServer(
+                llm=self.llm,
+                model=self.model_config.local_path,
+                tool_parser=None,
+                server_role=None,
+                metadata_server_cfg=None,
+                multimodal_server_config=multimodal_config,
+            )
+        else:
+            self.llm = await AsyncLLM(**llm_kwargs)
+            trtllm_server = OpenAIServer(
+                llm=self.llm,
+                model=self.model_config.local_path,
+                tool_parser=None,
+                server_role=None,
+                metadata_server_cfg=None,
+            )
+        
         app = trtllm_server.app
         self._server_port, self._server_task = await run_unvicorn(app, None, self._server_address)
 
+    @resume_on_abort
     async def generate(
         self,
         prompt_ids: list[int],
@@ -179,11 +206,7 @@ async def generate(
         image_data: Optional[list[Any]] = None,
         video_data: Optional[list[Any]] = None,
     ) -> TokenOutput:
-        """Generate sequence with token-in-token-out."""
-        assert image_data is None and video_data is None, "Multimodality is not yet supported in TRTLLMHttpServer."
-
         from tensorrt_llm.llmapi import SamplingParams
-
         max_tokens = min(self.config.response_length, self.config.max_model_len - len(prompt_ids))
         sampling_params["max_tokens"] = max_tokens
         sampling_params["logprobs"] = 1 if sampling_params.pop("logprobs", False) else None
@@ -192,15 +215,34 @@ async def generate(
         sampling_params.update(self.sampling_args)
 
         trt_llm_sampling_params = SamplingParams(**sampling_params)
-        outputs = await self.llm.generate_async(
-            inputs=prompt_ids,
-            sampling_params=trt_llm_sampling_params,
-        )
-
+        if self.is_vlm_model:
+            if image_data or video_data:
+                input_dict = {
+                    "prompt_token_ids": prompt_ids,
+                    "multi_modal_data": {},
+                }
+                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,
+                )
+            else:
+                outputs = await self.llm.generate_async(
+                    inputs=prompt_ids,
+                    sampling_params=trt_llm_sampling_params,
+                )
+        else:
+            outputs = await self.llm.generate_async(
+                inputs=prompt_ids,
+                sampling_params=trt_llm_sampling_params,
+            )
         token_ids = outputs.outputs[0].token_ids
         log_probs = None
-        if trt_llm_sampling_params.logprobs is not None:
-            log_probs = [list(d.values())[0].logprob for d in outputs.outputs[0].logprobs]
+        if outputs.outputs[0].logprobs is not None:
+            log_probs = [logprobs[token_ids[i]].logprob for i, logprobs in enumerate(outputs.outputs[0].logprobs)]
         return TokenOutput(token_ids=token_ids, log_probs=log_probs)
 
     async def wake_up(self):

verl/workers/rollout/trtllm_rollout/trtllm_rollout.py

@@ -281,6 +281,7 @@ def __init__(
         self.is_leader_rank = None
         self.replica_rank = None
         self.is_dp_rank = None
+        self._supports_partial_loading = None  
 
         # hybrid mode
         if self.device_mesh is not None:
@@ -312,6 +313,21 @@ def __init__(
 
         self.node_ip = ray.util.get_node_ip_address().strip("[]")
 
+    async def get_supports_partial_loading(self) -> bool:
+        """Query and cache whether the model supports partial weight loading."""
+        if self._supports_partial_loading is not None:
+            return self._supports_partial_loading
+
+        await self._init_server_adapter()
+        try:
+            self._supports_partial_loading = await self.server_actor.supports_partial_loading.remote()
+        except Exception as e:
+            logger.warning(f"Failed to query partial loading support: {e}, defaulting to False")
+            self._supports_partial_loading = False
+
+        logger.info(f"Model supports partial loading: {self._supports_partial_loading}")
+        return self._supports_partial_loading
+
     async def _init_server_adapter(self):
         if self._adapter is not None:
             return
@@ -405,16 +421,21 @@ async def flush():
             await self.update_weights_from_ipc_handles(serialized_device_handles)
             cur_available_bytes = total_available_bytes
             cur_handles = []
+            
+        # Query if model supports partial loading
+        supports_partial_loading = await self.get_supports_partial_loading()
 
         for name, param in weights:
-            size_in_bytes = param.element_size() * param.numel()
-            if size_in_bytes > cur_available_bytes:
-                await flush()
+            if supports_partial_loading:
+                size_in_bytes = param.element_size() * param.numel()
+                if size_in_bytes > cur_available_bytes:
+                    await flush()
+
+                assert cur_available_bytes >= size_in_bytes, (
+                    f"cur_available_bytes: {cur_available_bytes:,} size_in_bytes: {size_in_bytes:,} name: {name}"
+                )
+                cur_available_bytes -= size_in_bytes
 
-            assert cur_available_bytes >= size_in_bytes, (
-                f"cur_available_bytes: {cur_available_bytes:,} size_in_bytes: {size_in_bytes:,} name: {name}"
-            )
-            cur_available_bytes -= size_in_bytes
             handle = reduce_tensor(param.detach())
             cur_handles.append((name, handle))
 

verl/workers/rollout/trtllm_rollout/trtllm_worker_extension.py

@@ -0,0 +1,77 @@
+import base64
+import inspect
+import pickle
+from typing import Optional
+
+from tensorrt_llm._ray_utils import control_action_decorator
+from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import MoeLoadBalancer
+from tensorrt_llm._torch.utils import get_device_uuid
+from tensorrt_llm.logger import logger
+
+
+class WorkerExtension:
+
+    def __init__(self):
+        pass
+
+    @control_action_decorator
+    def supports_partial_loading(self) -> bool:
+        """Check if the model supports partial weight loading."""
+        try:
+            model = self.engine.model_engine.model
+            load_weights_args = inspect.getfullargspec(model.load_weights).args
+            return "allow_partial_loading" in load_weights_args
+        except Exception as e:
+            logger.warning(f"Failed to check partial loading support: {e}")
+            return False
+
+    @control_action_decorator
+    def update_weights(self, ipc_handles: Optional[dict] = None):
+        try:
+            if not hasattr(self.engine.model_engine.model, "first_pre_reload_weights"):
+                for module in self.engine.model_engine.model.modules():
+                    if hasattr(module, "pre_reload_weights") and not getattr(
+                        module, "_weights_removed", False
+                    ):
+                        module.pre_reload_weights()
+                setattr(self.engine.model_engine.model, "first_pre_reload_weights", True)
+
+            if ipc_handles is not None:
+                device_uuid = get_device_uuid()
+                handles = ipc_handles.get(device_uuid, None)
+                if handles is not None:
+                    weights = pickle.loads(base64.b64decode(handles))
+                    model = self.engine.model_engine.model
+                    load_weights_args = inspect.getfullargspec(model.load_weights).args
+                    supports_partial_loading = "allow_partial_loading" in load_weights_args
+
+                    if supports_partial_loading:
+                        self.engine.model_engine.model_loader.reload(
+                            model, weights, allow_partial_loading=True
+                        )
+                    else:
+                        self.engine.model_engine.model_loader.reload(
+                            model, weights, allow_partial_loading=False
+                        )
+            else:
+                for module in self.engine.model_engine.model.modules():
+                    if hasattr(module, "process_weights_after_loading") and not getattr(
+                        module, "_weights_removed", False
+                    ):
+                        module.process_weights_after_loading()
+                    if hasattr(module, "post_load_weights") and not getattr(
+                        module, "_weights_removed", False
+                    ):
+                        module.post_load_weights()
+                moe_load_balancer = getattr(self.engine.model_engine, "moe_load_balancer", None)
+                if isinstance(moe_load_balancer, MoeLoadBalancer):
+                    moe_load_balancer.register_weight_slots_after_to_cuda()
+                    logger.info("moe_load_balancer finalizing model...")
+                    moe_load_balancer.finalize_model()
+                    logger.info("moe_load_balancer finalize model done")
+                self.engine.reset_prefix_cache()
+                delattr(self.engine.model_engine.model, "first_pre_reload_weights")
+
+        except Exception as e:
+            logger.error("Encountered an error in update_weights")
+            raise e