Initialize slurm distributed training environment. If argument ``port`` is not specified, then the master port will be system environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system environment variable, then a default port ``29500`` will be used. Args: back
()
| 95 | return rt |
| 96 | |
| 97 | def init_dist_slurm(): |
| 98 | """Initialize slurm distributed training environment. |
| 99 | If argument ``port`` is not specified, then the master port will be system |
| 100 | environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system |
| 101 | environment variable, then a default port ``29500`` will be used. |
| 102 | Args: |
| 103 | backend (str): Backend of torch.distributed. |
| 104 | port (int, optional): Master port. Defaults to None. |
| 105 | """ |
| 106 | proc_id = int(os.environ['SLURM_PROCID']) |
| 107 | ntasks = int(os.environ['SLURM_NTASKS']) |
| 108 | node_list = os.environ['SLURM_NODELIST'] |
| 109 | num_gpus = torch.cuda.device_count() |
| 110 | torch.cuda.set_device(proc_id % num_gpus) |
| 111 | addr = subprocess.getoutput( |
| 112 | f'scontrol show hostname {node_list} | head -n1') |
| 113 | # specify master port |
| 114 | if 'MASTER_PORT' in os.environ: |
| 115 | pass # use MASTER_PORT in the environment variable |
| 116 | else: |
| 117 | # 29500 is torch.distributed default port |
| 118 | os.environ['MASTER_PORT'] = '29500' |
| 119 | # use MASTER_ADDR in the environment variable if it already exists |
| 120 | if 'MASTER_ADDR' not in os.environ: |
| 121 | os.environ['MASTER_ADDR'] = addr |
| 122 | |
| 123 | os.environ['WORLD_SIZE'] = str(ntasks) |
| 124 | os.environ['LOCAL_RANK'] = str(proc_id % num_gpus) |
| 125 | os.environ['RANK'] = str(proc_id) |
| 126 | |
| 127 | dist.init_process_group(backend='nccl') |