(launcher, backend='nccl', **kwargs)
| 27 | |
| 28 | |
| 29 | def init_dist(launcher, backend='nccl', **kwargs): |
| 30 | if mp.get_start_method(allow_none=True) is None: |
| 31 | mp.set_start_method('spawn') |
| 32 | if launcher == 'pytorch': |
| 33 | _init_dist_pytorch(backend, **kwargs) |
| 34 | elif launcher == 'mpi': |
| 35 | _init_dist_mpi(backend, **kwargs) |
| 36 | elif launcher == 'slurm': |
| 37 | _init_dist_slurm(backend, **kwargs) |
| 38 | else: |
| 39 | raise ValueError(f'Invalid launcher type: {launcher}') |
| 40 | |
| 41 | def _init_dist_pytorch(backend, **kwargs): |
| 42 | # TODO: use local_rank instead of rank % num_gpus |
no test coverage detected