A5 dtype: make checkSupportedDataType accept uint64/fp64 on Ascend950+-pytorch-AtomGit

Mtest/distributed/test_allreduce.py+44-5

	@@ -8,7 +8,7 @@ import torch.multiprocessing as mp
8		8
9	import torch_npu	9	import torch_npu
10	from torch_npu.testing.testcase import TestCase, run_tests	10	from torch_npu.testing.testcase import TestCase, run_tests
11	-from torch_npu.testing.common_utils import create_common_tensor	11	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
12	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	12	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
13		13
14		14
	@@ -64,7 +64,7 @@ class HcomAllReduceTest(TestCase):
64	self.assertEqual(p.exitcode, 0, "subprocess exit with abnormal code.")	64	self.assertEqual(p.exitcode, 0, "subprocess exit with abnormal code.")
65		65
66	# pylint:disable=huawei-too-many-arguments	66	# pylint:disable=huawei-too-many-arguments
67	- def _test_multiprocess(self, f, init_pg, expected, input1, world_size, reduce_op=dist.ReduceOp.SUM):	67	+ def _test_multiprocess(self, f, init_pg, expected, input1, world_size, reduce_op=dist.ReduceOp.SUM, use_equal=False):
68	ctx = mp.get_context('spawn')	68	ctx = mp.get_context('spawn')
69	c2p = ctx.Queue(world_size)	69	c2p = ctx.Queue(world_size)
70	done_event = ctx.Event()	70	done_event = ctx.Event()
	@@ -80,9 +80,17 @@ class HcomAllReduceTest(TestCase):
80	for _ in range(world_size):	80	for _ in range(world_size):
81	rank, dst, output = c2p.get()	81	rank, dst, output = c2p.get()
82	if rank == dst:	82	if rank == dst:
83	- self.assertEqual(output, expected,	83	+ # torch has no add/mul stub for UInt64, so assertEqual (which
84	- "rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(	84	+ # computes output - expected internally) raises. Use torch.equal
85	- rank, world_size, expected.dtype, expected.shape, expected, output))	85	+ # for uint64, which compares without subtraction.
		86	+ if use_equal:
		87	+ self.assertTrue(torch.equal(output, expected),
		88	+ "rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(
		89	+ rank, world_size, expected.dtype, expected.shape, expected, output))
		90	+ else:
		91	+ self.assertEqual(output, expected,
		92	+ "rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(
		93	+ rank, world_size, expected.dtype, expected.shape, expected, output))
86	done_event.set()	94	done_event.set()
87	for p in ps:	95	for p in ps:
88	p.join()	96	p.join()
	@@ -127,6 +135,37 @@ class HcomAllReduceTest(TestCase):
127	self._test_multiprocess(HcomAllReduceTest._test_all_reduce,	135	self._test_multiprocess(HcomAllReduceTest._test_all_reduce,
128	HcomAllReduceTest._init_dist_hccl, expected, input1, world_size)	136	HcomAllReduceTest._init_dist_hccl, expected, input1, world_size)
129		137
		138	+ # Ascend950 (Atlas A5) extends HCCL data type support with uint64/fp64.
		139	+ @SupportedDevices(["Ascend950"])
		140	+ @skipIfUnsupportMultiNPU(2)
		141	+ def test_dist_all_reduce_uint64(self):
		142	+ # CI currently supports only 2 devices
		143	+ ranks = [2]
		144	+ shape_format = [np.uint64, 2, [2, 3, 16]]
		145	+ for world_size in ranks:
		146	+ # uint64 is unsigned, use a non-negative range to avoid wrap-around on cast.
		147	+ exp_input, input1 = create_common_tensor(shape_format, 0, 10)
		148	+ # torch has no add/mul stub for UInt64, so compute the expected
		149	+ # result with numpy (which supports uint64 arithmetic) and cast
		150	+ # back to a tensor for comparison.
		151	+ expected_np = self._construct_excepted_result(exp_input.numpy(), world_size)
		152	+ expected = torch.from_numpy(expected_np)
		153	+ self._test_multiprocess(HcomAllReduceTest._test_all_reduce,
		154	+ HcomAllReduceTest._init_dist_hccl, expected, input1, world_size,
		155	+ use_equal=True)
		156	+
		157	+ @SupportedDevices(["Ascend950"])
		158	+ @skipIfUnsupportMultiNPU(2)
		159	+ def test_dist_all_reduce_fp64(self):
		160	+ # CI currently supports only 2 devices
		161	+ ranks = [2]
		162	+ shape_format = [np.float64, 2, [2, 3, 16]]
		163	+ for world_size in ranks:
		164	+ exp_input, input1 = create_common_tensor(shape_format, -10, 10)
		165	+ expected = self._construct_excepted_result(exp_input, world_size)
		166	+ self._test_multiprocess(HcomAllReduceTest._test_all_reduce,
		167	+ HcomAllReduceTest._init_dist_hccl, expected, input1, world_size)
		168	+
130	@skipIfUnsupportMultiNPU(2)	169	@skipIfUnsupportMultiNPU(2)
131	def test_dist_all_reduce_avg(self):	170	def test_dist_all_reduce_avg(self):
132	# CI currently supports only 2 devices	171	# CI currently supports only 2 devices

Mtest/distributed/test_reduce.py+45-4

	@@ -33,7 +33,7 @@ class HcclReduceTest(TestCase):
33	p2c.get()	33	p2c.get()
34		34
35	# pylint:disable=huawei-too-many-arguments	35	# pylint:disable=huawei-too-many-arguments
36	- def _test_multiprocess(self, f, init_pg, expected, input1, world_size, reduce_op=dist.ReduceOp.SUM):	36	+ def _test_multiprocess(self, f, init_pg, expected, input1, world_size, reduce_op=dist.ReduceOp.SUM, use_equal=False):
37	ctx = mp.get_context('spawn')	37	ctx = mp.get_context('spawn')
38	c2p = ctx.Queue(world_size)	38	c2p = ctx.Queue(world_size)
39	p2c = ctx.Queue(world_size)	39	p2c = ctx.Queue(world_size)
	@@ -49,9 +49,17 @@ class HcclReduceTest(TestCase):
49	for _ in range(world_size):	49	for _ in range(world_size):
50	rank, dst, output = c2p.get()	50	rank, dst, output = c2p.get()
51	if rank == dst:	51	if rank == dst:
52	- self.assertEqual(output, expected,	52	+ # torch has no add/mul stub for UInt64, so assertEqual (which
53	- "rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(	53	+ # computes output - expected internally) raises. Use torch.equal
54	- rank, world_size, expected.dtype, expected.shape, expected, output))	54	+ # for uint64, which compares without subtraction.
		55	+ if use_equal:
		56	+ self.assertTrue(torch.equal(output, expected),
		57	+ "rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(
		58	+ rank, world_size, expected.dtype, expected.shape, expected, output))
		59	+ else:
		60	+ self.assertEqual(output, expected,
		61	+ "rank {} world_size {} dtype {} shape {} Expect receive tensor {} but got {}.".format(
		62	+ rank, world_size, expected.dtype, expected.shape, expected, output))
55		63
56	for _ in range(world_size):	64	for _ in range(world_size):
57	p2c.put(0)	65	p2c.put(0)
	@@ -92,6 +100,39 @@ class HcclReduceTest(TestCase):
92	self._test_multiprocess(HcclReduceTest._test_reduce,	100	self._test_multiprocess(HcclReduceTest._test_reduce,
93	HcclReduceTest._init_dist_hccl, expected, input1, world_size)	101	HcclReduceTest._init_dist_hccl, expected, input1, world_size)
94		102
		103	+ # Ascend950 (Atlas A5) extends HCCL data type support with uint64/fp64.
		104	+ @SupportedDevices(["Ascend950"])
		105	+ @skipIfUnsupportMultiNPU(2)
		106	+ def test_reduce_uint64_dist(self):
		107	+ ranks = [2]
		108	+ shape_format = [np.uint64, 2, [12, 56, 256]]
		109	+ for world_size in ranks:
		110	+ if torch.npu.device_count() < world_size:
		111	+ continue
		112	+ # uint64 is unsigned, use a non-negative range to avoid wrap-around on cast.
		113	+ exp_input, input1 = create_common_tensor(shape_format, 0, 10)
		114	+ # torch has no add/mul stub for UInt64, so compute the expected
		115	+ # result with numpy (which supports uint64 arithmetic) and cast
		116	+ # back to a tensor for comparison.
		117	+ expected_np = self._construct_excepted_result(exp_input.numpy(), world_size)
		118	+ expected = torch.from_numpy(expected_np)
		119	+ self._test_multiprocess(HcclReduceTest._test_reduce,
		120	+ HcclReduceTest._init_dist_hccl, expected, input1, world_size,
		121	+ use_equal=True)
		122	+
		123	+ @SupportedDevices(["Ascend950"])
		124	+ @skipIfUnsupportMultiNPU(2)
		125	+ def test_reduce_fp64_dist(self):
		126	+ ranks = [2]
		127	+ shape_format = [np.float64, 2, [12, 56, 256]]
		128	+ for world_size in ranks:
		129	+ if torch.npu.device_count() < world_size:
		130	+ continue
		131	+ exp_input, input1 = create_common_tensor(shape_format, -10, 10)
		132	+ expected = self._construct_excepted_result(exp_input, world_size)
		133	+ self._test_multiprocess(HcclReduceTest._test_reduce,
		134	+ HcclReduceTest._init_dist_hccl, expected, input1, world_size)
		135	+
95	@skipIfUnsupportMultiNPU(2)	136	@skipIfUnsupportMultiNPU(2)
96	def test_reduce_uint8_dist(self):	137	def test_reduce_uint8_dist(self):
97	ranks = [2]	138	ranks = [2]

Mtest/distributed/test_reduce_scatter.py+47-3

	@@ -24,7 +24,7 @@ class HcclReduceScatterTestBase(TestCase):
24	return dist	24	return dist
25		25
26	# pylint:disable=huawei-too-many-arguments	26	# pylint:disable=huawei-too-many-arguments
27	- def _test_multiprocess(self, fn, init_pg, expected, input1, world_size, reduce_op=dist.ReduceOp.SUM):	27	+ def _test_multiprocess(self, fn, init_pg, expected, input1, world_size, reduce_op=dist.ReduceOp.SUM, use_equal=False):
28	ctx = mp.get_context('spawn')	28	ctx = mp.get_context('spawn')
29	c2p = ctx.Queue(world_size)	29	c2p = ctx.Queue(world_size)
30	p2c = ctx.Queue(world_size)	30	p2c = ctx.Queue(world_size)
	@@ -38,8 +38,15 @@ class HcclReduceScatterTestBase(TestCase):
38	ps.append(p)	38	ps.append(p)
39	for _ in range(world_size):	39	for _ in range(world_size):
40	rank, output = c2p.get()	40	rank, output = c2p.get()
41	- self.assertEqual(output, expected[rank],	41	+ # torch has no add/mul stub for UInt64, so assertEqual (which
42	- ("rank {} Expect receive tensor {} but got {}.").format(rank, expected[rank], output))	42	+ # computes output - expected internally) raises. Use torch.equal
		43	+ # for uint64, which compares without subtraction.
		44	+ if use_equal:
		45	+ self.assertTrue(torch.equal(output, expected[rank]),
		46	+ ("rank {} Expect receive tensor {} but got {}.").format(rank, expected[rank], output))
		47	+ else:
		48	+ self.assertEqual(output, expected[rank],
		49	+ ("rank {} Expect receive tensor {} but got {}.").format(rank, expected[rank], output))
43		50
44	for _ in range(world_size):	51	for _ in range(world_size):
45	p2c.put(0)	52	p2c.put(0)
	@@ -266,6 +273,43 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase):
266	self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted,	273	self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted,
267	HcclReduceScatterTest._init_dist_hccl, input_shapes, out_shape, world_size)	274	HcclReduceScatterTest._init_dist_hccl, input_shapes, out_shape, world_size)
268		275
		276	+ # Ascend950 (Atlas A5) extends HCCL data type support with uint64/fp64.
		277	+ @SupportedDevices(["Ascend950"])
		278	+ @skipIfUnsupportMultiNPU(2)
		279	+ def test_reduce_scatter_uint64(self):
		280	+ ranks = [2]
		281	+ shape_format = [[np.uint64, 2, [4, 9]]]
		282	+ for world_size in ranks:
		283	+ for shape in shape_format:
		284	+ input_list = []
		285	+ for _ in range(world_size):
		286	+ # uint64 is unsigned, use a non-negative range to avoid wrap-around on cast.
		287	+ _, input1 = create_common_tensor(shape, 0, 10)
		288	+ input_list.append(input1.cpu())
		289	+ # _construct_excepted_result uses input.cpu()*world_size (mul,
		290	+ # which works for uint64), so pass tensors directly. Only the
		291	+ # final comparison needs torch.equal (assertEqual does a-b which
		292	+ # uint64 lacks), via use_equal=True.
		293	+ expected = self._construct_excepted_result(input_list, world_size, dist.reduce_scatter)
		294	+ self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter,
		295	+ HcclReduceScatterTest._init_dist_hccl, expected, input_list, world_size,
		296	+ use_equal=True)
		297	+
		298	+ @SupportedDevices(["Ascend950"])
		299	+ @skipIfUnsupportMultiNPU(2)
		300	+ def test_reduce_scatter_fp64(self):
		301	+ ranks = [2]
		302	+ shape_format = [[np.float64, 2, [4, 9]]]
		303	+ for world_size in ranks:
		304	+ for shape in shape_format:
		305	+ input_list = []
		306	+ for _ in range(world_size):
		307	+ _, input1 = create_common_tensor(shape, -10, 10)
		308	+ input_list.append(input1.cpu())
		309	+ expected = self._construct_excepted_result(input_list, world_size, dist.reduce_scatter)
		310	+ self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter,
		311	+ HcclReduceScatterTest._init_dist_hccl, expected, input_list, world_size)
		312	+
269	@skipIfUnsupportMultiNPU(2)	313	@skipIfUnsupportMultiNPU(2)
270	def test_reduce_scatter_avg(self):	314	def test_reduce_scatter_avg(self):
271	ranks = [2]	315	ranks = [2]

Mtest/distributed/test_reduce_scatter_base.py+37-0

	@@ -64,6 +64,43 @@ class HcclReduceScatterBaseTest(HcclReduceScatterTestBase):
64	self._test_multiprocess(HcclReduceScatterBaseTest._test_reduce_scatter_base,	64	self._test_multiprocess(HcclReduceScatterBaseTest._test_reduce_scatter_base,
65	HcclReduceScatterBaseTest._init_dist_hccl, expected, input_list, world_size)	65	HcclReduceScatterBaseTest._init_dist_hccl, expected, input_list, world_size)
66		66
		67	+ # Ascend950 (Atlas A5) extends HCCL data type support with uint64/fp64.
		68	+ @SupportedDevices(["Ascend950"])
		69	+ @skipIfUnsupportMultiNPU(2)
		70	+ def test_reduce_scatter_base_uint64(self):
		71	+ ranks = [2]
		72	+ shape_format = [[np.uint64, 2, [4, 9]]]
		73	+ for world_size in ranks:
		74	+ for shape in shape_format:
		75	+ input_list = []
		76	+ for _ in range(world_size):
		77	+ # uint64 is unsigned, use a non-negative range to avoid wrap-around on cast.
		78	+ _, input1 = create_common_tensor(shape, 0, 10)
		79	+ input_list.append(input1.cpu())
		80	+ # _construct_excepted_result uses input.cpu()*world_size (mul,
		81	+ # which works for uint64), so pass tensors directly. Only the
		82	+ # final comparison needs torch.equal (assertEqual does a-b which
		83	+ # uint64 lacks), via use_equal=True.
		84	+ expected = self._construct_excepted_result(input_list, world_size, dist._reduce_scatter_base)
		85	+ self._test_multiprocess(HcclReduceScatterBaseTest._test_reduce_scatter_base,
		86	+ HcclReduceScatterBaseTest._init_dist_hccl, expected, input_list, world_size,
		87	+ use_equal=True)
		88	+
		89	+ @SupportedDevices(["Ascend950"])
		90	+ @skipIfUnsupportMultiNPU(2)
		91	+ def test_reduce_scatter_base_fp64(self):
		92	+ ranks = [2]
		93	+ shape_format = [[np.float64, 2, [4, 9]]]
		94	+ for world_size in ranks:
		95	+ for shape in shape_format:
		96	+ input_list = []
		97	+ for _ in range(world_size):
		98	+ _, input1 = create_common_tensor(shape, -10, 10)
		99	+ input_list.append(input1.cpu())
		100	+ expected = self._construct_excepted_result(input_list, world_size, dist._reduce_scatter_base)
		101	+ self._test_multiprocess(HcclReduceScatterBaseTest._test_reduce_scatter_base,
		102	+ HcclReduceScatterBaseTest._init_dist_hccl, expected, input_list, world_size)
		103	+
67	@skipIfUnsupportMultiNPU(2)	104	@skipIfUnsupportMultiNPU(2)
68	def test_reduce_scatter_base_avg(self):	105	def test_reduce_scatter_base_avg(self):
69	ranks = [2]	106	ranks = [2]

Mtest/distributed/test_reduce_scatter_tensor.py+37-0

	@@ -79,6 +79,43 @@ class HcclReduceScatterTensorTest(HcclReduceScatterTestBase):
79	self._test_multiprocess(HcclReduceScatterTensorTest._test_reduce_scatter_tensor,	79	self._test_multiprocess(HcclReduceScatterTensorTest._test_reduce_scatter_tensor,
80	HcclReduceScatterTensorTest._init_dist_hccl, expected, input_list, world_size)	80	HcclReduceScatterTensorTest._init_dist_hccl, expected, input_list, world_size)
81		81
		82	+ # Ascend950 (Atlas A5) extends HCCL data type support with uint64/fp64.
		83	+ @SupportedDevices(["Ascend950"])
		84	+ @skipIfUnsupportMultiNPU(2)
		85	+ def test_reduce_scatter_tensor_uint64(self):
		86	+ ranks = [2]
		87	+ shape_format = [[np.uint64, 2, [4, 9]]]
		88	+ for world_size in ranks:
		89	+ for shape in shape_format:
		90	+ input_list = []
		91	+ for _ in range(world_size):
		92	+ # uint64 is unsigned, use a non-negative range to avoid wrap-around on cast.
		93	+ _, input1 = create_common_tensor(shape, 0, 10)
		94	+ input_list.append(input1.cpu())
		95	+ # _construct_excepted_result uses input.cpu()*world_size (mul,
		96	+ # which works for uint64), so pass tensors directly. Only the
		97	+ # final comparison needs torch.equal (assertEqual does a-b which
		98	+ # uint64 lacks), via use_equal=True.
		99	+ expected = self._construct_excepted_result(input_list, world_size, dist.reduce_scatter_tensor)
		100	+ self._test_multiprocess(HcclReduceScatterTensorTest._test_reduce_scatter_tensor,
		101	+ HcclReduceScatterTensorTest._init_dist_hccl, expected, input_list, world_size,
		102	+ use_equal=True)
		103	+
		104	+ @SupportedDevices(["Ascend950"])
		105	+ @skipIfUnsupportMultiNPU(2)
		106	+ def test_reduce_scatter_tensor_fp64(self):
		107	+ ranks = [2]
		108	+ shape_format = [[np.float64, 2, [4, 9]]]
		109	+ for world_size in ranks:
		110	+ for shape in shape_format:
		111	+ input_list = []
		112	+ for _ in range(world_size):
		113	+ _, input1 = create_common_tensor(shape, -10, 10)
		114	+ input_list.append(input1.cpu())
		115	+ expected = self._construct_excepted_result(input_list, world_size, dist.reduce_scatter_tensor)
		116	+ self._test_multiprocess(HcclReduceScatterTensorTest._test_reduce_scatter_tensor,
		117	+ HcclReduceScatterTensorTest._init_dist_hccl, expected, input_list, world_size)
		118	+
82	@classmethod	119	@classmethod
83	# pylint:disable=huawei-too-many-arguments	120	# pylint:disable=huawei-too-many-arguments
84	def _test_reduce_scatter_tensor_uneven(cls, rank, input_list, world_size, init_pg, c2p, p2c, reduce_op=dist.ReduceOp.SUM):	121	def _test_reduce_scatter_tensor_uneven(cls, rank, input_list, world_size, init_pg, c2p, p2c, reduce_op=dist.ReduceOp.SUM):

Mtorch_npu/csrc/distributed/ProcessGroupHCCL.cpp+9-1

	@@ -170,6 +170,7 @@ HcclReduceOp getHcclReduceOp(const c10d::ReduceOp reduceOp, at::Tensor& input)
170	}	170	}
171		171
172	// AllGather & Broadcast support all data type, no need do more check.	172	// AllGather & Broadcast support all data type, no need do more check.
		173	+// On Ascend950 (Atlas A5), HCCL additionally supports uint64/fp64.
173	void checkSupportedDataType(HcclDataType type, std::string functionName)	174	void checkSupportedDataType(HcclDataType type, std::string functionName)
174	{	175	{
175	static std::set<HcclDataType> supportedDataTypes = {	176	static std::set<HcclDataType> supportedDataTypes = {
	@@ -180,8 +181,15 @@ void checkSupportedDataType(HcclDataType type, std::string functionName)
180	HCCL_DATA_TYPE_FP32,	181	HCCL_DATA_TYPE_FP32,
181	HCCL_DATA_TYPE_BFP16,	182	HCCL_DATA_TYPE_BFP16,
182	HCCL_DATA_TYPE_INT64};	183	HCCL_DATA_TYPE_INT64};
		184	+ // Ascend950 (Atlas A5) extends HCCL data type support with uint64/fp64.
		185	+ static std::set<HcclDataType> a5ExtraDataTypes = {
		186	+ HCCL_DATA_TYPE_UINT64,
		187	+ HCCL_DATA_TYPE_FP64};
		188	+ bool is_atlas_a5 = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950;
		189	+ bool supported = (supportedDataTypes.count(type) != 0) \|\|
		190	+ (is_atlas_a5 && a5ExtraDataTypes.count(type) != 0);
183	TORCH_CHECK(	191	TORCH_CHECK(
184	- supportedDataTypes.count(type) != 0,	192	+ supported,
185	"HCCL "+functionName+": Unsupported data type ",	193	"HCCL "+functionName+": Unsupported data type ",
186	getHcclDataTypeSerialString(type), DIST_ERROR(ErrCode::NOT_SUPPORT));	194	getHcclDataTypeSerialString(type), DIST_ERROR(ErrCode::NOT_SUPPORT));
187	}	195	}