(backend, **kwargs)
| 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 |
| 43 | rank = int(os.environ['RANK']) |
| 44 | num_gpus = torch.cuda.device_count() |
| 45 | torch.cuda.set_device(rank % num_gpus) |
| 46 | deepspeed.init_distributed(dist_backend=backend) |
| 47 | |
| 48 | def _init_dist_mpi(backend, **kwargs): |
| 49 | local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK']) |