在分布式 Tensor Parallelism(TP)训练场景中,CANN MC2 算子(AllGatherMatmul / MatmulReduceScatter)在内核内部自行完成 AllGather / ReduceScatter 通信,性能优于分离式通信+计算方案。然而 HyperParallel 当前没有这两个算子的分布式调度逻辑,用 DTensor 对象调用时会直接失败(CANN 内核无法处理 DTensor),即使绕过也无法推断输出 DTensor 的分布状态,下游算子无法正确组合分布式计算图。
为以下两个 MindSpore MC2 算子提供分布式调度支持:
preprocess()
to_local()
infer_layout()
Partial(sum)
LinearDistributedOp
gather_output=False
支持的分片模式:
两个算子均继承 DistributedOp 基类,实现 preprocess 和 infer_layout,通过 YAML 注册。 不需要 get_expand_impl 覆盖(基类默认返回 None)。
preprocess
infer_layout
get_expand_impl
Partial 说明:AllGatherMatmul 在 k 维切分时输出带 Partial(sum),full_tensor() 自动触发 AllReduce;MatmulReduceScatter 不需要 Partial(CANN 内核已内部完成 ReduceScatter,结果完整)。
full_tensor()
约束:
UT:test_parallel_all_gather_matmul.py(15 cases)、test_parallel_matmul_reduce_scatter.py(12 cases) ST:4 个 MindSpore 分布式测试文件(AGM 3 组,MRS 3 组,含 8 卡 mnk 全切场景)
test_parallel_all_gather_matmul.py
test_parallel_matmul_reduce_scatter.py
#656
ISSUE: 接入 MindSpore MC2 通算融合算子分布式调度(AllGatherMatmul / MatmulReduceScatter)
背景
在分布式 Tensor Parallelism(TP)训练场景中,CANN MC2 算子(AllGatherMatmul / MatmulReduceScatter)在内核内部自行完成 AllGather / ReduceScatter 通信,性能优于分离式通信+计算方案。然而 HyperParallel 当前没有这两个算子的分布式调度逻辑,用 DTensor 对象调用时会直接失败(CANN 内核无法处理 DTensor),即使绕过也无法推断输出 DTensor 的分布状态,下游算子无法正确组合分布式计算图。
目标
为以下两个 MindSpore MC2 算子提供分布式调度支持:
AllGatherMatmul
preprocess()的to_local()提取本地 tensor 传给 CANNinfer_layout()推断输出 layout:Partial(sum),调用方需 AllReduce(与LinearDistributedOpcontract_dim 分片语义一致)gather_output=False时 CANN 返回空 tensor,gather_out layout 强制全 Replicate,避免 Shard dim 越界支持的分片模式:
MatmulReduceScatter
preprocess()的to_local()提取本地 tensor 传给 CANNinfer_layout()推断输出 layout:支持的分片模式:
两个算子均继承 DistributedOp 基类,实现
preprocess和infer_layout,通过 YAML 注册。不需要
get_expand_impl覆盖(基类默认返回 None)。Partial 说明:AllGatherMatmul 在 k 维切分时输出带
Partial(sum),full_tensor()自动触发 AllReduce;MatmulReduceScatter 不需要 Partial(CANN 内核已内部完成 ReduceScatter,结果完整)。约束:
UT:
test_parallel_all_gather_matmul.py(15 cases)、test_parallel_matmul_reduce_scatter.py(12 cases)ST:4 个 MindSpore 分布式测试文件(AGM 3 组,MRS 3 组,含 8 卡 mnk 全切场景)
相关 PR
#656