| import os
|
| import subprocess
|
|
|
| import torch
|
| import torch.distributed as dist
|
|
|
|
|
| def setup_distributed(backend="nccl", port=None):
|
| """AdaHessian Optimizer
|
| Lifted from https://github.com/BIGBALLON/distribuuuu/blob/master/distribuuuu/utils.py
|
| Originally licensed MIT, Copyright (c) 2020 Wei Li
|
| """
|
| num_gpus = torch.cuda.device_count()
|
|
|
| if "SLURM_JOB_ID" in os.environ:
|
| rank = int(os.environ["SLURM_PROCID"])
|
| world_size = int(os.environ["SLURM_NTASKS"])
|
| node_list = os.environ["SLURM_NODELIST"]
|
| addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1")
|
|
|
| if port is not None:
|
| os.environ["MASTER_PORT"] = str(port)
|
| elif "MASTER_PORT" not in os.environ:
|
| os.environ["MASTER_PORT"] = "10685"
|
| if "MASTER_ADDR" not in os.environ:
|
| os.environ["MASTER_ADDR"] = addr
|
| os.environ["WORLD_SIZE"] = str(world_size)
|
| os.environ["LOCAL_RANK"] = str(rank % num_gpus)
|
| os.environ["RANK"] = str(rank)
|
| else:
|
| rank = int(os.environ["RANK"])
|
| world_size = int(os.environ["WORLD_SIZE"])
|
|
|
| torch.cuda.set_device(rank % num_gpus)
|
|
|
| dist.init_process_group(
|
| backend=backend,
|
| world_size=world_size,
|
| rank=rank,
|
| )
|
| return rank, world_size
|
|
|
