| from datetime import timedelta |
| import os |
|
|
| import torch |
| import torch.distributed as dist |
|
|
|
|
| def init_distributed_mode(): |
| if "MASTER_ADDR" in os.environ and "MASTER_PORT" in os.environ: |
| dist_url = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}" |
| else: |
| dist_url = "env://" |
|
|
| if "RANK" in os.environ and "WORLD_SIZE" in os.environ: |
| rank = int(os.environ["RANK"]) |
| world_size = int(os.environ["WORLD_SIZE"]) |
| local_rank = int(os.environ.get("LOCAL_RANK", 0)) |
| elif "SLURM_NODEID" in os.environ: |
| gpus_per_node = torch.cuda.device_count() |
| node_id = int(os.environ["SLURM_NODEID"]) |
| local_rank = int(os.environ["SLURM_LOCALID"]) |
| rank = node_id * gpus_per_node + local_rank |
| world_size = int(os.environ["SLURM_NTASKS"]) |
| else: |
| raise RuntimeError("Distributed environment not properly set.") |
|
|
| torch.cuda.set_device(local_rank) |
|
|
| dist.init_process_group( |
| backend="nccl", |
| init_method=dist_url, |
| world_size=world_size, |
| rank=rank, |
| timeout=timedelta(seconds=4800), |
| ) |
| dist.barrier() |
|
|
|
|
| def is_dist_avail_and_initialized(): |
| return dist.is_available() and dist.is_initialized() |
|
|
|
|
| def get_world_size(): |
| return dist.get_world_size() if is_dist_avail_and_initialized() else 1 |
|
|
|
|
| def get_rank(): |
| return dist.get_rank() if is_dist_avail_and_initialized() else 0 |
|
|
|
|
| def is_main_process(): |
| return get_rank() == 0 |
|
|
