在A5环境上输入数据类型是int类型时调用tensor.div_加rounding_mode=trunc；在A5环境跳过ND->NC1HWC0格式的转换。-pytorch-AtomGit

Mtest/distributed/test_allgather.py+5-4

	@@ -1,16 +1,15 @@
1	-import unittest
2	import os	1	import os
3	from random import randint	2	from random import randint
4	-
5	import numpy as np	3	import numpy as np
		4	+
6	import torch	5	import torch
7	import torch.distributed as dist	6	import torch.distributed as dist
8	import torch.multiprocessing as mp	7	import torch.multiprocessing as mp
9	-import torch_npu
10		8
		9	+import torch_npu
11	import torch_npu.distributed	10	import torch_npu.distributed
12	from torch_npu.testing.testcase import TestCase, run_tests	11	from torch_npu.testing.testcase import TestCase, run_tests
13	-from torch_npu.testing.common_utils import create_common_tensor	12	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
14	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	13	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
15		14
16		15
	@@ -155,6 +154,7 @@ class HcclAllGatherTest(HcclAllGatherTestBase):
155	pg.barrier()	154	pg.barrier()
156	p2c.get()	155	p2c.get()
157		156
		157	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
158	@skipIfUnsupportMultiNPU(2)	158	@skipIfUnsupportMultiNPU(2)
159	def test_all_gather_dist(self):	159	def test_all_gather_dist(self):
160	ranks = [2]	160	ranks = [2]
	@@ -174,6 +174,7 @@ class HcclAllGatherTest(HcclAllGatherTestBase):
174	self._test_multiprocess(HcclAllGatherTest._test_all_gather,	174	self._test_multiprocess(HcclAllGatherTest._test_all_gather,
175	HcclAllGatherTest._init_dist_hccl, expected, input1, world_size)	175	HcclAllGatherTest._init_dist_hccl, expected, input1, world_size)
176		176
		177	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
177	@skipIfUnsupportMultiNPU(2)	178	@skipIfUnsupportMultiNPU(2)
178	def test_all_gather_dist_different_shape(self):	179	def test_all_gather_dist_different_shape(self):
179	ranks = [2]	180	ranks = [2]

Mtest/distributed/test_allgather_base.py+4-8

	@@ -1,15 +1,10 @@
1	-import unittest
2	-import os
3	-from random import randint
4	-
5	import numpy as np	1	import numpy as np
		2	+
6	import torch	3	import torch
7	import torch.distributed as dist	4	import torch.distributed as dist
8	-import torch.multiprocessing as mp
9	-import torch_npu
10		5
11	-from torch_npu.testing.testcase import TestCase, run_tests	6	+from torch_npu.testing.testcase import run_tests
12	-from torch_npu.testing.common_utils import create_common_tensor	7	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
13	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	8	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
14		9
15	from test_allgather import HcclAllGatherTestBase	10	from test_allgather import HcclAllGatherTestBase
	@@ -29,6 +24,7 @@ class HcclAllGatherBaseTest(HcclAllGatherTestBase):
29	pg.barrier()	24	pg.barrier()
30	p2c.get()	25	p2c.get()
31		26
		27	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
32	@skipIfUnsupportMultiNPU(2)	28	@skipIfUnsupportMultiNPU(2)
33	def test_all_gather_base_dist(self):	29	def test_all_gather_base_dist(self):
34	ranks = [2]	30	ranks = [2]

Mtest/distributed/test_allgather_into_tensor.py+1-1

	@@ -57,6 +57,7 @@ class HcclAllGatherIntoTensorTest(HcclAllGatherTestBase):
57	with test_case.assertRaisesRegex(RuntimeError, error_expect):	57	with test_case.assertRaisesRegex(RuntimeError, error_expect):
58	pg.all_gather_into_tensor(gather_tensor, input1)	58	pg.all_gather_into_tensor(gather_tensor, input1)
59		59
		60	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
60	@skipIfUnsupportMultiNPU(2)	61	@skipIfUnsupportMultiNPU(2)
61	def test_all_gather_into_tensor_dist(self):	62	def test_all_gather_into_tensor_dist(self):
62	ranks = [2]	63	ranks = [2]
	@@ -74,7 +75,6 @@ class HcclAllGatherIntoTensorTest(HcclAllGatherTestBase):
74	self._test_multiprocess(HcclAllGatherIntoTensorTest._test_all_gather_into_tensor,	75	self._test_multiprocess(HcclAllGatherIntoTensorTest._test_all_gather_into_tensor,
75	HcclAllGatherIntoTensorTest._init_dist_hccl, expected, input1, world_size)	76	HcclAllGatherIntoTensorTest._init_dist_hccl, expected, input1, world_size)
76		77
77	- @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
78	@skipIfUnsupportMultiNPU(2)	78	@skipIfUnsupportMultiNPU(2)
79	def test_all_gather_into_tensor_dist_with_input_internal_format_and_offset(self):	79	def test_all_gather_into_tensor_dist_with_input_internal_format_and_offset(self):
80	ranks = [2]	80	ranks = [2]

Mtest/distributed/test_reduce.py+4-4

	@@ -1,14 +1,13 @@
1	-import unittest
2	import os	1	import os
3	-
4	import numpy as np	2	import numpy as np
		3	+
5	import torch	4	import torch
6	import torch.distributed as dist	5	import torch.distributed as dist
7	import torch.multiprocessing as mp	6	import torch.multiprocessing as mp
8	-import torch_npu
9		7
		8	+import torch_npu
10	from torch_npu.testing.testcase import TestCase, run_tests	9	from torch_npu.testing.testcase import TestCase, run_tests
11	-from torch_npu.testing.common_utils import create_common_tensor	10	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
12	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	11	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
13		12
14		13
	@@ -73,6 +72,7 @@ class HcclReduceTest(TestCase):
73		72
74	return expected	73	return expected
75		74
		75	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
76	@skipIfUnsupportMultiNPU(2)	76	@skipIfUnsupportMultiNPU(2)
77	def test_reduce_dist(self):	77	def test_reduce_dist(self):
78	ranks = [2, 4, 8]	78	ranks = [2, 4, 8]

Mtest/distributed/test_reduce_scatter.py+5-4

	@@ -1,15 +1,14 @@
1	-import unittest
2	import os	1	import os
3	from random import randint	2	from random import randint
4	-
5	import numpy as np	3	import numpy as np
		4	+
6	import torch	5	import torch
7	import torch.distributed as dist	6	import torch.distributed as dist
8	import torch.multiprocessing as mp	7	import torch.multiprocessing as mp
9	-import torch_npu
10		8
		9	+import torch_npu
11	from torch_npu.testing.testcase import TestCase, run_tests	10	from torch_npu.testing.testcase import TestCase, run_tests
12	-from torch_npu.testing.common_utils import create_common_tensor	11	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
13	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	12	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
14		13
15		14
	@@ -114,6 +113,7 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase):
114	with test_case.assertRaisesRegex(RuntimeError, error_expect):	113	with test_case.assertRaisesRegex(RuntimeError, error_expect):
115	pg.reduce_scatter(output, input_list_npu)	114	pg.reduce_scatter(output, input_list_npu)
116		115
		116	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
117	@skipIfUnsupportMultiNPU(2)	117	@skipIfUnsupportMultiNPU(2)
118	def test_reduce_scatter(self):	118	def test_reduce_scatter(self):
119	ranks = [2]	119	ranks = [2]
	@@ -135,6 +135,7 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase):
135	self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter,	135	self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter,
136	HcclReduceScatterTest._init_dist_hccl, expected, input_list, world_size)	136	HcclReduceScatterTest._init_dist_hccl, expected, input_list, world_size)
137		137
		138	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
138	@skipIfUnsupportMultiNPU(2)	139	@skipIfUnsupportMultiNPU(2)
139	def test_reduce_scatter_with_different_shape(self):	140	def test_reduce_scatter_with_different_shape(self):
140	ranks = [2]	141	ranks = [2]

Mtest/distributed/test_reduce_scatter_base.py+5-7

	@@ -1,14 +1,11 @@
1	-import unittest
2	-import os
3	-
4	import numpy as np	1	import numpy as np
		2	+
5	import torch	3	import torch
6	import torch.distributed as dist	4	import torch.distributed as dist
7	-import torch.multiprocessing as mp
8	-import torch_npu
9		5
10	-from torch_npu.testing.testcase import TestCase, run_tests	6	+import torch_npu
11	-from torch_npu.testing.common_utils import create_common_tensor	7	+from torch_npu.testing.testcase import run_tests
		8	+from torch_npu.testing.common_utils import create_common_tensor, SupportedDevices
12	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU	9	from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU
13		10
14	from test_reduce_scatter import HcclReduceScatterTestBase	11	from test_reduce_scatter import HcclReduceScatterTestBase
	@@ -28,6 +25,7 @@ class HcclReduceScatterBaseTest(HcclReduceScatterTestBase):
28	pg.barrier()	25	pg.barrier()
29	p2c.get()	26	p2c.get()
30		27
		28	+ @SupportedDevices(['Ascend910A', 'Ascend910B', 'Ascend910_93'])
31	@skipIfUnsupportMultiNPU(2)	29	@skipIfUnsupportMultiNPU(2)
32	def test_reduce_scatter_base(self):	30	def test_reduce_scatter_base(self):
33	ranks = [2]	31	ranks = [2]

Mtorch_npu/csrc/distributed/ProcessGroupHCCL.cpp+24-4

	@@ -4722,8 +4722,13 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::allreduce(
4722	[&](std::vector<c10_npu::NPUStream>& hcclStreams, c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>&) {	4722	[&](std::vector<c10_npu::NPUStream>& hcclStreams, c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>&) {
4723	if (opts.reduceOp == c10d::ReduceOp::AVG) {	4723	if (opts.reduceOp == c10d::ReduceOp::AVG) {
4724	c10_npu::NPUStreamGuard guard(hcclStreams[0]);	4724	c10_npu::NPUStreamGuard guard(hcclStreams[0]);
		4725	+ bool is_atlas_a5 = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950;
4725	for (auto& tensor : tensors_cp) {	4726	for (auto& tensor : tensors_cp) {
4726	- tensor.div_(getSize());	4727	+ if (is_atlas_a5 && at::isIntegralType(tensor.scalar_type(), /includeBool=/false)) {
		4728	+ tensor.div_(getSize(), "trunc");
		4729	+ } else {
		4730	+ tensor.div_(getSize());
		4731	+ }
4727	}	4732	}
4728	}	4733	}
4729	if (tensors_cp[0].scalar_type() != tensors[0].scalar_type()) {	4734	if (tensors_cp[0].scalar_type() != tensors[0].scalar_type()) {
	@@ -5057,8 +5062,13 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::reduce(
5057	[&](std::vector<c10_npu::NPUStream>& hcclStreams, c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>&) {	5062	[&](std::vector<c10_npu::NPUStream>& hcclStreams, c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>&) {
5058	if (opts.reduceOp == c10d::ReduceOp::AVG) {	5063	if (opts.reduceOp == c10d::ReduceOp::AVG) {
5059	c10_npu::NPUStreamGuard guard(hcclStreams[0]);	5064	c10_npu::NPUStreamGuard guard(hcclStreams[0]);
		5065	+ bool is_atlas_a5 = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950;
5060	for (auto& tensor : tensors_cp) {	5066	for (auto& tensor : tensors_cp) {
5061	- tensor.div_(getSize());	5067	+ if (is_atlas_a5 && at::isIntegralType(tensor.scalar_type(), /includeBool=/false)) {
		5068	+ tensor.div_(getSize(), "trunc");
		5069	+ } else {
		5070	+ tensor.div_(getSize());
		5071	+ }
5062	}	5072	}
5063	}	5073	}
5064	if (tensors_cp[0].scalar_type() != tensors[0].scalar_type()) {	5074	if (tensors_cp[0].scalar_type() != tensors[0].scalar_type()) {
	@@ -5875,8 +5885,13 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::reduce_scatter(
5875	[&](std::vector<c10_npu::NPUStream>& hcclStreams, c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>&) {	5885	[&](std::vector<c10_npu::NPUStream>& hcclStreams, c10::intrusive_ptr<ProcessGroupHCCL::WorkHCCL>&) {
5876	if (opts.reduceOp == c10d::ReduceOp::AVG) {	5886	if (opts.reduceOp == c10d::ReduceOp::AVG) {
5877	c10_npu::NPUStreamGuard guard(hcclStreams[0]);	5887	c10_npu::NPUStreamGuard guard(hcclStreams[0]);
		5888	+ bool is_atlas_a5 = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950;
5878	for (auto& tensor : outputTensors) {	5889	for (auto& tensor : outputTensors) {
5879	- tensor.div_(getSize());	5890	+ if (is_atlas_a5 && at::isIntegralType(tensor.scalar_type(), /includeBool=/false)) {
		5891	+ tensor.div_(getSize(), "trunc");
		5892	+ } else {
		5893	+ tensor.div_(getSize());
		5894	+ }
5880	}	5895	}
5881	}	5896	}
5882	},	5897	},
	@@ -5981,8 +5996,13 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::reduce_scatter(
5981	}	5996	}
5982	if (opts.reduceOp == c10d::ReduceOp::AVG) {	5997	if (opts.reduceOp == c10d::ReduceOp::AVG) {
5983	c10_npu::NPUStreamGuard guard(hcclStreams[0]);	5998	c10_npu::NPUStreamGuard guard(hcclStreams[0]);
		5999	+ bool is_atlas_a5 = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950;
5984	for (auto& tensor : outputTensors) {	6000	for (auto& tensor : outputTensors) {
5985	- tensor.div_(getSize());	6001	+ if (is_atlas_a5 && at::isIntegralType(tensor.scalar_type(), /includeBool=/false)) {
		6002	+ tensor.div_(getSize(), "trunc");
		6003	+ } else {
		6004	+ tensor.div_(getSize());
		6005	+ }
5986	}	6006	}
5987	}	6007	}
5988	},	6008	},