(
self,
master_address: str,
master_port: int,
world_size: int = None,
group_name: str = None,
timeout: int = None,
)
| 110 | self.logger.info("Rollout models are ready. Continue weight sync initialization.") |
| 111 | |
| 112 | async def setup_weight_sync_group( |
| 113 | self, |
| 114 | master_address: str, |
| 115 | master_port: int, |
| 116 | world_size: int = None, |
| 117 | group_name: str = None, |
| 118 | timeout: int = None, |
| 119 | ): |
| 120 | await self._wait_for_models_ready() |
| 121 | base_offset = 1 if self.use_nccl_sync else 0 |
| 122 | gpu_per_engine: int = self.config.explorer.rollout_model.gpu_per_engine |
| 123 | world_size = world_size or len(self.models) * gpu_per_engine + base_offset |
| 124 | timeout = timeout or self.config.synchronizer.sync_timeout |
| 125 | group_name = group_name or self.config.synchronizer.group_name |
| 126 | self.logger.info( |
| 127 | f"Initialize process group for weight synchronization, " |
| 128 | f"master_address={master_address}, master_port={master_port}, " |
| 129 | f"world_size={world_size}, rank_offset={base_offset}" |
| 130 | ) |
| 131 | |
| 132 | refs = [ |
| 133 | model.init_process_group( |
| 134 | master_address=master_address, |
| 135 | master_port=master_port, |
| 136 | rank_offset=i * gpu_per_engine + base_offset, |
| 137 | world_size=world_size, |
| 138 | group_name=group_name, |
| 139 | timeout=timeout, |
| 140 | ) |
| 141 | for i, model in enumerate(self.models) |
| 142 | ] |
| 143 | await asyncio.gather(*refs) |
| 144 | |
| 145 | async def set_state_dict_meta(self, state_dict_meta: List): |
| 146 | """Set the state_dict meta on all model workers for NCCL weight sync. |
no test coverage detected