(backend='nccl')
| 40 | |
| 41 | |
| 42 | def initialize(backend='nccl'): |
| 43 | port = "12345" |
| 44 | proc_id = int(os.environ['SLURM_PROCID']) |
| 45 | ntasks = int(os.environ['SLURM_NTASKS']) |
| 46 | node_list = os.environ['SLURM_NODELIST'] |
| 47 | if '[' in node_list: |
| 48 | beg = node_list.find('[') |
| 49 | pos1 = node_list.find('-', beg) |
| 50 | if pos1 < 0: |
| 51 | pos1 = 1000 |
| 52 | pos2 = node_list.find(',', beg) |
| 53 | if pos2 < 0: |
| 54 | pos2 = 1000 |
| 55 | node_list = node_list[:min(pos1, pos2)].replace('[', '') |
| 56 | addr = node_list[8:].replace('-', '.') |
| 57 | os.environ['MASTER_PORT'] = port |
| 58 | os.environ['MASTER_ADDR'] = addr |
| 59 | os.environ['WORLD_SIZE'] = str(ntasks) |
| 60 | os.environ['RANK'] = str(proc_id) |
| 61 | if backend == 'nccl': |
| 62 | dist.init_process_group(backend='nccl') |
| 63 | else: |
| 64 | dist.init_process_group(backend='gloo', rank=proc_id, world_size=ntasks) |
| 65 | rank = dist.get_rank() |
| 66 | device = rank % torch.cuda.device_count() |
| 67 | torch.cuda.set_device(device) |
| 68 | |
| 69 | |
| 70 | def finalize(): |
nothing calls this directly
no outgoing calls
no test coverage detected