fix recompute policy list overflow-MindSpeed-AtomGit

已合并

fix recompute policy list overflow #2657

wangyuansheng8创建于 2025年7月17日从refs/pull/2657/head 合入到 core_r0.8.0

fix recompute policy list overflow #2657

已合并

wangyuansheng8创建于 2025年7月17日

从refs/pull/2657/head合入到core_r0.8.0

共 2 个文件变更+5-4

Mmindspeed/core/memory/adaptive_recomputing/adaptive_recompute_solver.py+2-2

	@@ -152,10 +152,10 @@ class GraphSolver:
152	self.broadcast_in_mp_dp(policy_len_tensor, src)	152	self.broadcast_in_mp_dp(policy_len_tensor, src)
153	policy_len = tuple(policy_len_tensor.cpu().numpy().tolist())	153	policy_len = tuple(policy_len_tensor.cpu().numpy().tolist())
154	if global_rank == src:	154	if global_rank == src:
155	recompute_policy_tensor = torch.tensor(recompute_policy_list, dtype=torch.int8,	155	recompute_policy_tensor = torch.tensor(recompute_policy_list, dtype=torch.int32,
156	device=torch.npu.current_device())	156	device=torch.npu.current_device())
157	else:	157	else:
158	recompute_policy_tensor = torch.empty(policy_len, dtype=torch.int8,	158	recompute_policy_tensor = torch.empty(policy_len, dtype=torch.int32,
159	device=torch.npu.current_device())	159	device=torch.npu.current_device())
160		160
161	self.broadcast_in_mp_dp(recompute_policy_tensor, src)	161	self.broadcast_in_mp_dp(recompute_policy_tensor, src)

Mtests_extend/unit_tests/ops/cann/test_npu_grouped_mat_mul_all_reduce.py+3-2

	@@ -21,11 +21,11 @@ class TestNPUGroupedMatMulAllReduce(DistributedTest):
21	bias = None	21	bias = None
22	output_npu = torch_npu.npu_grouped_matmul(x, weight, bias=bias, scale=None,	22	output_npu = torch_npu.npu_grouped_matmul(x, weight, bias=bias, scale=None,
23	offset=None, antiquant_scale=None, antiquant_offset=None,	23	offset=None, antiquant_scale=None, antiquant_offset=None,
24	group_list=group_list, split_item=split_item)	24	group_list=group_list, split_item=split_item, group_type=-1)
25	else:	25	else:
26	output_npu = torch_npu.npu_grouped_matmul(x, weight, bias=bias, scale=[],	26	output_npu = torch_npu.npu_grouped_matmul(x, weight, bias=bias, scale=[],
27	offset=[], antiquant_scale=[], antiquant_offset=[],	27	offset=[], antiquant_scale=[], antiquant_offset=[],
28	group_list=group_list, split_item=split_item)	28	group_list=group_list, split_item=split_item, group_type=-1)
29		29
30	for i in range(outputs_num):	30	for i in range(outputs_num):
31	dist.all_reduce(output_npu[i], op=ReduceOp.SUM)	31	dist.all_reduce(output_npu[i], op=ReduceOp.SUM)
	@@ -48,6 +48,7 @@ class TestNPUGroupedMatMulAllReduce(DistributedTest):
48	hcomm_info = default_pg.get_hccl_comm_name(i)	48	hcomm_info = default_pg.get_hccl_comm_name(i)
49	return hcomm_info	49	return hcomm_info
50		50
		51	@pytest.mark.skip(reason="temporary skip for npu_grouped_matmul arguments change")
51	@pytest.mark.skipif(DEVICE_NAME != 'Ascend910B', reason='device type is not supported, skip this UT!')	52	@pytest.mark.skipif(DEVICE_NAME != 'Ascend910B', reason='device type is not supported, skip this UT!')
52	def test_npu_grouped_mat_mul_all_reduce(self):	53	def test_npu_grouped_mat_mul_all_reduce(self):
53	rank = int(os.environ["LOCAL_RANK"])	54	rank = int(os.environ["LOCAL_RANK"])