reduce scatter support tensorlist.size != world

Mtest/distributed/_fault_mode_cases/error_size.py+7-2

	@@ -10,8 +10,13 @@ def error_size():
10	dist.init_process_group(backend)	10	dist.init_process_group(backend)
11	rank = dist.get_rank()	11	rank = dist.get_rank()
12	torch.npu.set_device(rank)	12	torch.npu.set_device(rank)
13	- output = torch.tensor(2).npu()	13	+ ndev = torch.npu.device_count()
14	- input_list = [torch.tensor(var).npu() for var in range(3)]	14	+ # output on this rank's device, input_list tensors on a different device
		15	+ # -> reduce_scatter rejects input/output residing on different devices.
		16	+ out_dev = rank
		17	+ in_dev = (rank + 1) % ndev
		18	+ output = torch.zeros(4, dtype=torch.float32, device=f"npu:{out_dev}")
		19	+ input_list = [torch.zeros(4, dtype=torch.float32, device=f"npu:{in_dev}") for _ in range(2)]
15	dist.reduce_scatter(output, input_list)	20	dist.reduce_scatter(output, input_list)
16		21
17		22

Mtest/distributed/test_fault_mode.py+1-1

	@@ -117,7 +117,7 @@ class TestMode(TestCase):
117	process.terminate()	117	process.terminate()
118	process.wait()	118	process.wait()
119	self.assertIn(	119	self.assertIn(
120	- "Tensor list input to scatter/gather must match number of collective participants",	120	+ "Corresponding input/output tensors to reduce_scatter must all reside on the same device",
121	message	121	message
122	)	122	)
123		123

Mtest/distributed/test_reduce_scatter.py+105-0

	@@ -67,6 +67,51 @@ class HcclReduceScatterTestBase(TestCase):
67	return [input.cpu() for input in inputs]	67	return [input.cpu() for input in inputs]
68	return [input.cpu() * world_size for input in inputs]	68	return [input.cpu() * world_size for input in inputs]
69		69
		70	+ def _numel(self, shape):
		71	+ n = 1
		72	+ for d in shape:
		73	+ n *= d
		74	+ return n
		75	+
		76	+ # Expected output for _test_reduce_scatter_lifted. Input convention: the k-th
		77	+ # flattened element of rank r's i-th tensor is r*10000 + offset_i + k. After SUM,
		78	+ # global position p (< have) reduces to wsp + 5000ws*(ws-1); p >= have is 0.
		79	+ # rank r takes global [rout_numel, (r+1)out_numel).
		80	+ def _construct_lifted_expected(self, input_shapes, out_shape, world_size):
		81	+ offsets = [0]
		82	+ for s in input_shapes:
		83	+ offsets.append(offsets[-1] + self._numel(s))
		84	+ have = offsets[-1]
		85	+ out_numel = self._numel(out_shape)
		86	+ base = 5000 * world_size * (world_size - 1)
		87	+ expected = []
		88	+ for r in range(world_size):
		89	+ vals = []
		90	+ for k in range(out_numel):
		91	+ p = r * out_numel + k
		92	+ vals.append(world_size * p + base if p < have else 0.0)
		93	+ expected.append(torch.tensor(vals, dtype=torch.float32).reshape(out_shape))
		94	+ return expected
		95	+
		96	+ def _test_multiprocess_lifted(self, fn, init_pg, input_shapes, out_shape, world_size):
		97	+ ctx = mp.get_context('spawn')
		98	+ c2p = ctx.Queue(world_size)
		99	+ p2c = ctx.Queue(world_size)
		100	+ expected = self._construct_lifted_expected(input_shapes, out_shape, world_size)
		101	+ ps = []
		102	+ for i in range(world_size):
		103	+ p = ctx.Process(target=fn, args=(i, input_shapes, out_shape, world_size, init_pg, c2p, p2c))
		104	+ p.start()
		105	+ ps.append(p)
		106	+ for _ in range(world_size):
		107	+ rank, output = c2p.get()
		108	+ self.assertEqual(output, expected[rank],
		109	+ ("rank {} Expect receive tensor {} but got {}.").format(rank, expected[rank], output))
		110	+ for _ in range(world_size):
		111	+ p2c.put(0)
		112	+ for p in ps:
		113	+ p.join()
		114	+
70		115
71	class HcclReduceScatterTest(HcclReduceScatterTestBase):	116	class HcclReduceScatterTest(HcclReduceScatterTestBase):
72		117
	@@ -81,6 +126,28 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase):
81	pg.barrier()	126	pg.barrier()
82	p2c.get()	127	p2c.get()
83		128
		129	+ @classmethod
		130	+ # pylint:disable=huawei-too-many-arguments
		131	+ # input_shapes/out_shape are built per-rank with deterministic values so
		132	+ # the expected result can be computed locally.
		133	+ def _test_reduce_scatter_lifted(cls, rank, input_shapes, out_shape, world_size, init_pg, c2p, p2c,
		134	+ reduce_op=dist.ReduceOp.SUM):
		135	+ pg = init_pg(rank, world_size)
		136	+ input_list_npu = []
		137	+ offset = 0
		138	+ for s in input_shapes:
		139	+ n = 1
		140	+ for d in s:
		141	+ n *= d
		142	+ vals = torch.arange(offset, offset + n, dtype=torch.float32) + rank * 10000.0
		143	+ input_list_npu.append(vals.reshape(s).npu())
		144	+ offset += n
		145	+ output = torch.zeros(out_shape, dtype=torch.float32).npu()
		146	+ pg.reduce_scatter(output, input_list_npu, reduce_op)
		147	+ c2p.put((rank, output.cpu()))
		148	+ pg.barrier()
		149	+ p2c.get()
		150	+
84	@classmethod	151	@classmethod
85	# pylint:disable=huawei-too-many-arguments	152	# pylint:disable=huawei-too-many-arguments
86	def _test_reduce_scatter_with_input_internal_format_and_offset(cls, rank, input_list, world_size, init_pg):	153	def _test_reduce_scatter_with_input_internal_format_and_offset(cls, rank, input_list, world_size, init_pg):
	@@ -161,6 +228,44 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase):
161	self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter,	228	self._test_multiprocess(HcclReduceScatterTest._test_reduce_scatter,
162	HcclReduceScatterTest._init_dist_hccl, cpu_excepted_result, input_list, world_size)	229	HcclReduceScatterTest._init_dist_hccl, cpu_excepted_result, input_list, world_size)
163		230
		231	+ @SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
		232	+ @skipIfUnsupportMultiNPU(2)
		233	+ def test_reduce_scatter_single_tensor(self):
		234	+ # Single-tensor input list (length 1 != world_size). A single long tensor
		235	+ # of length world_size*out_numel is split evenly across ranks.
		236	+ ranks = [2]
		237	+ for world_size in ranks:
		238	+ out_shape = [4]
		239	+ input_shapes = [[world_size * 4]] # one tensor, len == world_size*out_numel
		240	+ self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted,
		241	+ HcclReduceScatterTest._init_dist_hccl, input_shapes, out_shape, world_size)
		242	+
		243	+ @SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
		244	+ @skipIfUnsupportMultiNPU(2)
		245	+ def test_reduce_scatter_input_list_not_equal_world_size(self):
		246	+ # input_list length differs from world_size: N < ws (zero pad) and N > ws (tail ignore).
		247	+ ranks = [2]
		248	+ for world_size in ranks:
		249	+ out_shape = [4]
		250	+ # N=1 < ws: have=4 < need=8, trailing rank gets zero pad
		251	+ self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted,
		252	+ HcclReduceScatterTest._init_dist_hccl, [[4]], out_shape, world_size)
		253	+ # N=3 > ws: have=12 > need=8, the 3rd tensor is ignored
		254	+ self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted,
		255	+ HcclReduceScatterTest._init_dist_hccl, [[4]] * 3, out_shape, world_size)
		256	+
		257	+ @SupportedDevices(['Ascend910B', 'Ascend910_93', 'Ascend950'])
		258	+ @skipIfUnsupportMultiNPU(2)
		259	+ def test_reduce_scatter_input_numel_not_equal_output(self):
		260	+ # N == world_size but per-tensor numel != output numel: have > need (tail ignore).
		261	+ ranks = [2]
		262	+ for world_size in ranks:
		263	+ out_shape = [4]
		264	+ # each tensor has 6 elements, output 4: have=12 > need=8, tail 2 elements per tensor ignored
		265	+ input_shapes = [[6]] * world_size
		266	+ self._test_multiprocess_lifted(HcclReduceScatterTest._test_reduce_scatter_lifted,
		267	+ HcclReduceScatterTest._init_dist_hccl, input_shapes, out_shape, world_size)
		268	+
164	@skipIfUnsupportMultiNPU(2)	269	@skipIfUnsupportMultiNPU(2)
165	def test_reduce_scatter_avg(self):	270	def test_reduce_scatter_avg(self):
166	ranks = [2]	271	ranks = [2]

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

	@@ -3320,6 +3320,32 @@ std::vector<at::Tensor> flatten_for_scatter_gather(
3320	return flattened;	3320	return flattened;
3321	}	3321	}
3322		3322
		3323	+// Flatten input tensor lists for reduce_scatter, allowing input list size
		3324	+// and per-tensor numel to differ from output (aligning with NCCL).
		3325	+std::vector<at::Tensor> flatten_for_reduce_scatter(
		3326	+ std::vector<std::vector<at::Tensor>>& tensor_lists,
		3327	+ std::vector<at::Tensor>& outputTensors)
		3328	+{
		3329	+ if (tensor_lists.size() != outputTensors.size()) {
		3330	+ TORCH_CHECK(false, "Tensor list operands to reduce_scatter must have the same length",
		3331	+ DIST_ERROR(ErrCode::VALUE));
		3332	+ }
		3333	+ const auto num_devices = tensor_lists.size();
		3334	+ std::vector<at::Tensor> flattened;
		3335	+ flattened.resize(num_devices);
		3336	+ for (const auto i : c10::irange(num_devices)) {
		3337	+ TORCH_CHECK(!tensor_lists[i].empty(),
		3338	+ "Tensor list operands to reduce_scatter must be non-empty",
		3339	+ DIST_ERROR(ErrCode::PARAM));
		3340	+ if (tensor_lists[i].front().get_device() != outputTensors[i].get_device()) {
		3341	+ TORCH_CHECK(false, "Corresponding input/output tensors to reduce_scatter must all reside"
		3342	+ " on the same device", DIST_ERROR(ErrCode::PARAM));
		3343	+ }
		3344	+ flattened[i] = c10d::newLikeFlat(tensor_lists, i);
		3345	+ }
		3346	+ return flattened;
		3347	+}
		3348	+
3323	void nslb_record_end()	3349	void nslb_record_end()
3324	{	3350	{
3325	std::string end_file_path;	3351	std::string end_file_path;
	@@ -5853,7 +5879,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupHCCL::reduce_scatter(
5853	}	5879	}
5854	bool same_size = check_same_size(inputTensors.back());	5880	bool same_size = check_same_size(inputTensors.back());
5855	if (same_size) {	5881	if (same_size) {
5856	- auto inputFlattened = flatten_for_scatter_gather(inputTensors, outputTensors, size_);	5882	+ auto inputFlattened = flatten_for_reduce_scatter(inputTensors, outputTensors);
5857	check_npu_tensors_different_devices(inputFlattened);	5883	check_npu_tensors_different_devices(inputFlattened);
5858	std::string functionName = __FUNCTION__;	5884	std::string functionName = __FUNCTION__;
5859	return collective(	5885	return collective(