AdaHessian Optimizer Lifted from https://github.com/BIGBALLON/distribuuuu/blob/master/distribuuuu/utils.py Originally licensed MIT, Copyright (c) 2020 Wei Li
(backend="nccl", port=None)
| 257 | |
| 258 | |
| 259 | def setup_distributed(backend="nccl", port=None): |
| 260 | """AdaHessian Optimizer |
| 261 | Lifted from https://github.com/BIGBALLON/distribuuuu/blob/master/distribuuuu/utils.py |
| 262 | Originally licensed MIT, Copyright (c) 2020 Wei Li |
| 263 | """ |
| 264 | num_gpus = torch.cuda.device_count() |
| 265 | # export ZHENSALLOC="hello boy!" |
| 266 | if "SLURM_JOB_ID" in os.environ and "ZHENSALLOC" not in os.environ: |
| 267 | _, rank, world_size = world_info_from_env() |
| 268 | node_list = os.environ["SLURM_NODELIST"] |
| 269 | addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1") |
| 270 | # specify master port |
| 271 | if port is not None: |
| 272 | os.environ["MASTER_PORT"] = str(port) |
| 273 | elif "MASTER_PORT" not in os.environ: |
| 274 | os.environ["MASTER_PORT"] = "10685" |
| 275 | if "MASTER_ADDR" not in os.environ: |
| 276 | os.environ["MASTER_ADDR"] = addr |
| 277 | os.environ["WORLD_SIZE"] = str(world_size) |
| 278 | os.environ["LOCAL_RANK"] = str(rank % num_gpus) |
| 279 | os.environ["RANK"] = str(rank) |
| 280 | else: |
| 281 | rank = int(os.environ["RANK"]) |
| 282 | world_size = int(os.environ["WORLD_SIZE"]) |
| 283 | |
| 284 | |
| 285 | torch.cuda.set_device(rank % num_gpus) |
| 286 | |
| 287 | dist.init_process_group( |
| 288 | backend=backend, |
| 289 | world_size=world_size, |
| 290 | rank=rank, |
| 291 | ) |
| 292 | |
| 293 | return rank, world_size |
| 294 | |
| 295 | |
| 296 |
no test coverage detected