| 85 | ) |
| 86 | |
| 87 | def init_torch_distributed(self): |
| 88 | if self.distributed: |
| 89 | # See: |
| 90 | # https://docs.nvidia.com/deeplearning/sdk/nccl-developer-guide/docs/env.html |
| 91 | os.environ.setdefault("NCCL_DEBUG", "INFO") |
| 92 | |
| 93 | # See: |
| 94 | # https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group |
| 95 | os.environ.setdefault("NCCL_BLOCKING_WAIT", "1") |
| 96 | |
| 97 | torch.distributed.init_process_group( |
| 98 | backend=self.dist_backend, |
| 99 | init_method=self.dist_init_method, |
| 100 | world_size=self.dist_world_size, |
| 101 | rank=self.dist_rank, |
| 102 | ) |
| 103 | |
| 104 | # About distributed model: |
| 105 | # if self.local_rank is not None and ngpu == 1 |
| 106 | # => Distributed with n-Process and n-GPU |
| 107 | # if self.local_rank is None and ngpu >= 1 |
| 108 | # => Distributed with 1-Process and n-GPU |
| 109 | if self.local_rank is not None and self.ngpu > 0: |
| 110 | torch.cuda.set_device(self.local_rank) |
| 111 | |
| 112 | def init_deepspeed(self): |
| 113 | try: |