Initialize torch.distributed.
(backend='nccl')
| 39 | |
| 40 | |
| 41 | def initialize_distributed(backend='nccl'): |
| 42 | """Initialize torch.distributed.""" |
| 43 | # Get local rank in case it is provided. |
| 44 | parser = argparse.ArgumentParser() |
| 45 | parser.add_argument('--local_rank', type=int, default=None, |
| 46 | help='local rank passed from distributed launcher') |
| 47 | args = parser.parse_args() |
| 48 | local_rank = args.local_rank |
| 49 | |
| 50 | # Get rank and world size. |
| 51 | rank = int(os.getenv('RANK', '0')) |
| 52 | world_size = int(os.getenv("WORLD_SIZE", '1')) |
| 53 | |
| 54 | print('> initializing torch.distributed with local rank: {}, ' |
| 55 | 'rank: {}, world size: {}'.format(local_rank, rank, world_size)) |
| 56 | |
| 57 | # Set the device id. |
| 58 | device = rank % torch.cuda.device_count() |
| 59 | if local_rank is not None: |
| 60 | device = local_rank |
| 61 | torch.cuda.set_device(device) |
| 62 | |
| 63 | # Call the init process. |
| 64 | init_method = 'tcp://' |
| 65 | master_ip = os.getenv('MASTER_ADDR', 'localhost') |
| 66 | master_port = os.getenv('MASTER_PORT', '6000') |
| 67 | init_method += master_ip + ':' + master_port |
| 68 | torch.distributed.init_process_group( |
| 69 | backend=backend, |
| 70 | world_size=world_size, |
| 71 | rank=rank, |
| 72 | init_method=init_method) |
| 73 | |
| 74 | |
| 75 | def print_separator(message): |
no outgoing calls
no test coverage detected