(args)
| 273 | |
| 274 | |
| 275 | def init_distributed_mode1(args): |
| 276 | if args.dist_on_itp: |
| 277 | args.rank = int(os.environ['OMPI_COMM_WORLD_RANK']) |
| 278 | args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE']) |
| 279 | args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK']) |
| 280 | args.dist_url = "tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT']) |
| 281 | os.environ['LOCAL_RANK'] = str(args.gpu) |
| 282 | os.environ['RANK'] = str(args.rank) |
| 283 | os.environ['WORLD_SIZE'] = str(args.world_size) |
| 284 | # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"] |
| 285 | elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: |
| 286 | args.rank = int(os.environ["RANK"]) |
| 287 | args.world_size = int(os.environ['WORLD_SIZE']) |
| 288 | args.gpu = int(os.environ['LOCAL_RANK']) |
| 289 | elif 'SLURM_PROCID' in os.environ: |
| 290 | args.rank = int(os.environ['SLURM_PROCID']) |
| 291 | args.gpu = args.rank % torch.cuda.device_count() |
| 292 | else: |
| 293 | print('Not using distributed mode') |
| 294 | setup_for_distributed(is_master=True) # hack |
| 295 | args.distributed = False |
| 296 | return |
| 297 | |
| 298 | args.distributed = True |
| 299 | |
| 300 | torch.cuda.set_device(args.gpu) |
| 301 | args.dist_backend = 'nccl' |
| 302 | print('| distributed init (rank {}): {}, gpu {}'.format( |
| 303 | args.rank, args.dist_url, args.gpu), flush=True) |
| 304 | torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, |
| 305 | world_size=args.world_size, rank=args.rank) |
| 306 | torch.distributed.barrier() |
| 307 | setup_for_distributed(args.rank == 0) |
| 308 | |
| 309 | |
| 310 | class NativeScalerWithGradNormCount: |
nothing calls this directly
no test coverage detected