已合并
add datatypes for HCCLUtils #30224
Kuteriod创建于 1月29日
add datatypes for HCCLUtils #30224
已合并
共 2 个文件变更+57-2
| @@ -71,10 +71,22 @@ class ProcessGroupHCCLTest(TestCase): | |||
| 71 | 71 | ||
| 72 | def _test_broadcast_process( | 72 | def _test_broadcast_process( |
| 73 | cls, rank, shared_tensors, world_size, init_pg, c2p, p2c): | 73 | cls, rank, shared_tensors, world_size, init_pg, c2p, p2c): |
| 74 | + dtype_name = None | ||
| 75 | + if isinstance(shared_tensors, tuple): | ||
| 76 | + shared_tensors, dtype_name = shared_tensors | ||
| 74 | pg = init_pg(rank, world_size) | 77 | pg = init_pg(rank, world_size) |
| 75 | xs = [shared_tensors[rank].to(f"npu:{rank}")] | 78 | xs = [shared_tensors[rank].to(f"npu:{rank}")] |
| 79 | + if dtype_name is not None: | ||
| 80 | + expected = shared_tensors[0].to(torch.float32) | ||
| 81 | + result = xs[0].to(torch.float32).cpu() | ||
| 76 | pg.broadcast(xs).wait() | 82 | pg.broadcast(xs).wait() |
| 77 | - c2p.put((rank, torch.zeros(2, 2), xs[0].to("cpu"))) | 83 | + if dtype_name is not None: |
| 84 | + expected = shared_tensors[0].to(torch.float32) | ||
| 85 | + result = xs[0].to(torch.float32).cpu() | ||
| 86 | + else: | ||
| 87 | + expected = shared_tensors[0] | ||
| 88 | + result = xs[0].to("cpu") | ||
| 89 | + c2p.put((rank, expected, result)) | ||
| 78 | p2c.get() | 90 | p2c.get() |
| 79 | 91 | ||
| 80 | def test_shared_broadcast_hccl(self): | 92 | def test_shared_broadcast_hccl(self): |
| @@ -103,12 +115,24 @@ class ProcessGroupHCCLTest(TestCase): | |||
| 103 | 115 | ||
| 104 | def _test_allgather_process( | 116 | def _test_allgather_process( |
| 105 | cls, rank, shared_tensors, world_size, init_pg, c2p, p2c): | 117 | cls, rank, shared_tensors, world_size, init_pg, c2p, p2c): |
| 118 | + dtype_name = None | ||
| 119 | + if isinstance(shared_tensors, tuple): | ||
| 120 | + shared_tensors, dtype_name = shared_tensors | ||
| 106 | pg = init_pg(rank, world_size) | 121 | pg = init_pg(rank, world_size) |
| 107 | xs = [shared_tensors[rank].to(f"npu:{rank}")] | 122 | xs = [shared_tensors[rank].to(f"npu:{rank}")] |
| 123 | + if dtype_name is not None: | ||
| 124 | + dtype = getattr(torch, dtype_name, getattr(torch_npu, dtype_name, None)) | ||
| 125 | + xs[0] = xs[0].to(dtype) | ||
| 108 | ys = [[torch.zeros_like(xs[0]) for i in range(world_size)]] | 126 | ys = [[torch.zeros_like(xs[0]) for i in range(world_size)]] |
| 109 | pg.allgather(ys, xs).wait() | 127 | pg.allgather(ys, xs).wait() |
| 110 | for i in range(world_size): | 128 | for i in range(world_size): |
| 111 | - c2p.put((rank, torch.ones(2, 2) * i, ys[0][i].to("cpu"))) | 129 | + if dtype_name is not None: |
| 130 | + expected = shared_tensors[i].to(torch.float32) | ||
| 131 | + result = ys[0][i].to(torch.float32).cpu() | ||
| 132 | + else: | ||
| 133 | + expected = shared_tensors[i] | ||
| 134 | + result = ys[0][i].to("cpu") | ||
| 135 | + c2p.put((rank, expected, result)) | ||
| 112 | 136 | ||
| 113 | p2c.get() | 137 | p2c.get() |
| 114 | 138 | ||
| @@ -119,6 +143,31 @@ class ProcessGroupHCCLTest(TestCase): | |||
| 119 | ProcessGroupHCCLTest._init_pg_hccl, | 143 | ProcessGroupHCCLTest._init_pg_hccl, |
| 120 | self.world_size) | 144 | self.world_size) |
| 121 | 145 | ||
| 146 | + def test_shared_broadcast_hccl_uint_dtypes(self): | ||
| 147 | + base = torch.tensor([1, 2, 4, 8], dtype=torch.float16).reshape(2, 2) | ||
| 148 | + shared_tensors = [base * (2 ** (i * 4)) for i in range(self.world_size)] | ||
| 149 | + for dtype_name in ["uint16", "uint32", "uint64"]: | ||
| 150 | + dtype = getattr(torch, dtype_name, getattr(torch_npu, dtype_name, None)) | ||
| 151 | + self.assertIsNotNone(dtype, f"{dtype_name} not available") | ||
| 152 | + self._test_multiprocess( | ||
| 153 | + ProcessGroupHCCLTest._test_broadcast_process, | ||
| 154 | + (shared_tensors, dtype_name), | ||
| 155 | + ProcessGroupHCCLTest._init_pg_hccl, | ||
| 156 | + 1) | ||
| 157 | + | ||
| 158 | + | ||
| 159 | + def test_shared_allgather_hccl_uint_dtypes(self): | ||
| 160 | + base = torch.tensor([1, 2, 4, 8], dtype=torch.float16).reshape(2, 2) | ||
| 161 | + shared_tensors = [base * (2 ** (i * 4)) for i in range(self.world_size)] | ||
| 162 | + for dtype_name in ["uint16", "uint32", "uint64"]: | ||
| 163 | + dtype = getattr(torch, dtype_name, getattr(torch_npu, dtype_name, None)) | ||
| 164 | + self.assertIsNotNone(dtype, f"{dtype_name} not available") | ||
| 165 | + self._test_multiprocess( | ||
| 166 | + ProcessGroupHCCLTest._test_allgather_process, | ||
| 167 | + (shared_tensors, dtype_name), | ||
| 168 | + ProcessGroupHCCLTest._init_pg_hccl, | ||
| 169 | + self.world_size) | ||
| 170 | + | ||
| 122 | 171 | ||
| 123 | class ComputeBucketAssignmentTest(TestCase): | 172 | class ComputeBucketAssignmentTest(TestCase): |
| 124 | def test_single_limit_single_dtype(self): | 173 | def test_single_limit_single_dtype(self): |
| @@ -57,6 +57,9 @@ std::map<at::ScalarType, HcclDataType> kScalarTypeToHcclDataType = { | |||
| 57 | {at::kLong, HCCL_DATA_TYPE_INT64}, | 57 | {at::kLong, HCCL_DATA_TYPE_INT64}, |
| 58 | {at::kHalf, HCCL_DATA_TYPE_FP16}, | 58 | {at::kHalf, HCCL_DATA_TYPE_FP16}, |
| 59 | {at::kFloat, HCCL_DATA_TYPE_FP32}, | 59 | {at::kFloat, HCCL_DATA_TYPE_FP32}, |
| 60 | + {at::ScalarType::UInt16, HCCL_DATA_TYPE_UINT16}, | ||
| 61 | + {at::ScalarType::UInt32, HCCL_DATA_TYPE_UINT32}, | ||
| 62 | + {at::ScalarType::UInt64, HCCL_DATA_TYPE_UINT64}, | ||
| 60 | {at::kDouble, HCCL_DATA_TYPE_FP64}, | 63 | {at::kDouble, HCCL_DATA_TYPE_FP64}, |
| 61 | {at::kBool, HCCL_DATA_TYPE_UINT8}, | 64 | {at::kBool, HCCL_DATA_TYPE_UINT8}, |
| 62 | {at::kBFloat16, HCCL_DATA_TYPE_BFP16}, | 65 | {at::kBFloat16, HCCL_DATA_TYPE_BFP16}, |
| @@ -69,6 +72,9 @@ std::map<HcclDataType, std::string> kHcclDataTypeToStringMap = { | |||
| 69 | {HCCL_DATA_TYPE_INT8, "at::kChar"}, | 72 | {HCCL_DATA_TYPE_INT8, "at::kChar"}, |
| 70 | {HCCL_DATA_TYPE_INT16, "at::kShort"}, | 73 | {HCCL_DATA_TYPE_INT16, "at::kShort"}, |
| 71 | {HCCL_DATA_TYPE_INT32, "at::kInt"}, | 74 | {HCCL_DATA_TYPE_INT32, "at::kInt"}, |
| 75 | + {HCCL_DATA_TYPE_UINT16, "at::ScalarType::UInt16"}, | ||
| 76 | + {HCCL_DATA_TYPE_UINT32, "at::ScalarType::UInt32"}, | ||
| 77 | + {HCCL_DATA_TYPE_UINT64, "at::ScalarType::UInt64"}, | ||
| 72 | {HCCL_DATA_TYPE_INT64, "at::kLong"}, | 78 | {HCCL_DATA_TYPE_INT64, "at::kLong"}, |
| 73 | {HCCL_DATA_TYPE_FP16, "at::kHalf"}, | 79 | {HCCL_DATA_TYPE_FP16, "at::kHalf"}, |
| 74 | {HCCL_DATA_TYPE_FP32, "at::kFloat"}, | 80 | {HCCL_DATA_TYPE_FP32, "at::kFloat"}, |