| 270 | |
| 271 | |
| 272 | def _distributed_worker( |
| 273 | world_rank: int, |
| 274 | world_size: int, |
| 275 | fn: Callable, |
| 276 | args: Any, |
| 277 | local_rank: Optional[int] = None, |
| 278 | verbose: bool = False, |
| 279 | ) -> bool: |
| 280 | if local_rank is None: # single Node |
| 281 | local_rank = world_rank |
| 282 | if verbose: |
| 283 | print("Distributed worker: %d / %d" % (world_rank + 1, world_size)) |
| 284 | distributed = world_size > 1 |
| 285 | if distributed: |
| 286 | torch.cuda.set_device(local_rank) |
| 287 | torch.distributed.init_process_group( |
| 288 | backend="nccl", world_size=world_size, rank=world_rank |
| 289 | ) |
| 290 | # Dump collection that participates all ranks. |
| 291 | # This initializes the communicator required by `batch_isend_irecv`. |
| 292 | # See: https://github.com/pytorch/pytorch/pull/74701 |
| 293 | _ = [None for _ in range(world_size)] |
| 294 | torch.distributed.all_gather_object(_, 0) |
| 295 | fn(local_rank, world_rank, world_size, args) |
| 296 | if distributed: |
| 297 | torch.distributed.barrier() |
| 298 | torch.distributed.destroy_process_group() |
| 299 | if verbose: |
| 300 | print("Job Done for worker: %d / %d" % (world_rank + 1, world_size)) |
| 301 | return True |
| 302 | |
| 303 | |
| 304 | def cli(fn: Callable, args: Any, verbose: bool = False) -> bool: |