Fix colocate mode

trinity/common/config.py

@@ -730,6 +730,8 @@ class ExplorerConfig:
     """Config for explorer."""
 
     name: str = EXPLORER_NAME
+    # ! DO NOT SET, automatically copied from Config.mode
+    mode: str = "both"
     # for workflow runner
     # number of workflow runners.
     runner_per_model: int = 8  # number of runners per each rollout model

trinity/common/config_validator.py

@@ -253,6 +253,8 @@ def _set_gpu_allocation_info(self, config: Config) -> None:
                 )
         elif config.mode == "colocate":
             self.logger.warning("`colocate` is only for single GPU scenario.")
+            if not config.explorer.rollout_model.engine_type.startswith("vllm"):
+                raise ValueError("Colocate mode only supports vLLM.")
             if cluster.total_gpu_num != 1:
                 raise ValueError(
                     f"Colocate mode requires exactly 1 GPU, but got {cluster.total_gpu_num} GPUs. Please use `both` mode instead."
@@ -689,6 +691,7 @@ def validate(self, config: Config) -> None:
         config.explorer.rollout_model.enable_return_routed_experts = (
             config.algorithm.enable_router_replay
         )
+        config.explorer.mode = config.mode
         config.explorer.rollout_model.ray_namespace = config.ray_namespace
         config.explorer.rollout_model.sync_method = config.synchronizer.sync_method
         config.explorer.rollout_model.checkpoint_job_dir = config.checkpoint_job_dir

trinity/common/models/allocator.py

@@ -41,6 +41,11 @@ def get_actor_name(self, role: str, engine_id: int, node_id: int) -> str:
         """Generate a unique actor name based on the model config, engine ID, and node ID."""
         return f"{self.config.name}_{role}_model_{engine_id}_{node_id}"
 
+    @property
+    def placement_group_name(self) -> str:
+        """Generate a descriptive name for the inference model placement group."""
+        return f"{self.config.name}_rollout_model"
+
     def allocate_bundles(self) -> BundleResult:
         """Allocate bundles for the rollout model and auxiliary models based on the configuration."""
         rollout_model = self.config.rollout_model
@@ -112,8 +117,15 @@ async def create_engine(
 
     async def create_all_models(self) -> Tuple[List[ModelWrapper], List[List[ModelWrapper]]]:
         """Create all model actors for the rollout model and auxiliary models based on the configuration."""
+        if self.config.mode == "colocate":
+            return await self._create_colocate_model()
+
         self.bundle_result = self.allocate_bundles()
-        self.pg = placement_group(self.bundle_result.bundles, strategy="PACK")
+        self.pg = placement_group(
+            self.bundle_result.bundles,
+            strategy="PACK",
+            name=self.placement_group_name,
+        )
         await self.pg.ready()
         self.analyze_placement_group(self.pg, self.bundle_result)
         # create rollout_models
@@ -147,6 +159,47 @@ async def create_all_models(self) -> Tuple[List[ModelWrapper], List[List[ModelWr
         ]
         return rollout_models, auxiliary_models
 
+    async def _create_colocate_model(
+        self,
+    ) -> Tuple[List[ModelWrapper], List[List[ModelWrapper]]]:
+        """Create the rollout model without reserving a placement group.
+
+        In colocate mode the trainer and rollout model share the same GPU. Reserving
+        that GPU in a placement group would prevent the trainer from acquiring it.
+        """
+        config = deepcopy(self.config.rollout_model)
+        config.engine_id = 0
+        actor_name = self.get_actor_name("rollout", 0, 0)
+        config.ray_actor_name = actor_name
+
+        if not config.engine_type.startswith("vllm"):
+            raise ValueError(
+                f"Colocate mode only supports vLLM, but got engine type: {config.engine_type}"
+            )
+
+        from trinity.common.models.vllm_model import vLLMRolloutModel
+
+        self.logger.info(
+            "Creating colocated inference_model %s in %s.",
+            actor_name,
+            config.ray_namespace,
+        )
+        handler = (
+            ray.remote(vLLMRolloutModel)
+            .options(
+                name=actor_name,
+                num_cpus=0,
+                num_gpus=0,
+                namespace=config.ray_namespace,
+            )
+            .remote(config=config)
+        )
+        await handler.prepare.remote()
+        server_address = await handler.get_api_server_url.remote()
+        wrapper = ModelWrapper(models=[handler], config=config, api_address=server_address)
+        await wrapper.prepare()
+        return [wrapper], []
+
     def get_model(self, config: InferenceModelConfig, role: str, engine_id: int) -> ModelWrapper:
         """Get the model actor for the given role and engine ID."""
         actor_name = self.get_actor_name(role, engine_id, 0)