(args=SimpleNamespace())
| 92 | |
| 93 | |
| 94 | def init_distributed_mode(args=SimpleNamespace()): |
| 95 | random_seed() |
| 96 | if getattr(args, 'dist_on_itp', False): |
| 97 | args.rank = int(os.environ['OMPI_COMM_WORLD_RANK']) |
| 98 | args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE']) |
| 99 | args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK']) |
| 100 | args.local_rank = args.gpu |
| 101 | args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT']) |
| 102 | os.environ['LOCAL_RANK'] = str(args.gpu) |
| 103 | os.environ['RANK'] = str(args.rank) |
| 104 | os.environ['WORLD_SIZE'] = str(args.world_size) |
| 105 | # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"] |
| 106 | elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ and "LOCAL_RANK" in os.environ: |
| 107 | args.world_size = int(os.environ['WORLD_SIZE']) |
| 108 | args.rank = int(os.environ["RANK"]) |
| 109 | args.gpu = int(os.environ['LOCAL_RANK']) |
| 110 | args.local_rank = args.gpu |
| 111 | args.dist_url = 'env://' |
| 112 | elif 'SLURM_PROCID' in os.environ: |
| 113 | os.environ['MASTER_PORT'] = '8964' |
| 114 | while 'MASTER_ADDR' not in os.environ or len(os.environ['MASTER_ADDR'].strip()) == 0: |
| 115 | os.environ['MASTER_ADDR'] = subprocess.check_output('sinfo -Nh -n %s | head -n 1 | awk \'{print $1}\'' % os.environ['SLURM_NODELIST'], shell=True, ).decode().strip() |
| 116 | time.sleep(1) |
| 117 | print(os.environ['MASTER_ADDR']) |
| 118 | args.world_size = int(os.environ['SLURM_NPROCS']) |
| 119 | args.rank = int(os.environ['SLURM_PROCID']) |
| 120 | args.gpu = args.rank % torch.cuda.device_count() |
| 121 | args.local_rank = args.gpu |
| 122 | args.dist_url = 'env://' |
| 123 | os.environ['LOCAL_RANK'] = str(args.gpu) |
| 124 | os.environ['WORLD_SIZE'] = str(args.world_size) |
| 125 | os.environ['RANK'] = str(args.rank) |
| 126 | else: |
| 127 | os.environ['MASTER_ADDR'] = "127.0.0.1" |
| 128 | os.environ['MASTER_PORT'] = str(find_free_port(9000, 10000)) |
| 129 | os.environ['RANK'] = '0' |
| 130 | os.environ['LOCAL_RANK'] = '0' |
| 131 | os.environ['WORLD_SIZE'] = '1' |
| 132 | args.rank = 0 |
| 133 | args.gpu = args.local_rank = 0 |
| 134 | args.world_size = 1 |
| 135 | args.dist_url = 'env://' |
| 136 | |
| 137 | args.distributed = True |
| 138 | |
| 139 | torch.cuda.set_device(args.gpu) |
| 140 | args.dist_backend = 'nccl' |
| 141 | print('| distributed init (rank {}): {}, gpu {}'.format( |
| 142 | args.rank, args.dist_url, args.gpu), flush=True) |
| 143 | torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, |
| 144 | world_size=args.world_size, rank=args.rank) |
| 145 | torch.distributed.barrier() |
| 146 | setup_for_distributed(args.rank == 0) |
| 147 | |
| 148 | |
| 149 | class SmoothedValue(object): |
nothing calls this directly
no test coverage detected