From 90b4a415e7440fab01e0be345b6a31801df8e10f Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Sun, 12 Jul 2026 22:12:07 +0800 Subject: [PATCH 1/4] fix colocate --- trinity/common/config.py | 2 ++ trinity/common/config_validator.py | 1 + trinity/common/models/allocator.py | 52 ++++++++++++++++++++++++++++++ 3 files changed, 55 insertions(+) diff --git a/trinity/common/config.py b/trinity/common/config.py index 9d0ca6382a5..6d27cb3af87 100644 --- a/trinity/common/config.py +++ b/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 diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index d4282d8309a..a5bb3a81e08 100644 --- a/trinity/common/config_validator.py +++ b/trinity/common/config_validator.py @@ -689,6 +689,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 diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 551ab09641d..23dead8d1cc 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -112,6 +112,9 @@ 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") await self.pg.ready() @@ -147,6 +150,55 @@ 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 config.engine_type.startswith("vllm"): + from trinity.common.models.vllm_model import vLLMRolloutModel + + model_cls = vLLMRolloutModel + num_gpus = 0 + elif config.engine_type == "sglang": + from trinity.common.models.sglang_model import SGLangRolloutModel + + model_cls = SGLangRolloutModel + num_gpus = config.tensor_parallel_size + else: + raise ValueError( + f"Unsupported engine type in colocate mode: {config.engine_type}" + ) + + self.logger.info( + "Creating colocated inference_model %s in %s.", + actor_name, + config.ray_namespace, + ) + handler = ( + ray.remote(model_cls) + .options( + name=actor_name, + num_cpus=0, + num_gpus=num_gpus, + 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) From 2536db1f7542952be29c4de48f8cd42e03d06903 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Sun, 12 Jul 2026 22:15:18 +0800 Subject: [PATCH 2/4] fix pre-commit --- trinity/common/config_validator.py | 2 ++ trinity/common/models/allocator.py | 21 +++++++-------------- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/trinity/common/config_validator.py b/trinity/common/config_validator.py index a5bb3a81e08..24663f9e094 100644 --- a/trinity/common/config_validator.py +++ b/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." diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 23dead8d1cc..16e1a27c28a 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -163,32 +163,25 @@ async def _create_colocate_model( actor_name = self.get_actor_name("rollout", 0, 0) config.ray_actor_name = actor_name - if config.engine_type.startswith("vllm"): - from trinity.common.models.vllm_model import vLLMRolloutModel - - model_cls = vLLMRolloutModel - num_gpus = 0 - elif config.engine_type == "sglang": - from trinity.common.models.sglang_model import SGLangRolloutModel - - model_cls = SGLangRolloutModel - num_gpus = config.tensor_parallel_size - else: + if not config.engine_type.startswith("vllm"): raise ValueError( - f"Unsupported engine type in colocate mode: {config.engine_type}" + "Colocate mode only supports vLLM, " + f"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(model_cls) + ray.remote(vLLMRolloutModel) .options( name=actor_name, num_cpus=0, - num_gpus=num_gpus, + num_gpus=0, namespace=config.ray_namespace, ) .remote(config=config) From 95dfad9ee41e9b47082056ec8acae94f780788fb Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Sun, 12 Jul 2026 22:19:11 +0800 Subject: [PATCH 3/4] fix pre-commit --- trinity/common/models/allocator.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 16e1a27c28a..4910dad0579 100644 --- a/trinity/common/models/allocator.py +++ b/trinity/common/models/allocator.py @@ -165,8 +165,7 @@ async def _create_colocate_model( if not config.engine_type.startswith("vllm"): raise ValueError( - "Colocate mode only supports vLLM, " - f"but got engine type: {config.engine_type}" + f"Colocate mode only supports vLLM, but got engine type: {config.engine_type}" ) from trinity.common.models.vllm_model import vLLMRolloutModel From 3c2afda2cdcb5daff793df369b0ac30507904605 Mon Sep 17 00:00:00 2001 From: "panxuchen.pxc" Date: Mon, 13 Jul 2026 10:49:01 +0800 Subject: [PATCH 4/4] add placement group name --- trinity/common/models/allocator.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 4910dad0579..50d03eb881d 100644 --- a/trinity/common/models/allocator.py +++ b/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 @@ -116,7 +121,11 @@ async def create_all_models(self) -> Tuple[List[ModelWrapper], List[List[ModelWr 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