(args)
| 297 | |
| 298 | |
| 299 | def init_distributed_mode(args): |
| 300 | if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: |
| 301 | args.rank = int(os.environ["RANK"]) |
| 302 | args.world_size = int(os.environ['WORLD_SIZE']) |
| 303 | args.gpu = int(os.environ['LOCAL_RANK']) |
| 304 | elif 'SLURM_PROCID' in os.environ: |
| 305 | args.rank = int(os.environ['SLURM_PROCID']) |
| 306 | args.gpu = args.rank % torch.cuda.device_count() |
| 307 | elif hasattr(args, "rank"): |
| 308 | pass |
| 309 | else: |
| 310 | print('Not using distributed mode') |
| 311 | args.distributed = False |
| 312 | return |
| 313 | |
| 314 | args.distributed = True |
| 315 | |
| 316 | torch.cuda.set_device(args.gpu) |
| 317 | args.dist_backend = 'nccl' |
| 318 | print('| distributed init (rank {}): {}'.format( |
| 319 | args.rank, args.dist_url), flush=True) |
| 320 | torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, |
| 321 | world_size=args.world_size, rank=args.rank) |
| 322 | setup_for_distributed(args.rank == 0) |
no test coverage detected