(gpu_id, global_rank, world_size, dist_url, dist_backend)
| 49 | |
| 50 | |
| 51 | def init_distributed(gpu_id, global_rank, world_size, dist_url, dist_backend): |
| 52 | torch.cuda.set_device(gpu_id) |
| 53 | print( |
| 54 | f"| distributed init (rank {global_rank}) (world {world_size}): {dist_url}", |
| 55 | flush=True, |
| 56 | ) |
| 57 | torch.distributed.init_process_group( |
| 58 | backend=dist_backend, |
| 59 | init_method=dist_url, |
| 60 | world_size=world_size, |
| 61 | rank=global_rank, |
| 62 | ) |
| 63 | torch.distributed.barrier() |
| 64 | setup_print_for_distributed(is_primary()) |
| 65 | |
| 66 | |
| 67 | def all_reduce_sum(tensor): |
nothing calls this directly
no test coverage detected