Initialize the parameter server. env RANK, WORLD_SIZE and MASTER_ADDR must be set. Args: auto_pg: Whether to automatically initialize the process group. Notice that if auto_pg is True, will destroy the process group after update. It is recommended to set
(
self,
*,
rank: int | None = None,
world_size: int | None = None,
auto_pg: bool = True,
gpu_count: int | None = None,
mem_fraction: float | None = None,
master_addr: str | None = None,
master_port: int | None = None,
)
| 170 | shared_memory_pool_name = "__shared_memory_pool__" |
| 171 | |
| 172 | def __init__( |
| 173 | self, |
| 174 | *, |
| 175 | rank: int | None = None, |
| 176 | world_size: int | None = None, |
| 177 | auto_pg: bool = True, |
| 178 | gpu_count: int | None = None, |
| 179 | mem_fraction: float | None = None, |
| 180 | master_addr: str | None = None, |
| 181 | master_port: int | None = None, |
| 182 | ): |
| 183 | """ |
| 184 | Initialize the parameter server. env RANK, WORLD_SIZE and MASTER_ADDR must be set. |
| 185 | |
| 186 | Args: |
| 187 | auto_pg: Whether to automatically initialize the process group. |
| 188 | Notice that if auto_pg is True, will destroy the process group after update. It is recommended to set auto_pg to True! |
| 189 | mem_fraction: The proportion (as a fraction) of the current free device memory for allocation. |
| 190 | """ |
| 191 | self._rank = rank or int(os.environ["RANK"]) |
| 192 | self._world_size = world_size or int(os.environ["WORLD_SIZE"]) |
| 193 | self.device_manager = DeviceManager() |
| 194 | self._gpu_count = gpu_count or self.device_manager.device_module.device_count() |
| 195 | self._local_rank = self._rank % self._gpu_count |
| 196 | self._auto_pg = auto_pg |
| 197 | self._all_hosts = [] |
| 198 | self._global_device_uuids: list[str] = [] |
| 199 | self._local_rdma_devices: dict[str, set[int]] = defaultdict(set) |
| 200 | self._remote_rdma_devices: dict[str, set[int]] = defaultdict(set) |
| 201 | self._mem_fraction = mem_fraction or float(os.getenv("PS_MEM_FRACTION", "0.9")) |
| 202 | |
| 203 | assert self._rank is not None and self._rank >= 0, self._rank |
| 204 | assert self._world_size and self._world_size > 0, self._world_size |
| 205 | assert ( |
| 206 | self._gpu_count is not None |
| 207 | and self._gpu_count > 0 |
| 208 | and self._gpu_count <= self.device_manager.device_module.device_count() |
| 209 | ), self._gpu_count |
| 210 | assert ( |
| 211 | self._mem_fraction is not None and self._mem_fraction > 0 and self._mem_fraction <= 1 |
| 212 | ), self._mem_fraction |
| 213 | |
| 214 | self._zmq_ctx = zmq.Context() |
| 215 | self._zmq_addr_counter = 0 |
| 216 | |
| 217 | # stores the name of the checkpoint currently using the shared memory pool, or empty string if none |
| 218 | self._current_shared_memory_pool_user: str = "" |
| 219 | self._memory_pool: dict[str, list[MemoryBuffer]] = {} |
| 220 | self._memory_pool[self.shared_memory_pool_name] = [] |
| 221 | # dict key is owner_rank, value is a bucket metas list in owner_rank |
| 222 | self._current_global_parameter_metas: dict[int, MemoryBufferMetaList] = {} |
| 223 | # NPU transfer engine initialization requires prior set_device. |
| 224 | device_index = self._local_rank |
| 225 | self.device_manager.device_module.set_device(device_index) |
| 226 | try: |
| 227 | self._p2p_store = P2PStore(self.device_manager) |
| 228 | except ImportError as e: |
| 229 | logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}") |
nothing calls this directly
no test coverage detected