(args)
| 215 | |
| 216 | |
| 217 | def init_distributed_mode(args): |
| 218 | if args.dist_on_itp: |
| 219 | args.rank = int(os.environ['OMPI_COMM_WORLD_RANK']) |
| 220 | args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE']) |
| 221 | args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK']) |
| 222 | args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT']) |
| 223 | os.environ['LOCAL_RANK'] = str(args.gpu) |
| 224 | os.environ['RANK'] = str(args.rank) |
| 225 | os.environ['WORLD_SIZE'] = str(args.world_size) |
| 226 | # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"] |
| 227 | elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: |
| 228 | args.rank = int(os.environ["RANK"]) |
| 229 | args.world_size = int(os.environ['WORLD_SIZE']) |
| 230 | args.gpu = int(os.environ['LOCAL_RANK']) |
| 231 | elif 'SLURM_PROCID' in os.environ: |
| 232 | args.rank = int(os.environ['SLURM_PROCID']) |
| 233 | args.gpu = args.rank % torch.cuda.device_count() |
| 234 | else: |
| 235 | print('Not using distributed mode') |
| 236 | setup_for_distributed(is_master=True) # hack |
| 237 | args.distributed = False |
| 238 | return |
| 239 | |
| 240 | args.distributed = True |
| 241 | |
| 242 | torch.cuda.set_device(args.gpu) |
| 243 | args.dist_backend = 'nccl' |
| 244 | print('| distributed init (rank {}): {}, gpu {}'.format( |
| 245 | args.rank, args.dist_url, args.gpu), flush=True) |
| 246 | torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, |
| 247 | world_size=args.world_size, rank=args.rank) |
| 248 | torch.distributed.barrier() |
| 249 | setup_for_distributed(args.rank == 0) |
| 250 | |
| 251 | |
| 252 | class NativeScalerWithGradNormCount: |
nothing calls this directly
no test coverage detected