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..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." @@ -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 diff --git a/trinity/common/models/allocator.py b/trinity/common/models/allocator.py index 551ab09641d..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 @@ -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)