(fn: Callable, args: tuple)
| 63 | |
| 64 | |
| 65 | def _torchrun_worker(fn: Callable, args: tuple) -> None: |
| 66 | rank = int(os.environ["RANK"]) |
| 67 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 68 | |
| 69 | if rank == 0: |
| 70 | _print_gpu_topology() |
| 71 | torch.cuda.set_device(local_rank) |
| 72 | backend = "nccl" |
| 73 | if rank == 0: |
| 74 | print(f"Using distributed backend: '{backend}' (torchrun)") |
| 75 | |
| 76 | dist.init_process_group(backend=backend, init_method="env://", device_id=local_rank) |
| 77 | try: |
| 78 | fn(*args) |
| 79 | finally: |
| 80 | dist.destroy_process_group() |
| 81 | |
| 82 | |
| 83 | def _distributed_worker(rank: int, world_size: int, port: int, fn: Callable, args: tuple) -> None: |
no test coverage detected