已合并
Add npu_all_gather_base_mm testcase #3339
wang-guangbin创建于 2025年10月13日
Add npu_all_gather_base_mm testcase #3339
已合并
共 1 个文件变更+53-13
| @@ -11,7 +11,6 @@ from torch_npu.testing.common_utils import create_common_tensor, SupportedDevice | |||
| 11 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | 11 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU |
| 12 | 12 | ||
| 13 | 13 | ||
| 14 | - | ||
| 15 | class TestAllGatherBaseMm(TestCase): | 14 | class TestAllGatherBaseMm(TestCase): |
| 16 | 15 | ||
| 17 | 16 | ||
| @@ -25,7 +24,7 @@ class TestAllGatherBaseMm(TestCase): | |||
| 25 | 24 | ||
| 26 | 25 | ||
| 27 | def _test_npu_all_gather_base_mm(cls, rank, input_list): | 26 | def _test_npu_all_gather_base_mm(cls, rank, input_list): |
| 28 | - x1_list, x2_list, world_size, init_pg, c2p = input_list | 27 | + x1_list, x2_list, x1_scale_list, x2_scale_list, world_size, comm_mode, output_dtype, init_pg, c2p = input_list |
| 29 | x1 = x1_list[rank] | 28 | x1 = x1_list[rank] |
| 30 | x2 = x2_list[rank] | 29 | x2 = x2_list[rank] |
| 31 | pg = init_pg(rank, world_size) | 30 | pg = init_pg(rank, world_size) |
| @@ -37,20 +36,25 @@ class TestAllGatherBaseMm(TestCase): | |||
| 37 | 36 | ||
| 38 | x1 = x1.npu() | 37 | x1 = x1.npu() |
| 39 | x2 = x2.npu() | 38 | x2 = x2.npu() |
| 39 | + x1_scale = x1_scale_list[rank].npu() if x1_scale_list else None | ||
| 40 | + x2_scale = x2_scale_list[rank].npu() if x2_scale_list else None | ||
| 40 | out, gather_out = torch_npu.npu_all_gather_base_mm(x1, | 41 | out, gather_out = torch_npu.npu_all_gather_base_mm(x1, |
| 41 | x2, | 42 | x2, |
| 42 | hcom_name, | 43 | hcom_name, |
| 43 | world_size, | 44 | world_size, |
| 44 | bias=None, | 45 | bias=None, |
| 46 | + x1_scale=x1_scale, | ||
| 47 | + x2_scale=x2_scale, | ||
| 45 | gather_index=0, | 48 | gather_index=0, |
| 46 | gather_output=True, | 49 | gather_output=True, |
| 47 | - comm_turn=0) | 50 | + output_dtype=output_dtype, |
| 48 | - | 51 | + comm_turn=0, |
| 49 | - c2p.put((rank, out.cpu(), gather_out.cpu())) | 52 | + comm_mode=comm_mode) |
| 53 | + c2p.put((rank, out.cpu().numpy(), gather_out.cpu().numpy())) | ||
| 50 | pg.barrier() | 54 | pg.barrier() |
| 51 | 55 | ||
| 52 | def _test_multiprocess(self, f, init_pg, input_list): | 56 | def _test_multiprocess(self, f, init_pg, input_list): |
| 53 | - expt_out_list, expt_gather, x1, x2, world_size = input_list | 57 | + expt_out_list, expt_gather, x1, x2, x1_scale_list, x2_scale_list, world_size, comm_mode, output_dtype = input_list |
| 54 | ctx = mp.get_context('spawn') | 58 | ctx = mp.get_context('spawn') |
| 55 | c2p = ctx.Queue(world_size) | 59 | c2p = ctx.Queue(world_size) |
| 56 | ps = [] | 60 | ps = [] |
| @@ -58,27 +62,39 @@ class TestAllGatherBaseMm(TestCase): | |||
| 58 | for i in range(world_size): | 62 | for i in range(world_size): |
| 59 | p = ctx.Process( | 63 | p = ctx.Process( |
| 60 | target=f, | 64 | target=f, |
| 61 | - args=(i, [x1, x2, world_size, init_pg, c2p])) | 65 | + args=(i, [x1, x2, x1_scale_list, x2_scale_list, world_size, comm_mode, output_dtype, init_pg, c2p])) |
| 62 | p.start() | 66 | p.start() |
| 63 | ps.append(p) | 67 | ps.append(p) |
| 64 | 68 | ||
| 65 | for _ in range(world_size): | 69 | for _ in range(world_size): |
| 66 | rank, output, gather_output = c2p.get() | 70 | rank, output, gather_output = c2p.get() |
| 71 | + output, gather_output = torch.from_numpy(output), torch.from_numpy(gather_output) | ||
| 67 | self.assertEqual(output, expt_out_list[rank], | 72 | self.assertEqual(output, expt_out_list[rank], |
| 68 | ("rank {} Expect receive tensor {} but got {}.").format(rank, expt_out_list[rank], output)) | 73 | ("rank {} Expect receive tensor {} but got {}.").format(rank, expt_out_list[rank], output)) |
| 69 | self.assertEqual(gather_output, expt_gather, | 74 | self.assertEqual(gather_output, expt_gather, |
| 70 | ("rank {} Expect receive tensor {} but got {}.").format(rank, expt_gather, gather_output)) | 75 | ("rank {} Expect receive tensor {} but got {}.").format(rank, expt_gather, gather_output)) |
| 71 | - | ||
| 72 | for p in ps: | 76 | for p in ps: |
| 73 | p.join() | 77 | p.join() |
| 74 | 78 | ||
| 75 | - def _construct_excepted_result(self, x1_list, x2_list, world_size): | 79 | + def _construct_excepted_result(self, x1_list, x2_list, world_size, x1_scale_list=None, x2_scale_list=None, output_dtype=None): |
| 76 | gather_out = torch.cat(x1_list) | 80 | gather_out = torch.cat(x1_list) |
| 81 | + if x1_scale_list: | ||
| 82 | + x1_scale = torch.cat(x1_scale_list) | ||
| 77 | out_list = [] | 83 | out_list = [] |
| 78 | - out_dtype = gather_out.dtype | 84 | + if output_dtype: |
| 85 | + out_dtype = output_dtype | ||
| 86 | + else: | ||
| 87 | + out_dtype = gather_out.dtype | ||
| 79 | for i in range(world_size): | 88 | for i in range(world_size): |
| 80 | - out_list.append(torch.matmul(gather_out.npu(), x2_list[i].npu()).to(out_dtype).cpu()) | 89 | + gather_out_npu, x2_list_npu = gather_out.npu(), x2_list[i].npu() |
| 81 | - return out_list, gather_out | 90 | + if x1_scale_list: |
| 91 | + mm_res = torch_npu.npu_quant_matmul(x1=gather_out_npu, x2=x2_list_npu, scale=x2_scale_list[i].squeeze(0).npu(), pertoken_scale=x1_scale.squeeze(-1).npu(), output_dtype=out_dtype) | ||
| 92 | + elif x2_scale_list: | ||
| 93 | + mm_res = torch_npu.npu_quant_matmul(x1=gather_out_npu, x2=x2_list_npu, scale=x2_scale_list[i].squeeze(0).npu(), output_dtype=out_dtype) | ||
| 94 | + else: | ||
| 95 | + mm_res = torch.matmul(gather_out_npu, x2_list_npu) | ||
| 96 | + out_list.append(mm_res.to(out_dtype).cpu()) | ||
| 97 | + return out_list, gather_out_npu.cpu() | ||
| 82 | 98 | ||
| 83 | 99 | ||
| 84 | 100 | ||
| @@ -96,8 +112,32 @@ class TestAllGatherBaseMm(TestCase): | |||
| 96 | x1_list.append(x1) | 112 | x1_list.append(x1) |
| 97 | x2_list.append(x2) | 113 | x2_list.append(x2) |
| 98 | expt_out_list, expt_gather = self._construct_excepted_result(x1_list, x2_list, world_size) | 114 | expt_out_list, expt_gather = self._construct_excepted_result(x1_list, x2_list, world_size) |
| 115 | + for comm_mode in ['aiv', 'ai_cpu']: | ||
| 116 | + self._test_multiprocess(TestAllGatherBaseMm._test_npu_all_gather_base_mm, | ||
| 117 | + TestAllGatherBaseMm._init_dist_hccl, [expt_out_list, expt_gather, x1_list, x2_list, None, None, world_size, comm_mode, None]) | ||
| 118 | + | ||
| 119 | + | ||
| 120 | + | ||
| 121 | + def test_npu_all_gather_quant_mm(self): | ||
| 122 | + world_size = 8 | ||
| 123 | + m, k, n = 16, 512, 256 | ||
| 124 | + output_dtype = torch.float16 | ||
| 125 | + x1_list = [] | ||
| 126 | + x2_list = [] | ||
| 127 | + x1_scale_list = [] | ||
| 128 | + x2_scale_list = [] | ||
| 129 | + for _ in range(world_size): | ||
| 130 | + x1 = torch.randint(-10, 10, size=(m, k), dtype=torch.int8) | ||
| 131 | + x2 = torch.randint(-10, 10, size=(k, n), dtype=torch.int8) | ||
| 132 | + x1_scale = torch.randn((m, 1), dtype=torch.float32) | ||
| 133 | + x2_scale = torch.randn((1, n), dtype=torch.float32) | ||
| 134 | + x1_list.append(x1) | ||
| 135 | + x2_list.append(x2) | ||
| 136 | + x1_scale_list.append(x1_scale) | ||
| 137 | + x2_scale_list.append(x2_scale) | ||
| 138 | + expt_out_list, expt_gather = self._construct_excepted_result(x1_list, x2_list, world_size, x1_scale_list, x2_scale_list, output_dtype) | ||
| 99 | self._test_multiprocess(TestAllGatherBaseMm._test_npu_all_gather_base_mm, | 139 | self._test_multiprocess(TestAllGatherBaseMm._test_npu_all_gather_base_mm, |
| 100 | - TestAllGatherBaseMm._init_dist_hccl, [expt_out_list, expt_gather, x1_list, x2_list, world_size]) | 140 | + TestAllGatherBaseMm._init_dist_hccl, [expt_out_list, expt_gather, x1_list, x2_list, x1_scale_list, x2_scale_list, world_size, 'aiv', output_dtype]) |
| 101 | 141 | ||
| 102 | 142 | ||
| 103 | if __name__ == '__main__': | 143 | if __name__ == '__main__': |