(backend: str, **kwargs)
| 70 | |
| 71 | |
| 72 | def _init_dist_mpi(backend: str, **kwargs) -> None: |
| 73 | local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK']) |
| 74 | torch.cuda.set_device(local_rank) |
| 75 | if 'MASTER_PORT' not in os.environ: |
| 76 | # 29500 is torch.distributed default port |
| 77 | os.environ['MASTER_PORT'] = '29500' |
| 78 | if 'MASTER_ADDR' not in os.environ: |
| 79 | raise KeyError('The environment variable MASTER_ADDR is not set') |
| 80 | os.environ['WORLD_SIZE'] = os.environ['OMPI_COMM_WORLD_SIZE'] |
| 81 | os.environ['RANK'] = os.environ['OMPI_COMM_WORLD_RANK'] |
| 82 | dist.init_process_group(backend=backend, **kwargs) |
| 83 | |
| 84 | |
| 85 | def _init_dist_slurm(backend: str, port: Optional[int] = None) -> None: |
no test coverage detected
searching dependent graphs…