| 22 | |
| 23 | |
| 24 | class NCCLCommunicator: |
| 25 | def __init__(self, |
| 26 | comm_rank: int, |
| 27 | cuda_id: int, |
| 28 | comm_group_size: int, |
| 29 | comm_name: str): |
| 30 | self.comm_rank = comm_rank |
| 31 | cupy.cuda.Device(cuda_id).use() |
| 32 | self.comm_group_size = comm_group_size |
| 33 | print("Initialize NCCLCommunicator: <", comm_name, ">; rank:", comm_rank) |
| 34 | self.dist_store = dist.distributed_c10d._get_default_store() |
| 35 | |
| 36 | if self.comm_rank == 0: |
| 37 | cuda_id = cupy.cuda.nccl.get_unique_id() |
| 38 | # print(cuda_id) |
| 39 | cuda_id_str = np.array(cuda_id).tobytes() |
| 40 | self.dist_store.set('group-'+comm_name+'-unique-id', cuda_id_str) |
| 41 | # print("Master put <group-"+comm_name+"-unique-id: ", cuda_id_str, ">.") |
| 42 | else: |
| 43 | cuda_id_str = self.dist_store.get('group-'+comm_name+'-unique-id') |
| 44 | |
| 45 | comm_id = tuple(np.frombuffer(cuda_id_str, dtype=int)) |
| 46 | # comm_id = cupy.cuda.nccl.get_unique_id() |
| 47 | # print(comm_id) |
| 48 | self.comm = cupy.cuda.nccl.NcclCommunicator(comm_group_size, comm_id, comm_rank) |
| 49 | |
| 50 | @staticmethod |
| 51 | def barrier(): |
| 52 | dist.barrier() |
| 53 | |
| 54 | def store_set(self, key, value): |
| 55 | self.dist_store.set(key, value) |
| 56 | |
| 57 | def store_get(self, key): |
| 58 | return self.dist_store.get(key) |
| 59 | |
| 60 | def send(self, |
| 61 | tensor: torch.Tensor, |
| 62 | dst: int, |
| 63 | stream=cupy.cuda.Stream.null): |
| 64 | # print("Send tensor of size:", torch.numel(tensor)) |
| 65 | self.comm.send( |
| 66 | tensor.data_ptr(), |
| 67 | torch.numel(tensor), |
| 68 | _type_torch_to_cupy(tensor.dtype), |
| 69 | dst, |
| 70 | stream.ptr |
| 71 | ) |
| 72 | |
| 73 | def recv(self, |
| 74 | tensor: torch.Tensor, |
| 75 | src: int, |
| 76 | stream=cupy.cuda.Stream.null): |
| 77 | # print("Recv tensor of size:", torch.numel(tensor)) |
| 78 | # print("mean:", torch.mean(tensor).item(), " std:", torch.std(tensor).item()) |
| 79 | self.comm.recv( |
| 80 | tensor.data_ptr(), |
| 81 | torch.numel(tensor), |
no outgoing calls
no test coverage detected