Initialize the underlying communication collective library. Parameters ---------- ccl : str The name of the communication collective library. Currently supported libraries are: - nccl - rccl - mpi *device_ids : int
(self, ccl: str, *device_ids)
| 323 | return DModule(func(path, device), self) |
| 324 | |
| 325 | def init_ccl(self, ccl: str, *device_ids): |
| 326 | """Initialize the underlying communication collective library. |
| 327 | |
| 328 | Parameters |
| 329 | ---------- |
| 330 | ccl : str |
| 331 | The name of the communication collective library. Currently supported libraries are: |
| 332 | - nccl |
| 333 | - rccl |
| 334 | - mpi |
| 335 | |
| 336 | *device_ids : int |
| 337 | The device IDs to be used by the underlying communication library. |
| 338 | """ |
| 339 | assert ccl in ("nccl", "rccl"), f"Unsupported CCL backend: {ccl}" |
| 340 | _ffi_api.SessionInitCCL(self, ccl, Shape(device_ids)) # type: ignore # pylint: disable=no-member |
| 341 | self._clear_ipc_memory_pool() |
| 342 | |
| 343 | def broadcast( |
| 344 | self, |