()
| 65 | |
| 66 | |
| 67 | def _init_distributed() -> tuple[int, int, int, torch.device, str]: |
| 68 | if not dist.is_available(): |
| 69 | raise RuntimeError("torch.distributed is not available in this PyTorch build.") |
| 70 | |
| 71 | rank, local_rank, world_size = _rank_env() |
| 72 | if world_size < 2: |
| 73 | raise RuntimeError("Run this script with torchrun and at least 2 processes.") |
| 74 | |
| 75 | if torch.cuda.is_available(): |
| 76 | if local_rank >= torch.cuda.device_count(): |
| 77 | raise RuntimeError( |
| 78 | f"LOCAL_RANK={local_rank} but only {torch.cuda.device_count()} CUDA devices exist." |
| 79 | ) |
| 80 | torch.cuda.set_device(local_rank) |
| 81 | device = torch.device("cuda", local_rank) |
| 82 | backend = "nccl" |
| 83 | else: |
| 84 | device = torch.device("cpu") |
| 85 | backend = "gloo" |
| 86 | |
| 87 | if not dist.is_initialized(): |
| 88 | dist.init_process_group(backend=backend) |
| 89 | return rank, local_rank, world_size, device, backend |
| 90 | |
| 91 | |
| 92 | def _print_rank0(rank: int, message: str) -> None: |
no test coverage detected