(rank: int, world_size: int, master_addr: str = 'localhost', port: int = 12355, backend: str = 'nccl')
| 146 | |
| 147 | |
| 148 | def init_distributed(rank: int, world_size: int, master_addr: str = 'localhost', port: int = 12355, backend: str = 'nccl'): |
| 149 | print(rank, 'initializing distributed') |
| 150 | os.environ["MASTER_ADDR"] = master_addr |
| 151 | os.environ["MASTER_PORT"] = str(port) |
| 152 | dist.init_process_group(backend, rank=rank, world_size=world_size) |
| 153 | torch.cuda.set_device(rank) |
| 154 | |
| 155 | |
| 156 | class TemporarilySeededRandom: |