基于Ascend C的MC²通算融合算子性能优化最佳实践

举报
昇腾CANN 发表于 2025/04/08 09:58:11 2025/04/08
【摘要】 根据不同的应用场景,昇腾CANN的算子加速库内置了多种MC²通算融合类算子,典型算子包括AllGatherMatMul、MatMulReduceScatter、MatMulAllReduce等。这些算子将AllReduce、ReduceScatter、AllGather等通信过程与MatMul计算进行融合,使SP、TP等并行场景下的通信和矩阵计算实现流水并行,从而加速大模型的执行。

近年来,随着GPT-3、LLaMA、DeepSeek等模型的参数规模突破千亿甚至万亿级别,AI大模型已成为驱动人工智能技术革新的核心引擎。然而,模型规模的指数级增长使得单设备训练在计算能力、存储容量和能效等方面面临根本性瓶颈。以1750亿参数的GPT-3模型为例,其训练至少需要数百GB的内存资源,这使得分布式并行推理/训练从可选方案转变为必选技术路径。

在分布式推理/训练框架中,如何平衡计算效率与通信开销成为关键挑战。传统的单维度并行策略(如数据并行)已无法满足需求,而混合并行(Hybrid Parallelism)通过同时切分计算图(模型并行)、数据批次(数据并行)以及流水线阶段(流水线并行),成为当前主流解决方案。其中,Cube计算和集合通信算子的协同切分与并行策略,是决定系统吞吐量的关键因素。昇腾CANN通过将需要切分并行的计算和通信算子进行融合(又称:通算融合),实现了计算和通信的流水,从而并行提升性能。融合后的算子简称为MC2通算融合算子(Matrix Computation & Communication)。

根据不同的应用场景,昇腾CANN的算子加速库内置了多种MC²通算融合类算子,典型算子包括AllGatherMatMul、MatMulReduceScatter、MatMulAllReduce等。这些算子将AllReduce、ReduceScatter、AllGather等通信过程与MatMul计算进行融合,使SP、TP等并行场景下的通信和矩阵计算实现流水并行,从而加速大模型的执行。

此外,基于Ascend C算子编程语言,开发者可以进一步对MC²算子实施性能优化,以满足更加严苛的需求场景。

本文将从开发者角度出发,系统阐述面向大模型推理/训练场景的MC2算子性能优化方法,重点解决以下问题:

  1. 如何进行通算融合性能评估:根据输入张量的Shape、融合前单独计算和通信的耗时情况等因素,分析通算融合后的预期性能收益及可能存在的瓶颈点。
  2. 如何制定数据切分策略:通过理论分析和实际示例,说明如何将输入张量进行切分,形成多条数据流水并行,以达到最优性能目标。

优化后的完整样例可参考MatmulAllReduce样例MatmulReduceScatter样例AllGatherMatmul样例

MC2算子性能瓶颈点分析

MC2通算融合算子的性能收益主要来自于通信、计算的并行执行,即将输入数据切分为多个子块,子块的计算和通信任务形成两条流水线,通过两条流水线上任务的并行执行,实现流水掩盖,从而提升算子性能。如下图所示,MC2算子先做Matmul计算、后通信的场景,输入矩阵沿M轴被切分为两块,第二块数据的Matmul计算和第一块数据的通信可以并行执行,从而达到计算和通信时间相互掩盖的目的。本节的所有图示中MM代表Matmul计算,hcom代表通信任务。

 1.PNG

MC2算子性能收益公式为:

1-1.png

融合后MC2算子的执行耗时,受以下因素制约,从而影响算子性能收益。

1、因素一:计算和通信的执行时间差异

若计算和通信任务的执行时间相差不大,则融合后,MC2算子的计算和通信并行执行,能得到较好的流水掩盖,性能收益较大。

若计算和通信任务的执行时间差异较大,则融合后,MC2算子内计算和通信并行执行,能够掩盖的时间较少,算子整体执行耗时与未切分串行时的算子执行耗时接近,此时无法获得较大性能收益。

2.png

2、因素二:数据切分导致的计算或通信的执行时间膨胀

当对输入数据进行切分后,原本的整块数据被切分成若干小数据块,对若干小数据块分别做Matmul计算或者执行通信任务,此时相比切分前,计算或者通信任务的执行时间可能发生膨胀(即执行时间变长)。该膨胀产生的原因包括:切分后的数据块过小导致计算或通信的效率降低、切分的数据块过多导致增加额外的调度开销、并行执行后计算和通信对L2Cache或device存储内存的访问冲突等。以Matmul计算为例,简单说明数据切分后执行时间可能发生的膨胀情况。

(1)未发生膨胀:

数据切分前,Matmul执行时间为200us,将Matmul的输入均匀切分为两块,假设切分后,每块数据的Matmul执行时间都是100us,通过计算的并行执行,下图实际性能收益为100us。

3.png

(2)发生一般程度的膨胀:

数据切分前,Matmul执行时间为200us,将Matmul的输入均匀切分为两块,假设切分后,每块数据的Matmul执行时间都是150us,通过计算的并行执行,下图实际性能收益为50us。

4.png

(3)发生严重程度的膨胀:

数据切分前,Matmul执行时间为200us,将Matmul的输入均匀切分为两块,假设切分后,每块数据的Matmul执行时间都是200us,通过计算的并行执行,下图实际性能收益为劣化50us。

5.png

综合上述分析,计算和通信执行时间较均衡的场景,有更好的流水掩盖和性能收益;同时,性能收益也受到数据切分导致的执行时间膨胀的影响。下文将介绍如何制定数据切分策略,以达到最佳流水掩盖效果。

设计优化方案

以Atlas A2 训练系列产品/Atlas 800I A2 推理产品,输入数据格式为ND、数据类型为half的MatmulAllreduce算子为例 ,该算子中计算执行在前,通信任务执行在后。假定Matmul计算中左矩阵的形状为[M, K],右矩阵的形状为[K, N],算子中的通信对象为Matmul的输出矩阵,则通信任务的输入shape为[M, N]。因为K轴只在Matmul计算中存在,当K轴较大时,计算量大,计算执行时间大于通信执行时间,此时计算为算子的瓶颈(bound);反之,当K轴较小时,计算量小,计算执行时间小于通信执行时间,此时通信为算子的瓶颈。在制定数据切分策略前,对原始矩阵分别执行计算和通信任务,根据两个任务的执行时间判定bound场景。

该算子的数据切分策略应满足如下要求:

  • 只切分M轴。因为通信任务调用的Hccl API要求分块数据内存连续,若按N轴切分,则每行数据都被切断,导致通信数据的内存不连续,不满足通信要求;若按M轴切分,则每行数据都是内存连续的,满足通信要求。
  • 若A表示长块、B表示短块,只能切出A或B连续排布的形式,例如AAAB、BAAA等情况。

如上文所述,数据切分的目标是达成尽可能多的流水掩盖。根据计算与通信任务的执行时间差异,实际场景可以分解为如下两个具体场景,两个场景有各自细分的切分目标。

1、计算bound:

对于同一个切分后的数据块,计算执行耗时大于通信执行耗时,此时,计算连续,且通信的尾块要短,如下图所示。

6.png

2、通信bound:

对于同一个切分后的数据块,通信执行耗时大于计算执行耗时,此时,通信连续,且计算的头快要短,如下图所示。

7.png

在进行最终的数据切分前,需要做的前置工作有:判定bound场景、分别对Matmul计算和AllReduce通信的数据量与执行时间关系做公式拟合。具体步骤如下。

1、将输入数据分别单独执行Matmul计算和AllReduce通信任务,利用msProf工具分别采集执行时间,判定时间较长的任务为对应的bound场景。例如通信执行时间大于计算执行时间,则为通信bound场景。

2、将输入数据按M轴切分,分别切成M轴为256、512、768、1024、2048的若干数据块,这里可以根据实际情况调整切块大小。

3、将步骤2得到的数据块分别做AllReduce通信,利用msProf工具采集执行时间,得到每个数据块的执行时间t1, t2, ..., tn。然后,作图分析数据量与对应的执行时间的关系,拟合得到公式t = CostComm(m),其中m表示数据块的M轴长度,t表示数据块的通信执行时间,CostComm表示拟合得到的m和t的映射关系。该映射关系一般为线性,若不满足线性关系,可以采用分段拟合的方式。示例如下:

         8.png

9.png

数据量x = m * N * sizeof(dataType),单位是Bytes。该拟合公式表示为:

  • x小于8MB:t = -0.9698202 * x * x + 27.0622573 * x + 14.769,单位是us。
  • x大于等于8MB:t = 13.58491263 * x + 61.508333,单位是us。

4、将步骤2切分得到的各数据块做Matmul计算,按与步骤3相同的方式采集各数据块的计算执行时间t,并拟合得到M轴长度和计算执行时间关系的公式t = CostMM(m),CostMM示拟合得到的m和t的映射关系。

切分算法步骤:

1、根据输入矩阵的shape:M、K、N,按照经验值设置合适的M轴方向切分的短块长度。如下表达式中,a、b、c表示根据经验值给出的短块长度的备选,将K、N带入如下三个表达式后,取a、b、c的最小值m0为选定的短块长度。

  • a * K * N >= 4 * 1024 * 1024 * 1024,a取不等式的最小值
  • b * K * N / 1024 + b * N >= 6 * 1024 * 1024,b取不等式的最小值
  • c >= 3 * 128,c取不等式的最小值
  • m0 = min(a, b, c)

2、根据短块长度m0和拟合公式,分别得到计算执行时间t0 = CostMM(m0) * 1.15,通信执行时间t1 = CostComm(m0) * 1.15。

注意:通信和计算并行执行时,可能出现抢占内存带宽的情况,导致执行时间增加,一般按经验在拟合公式中乘以1.15的系数,用户可以根据实测情况调整该系数。

3、根据短块的长度,按计算bound或通信bound配平通算,得到长块的长度,长块长度尽量对齐128个元素,以保证计算亲和。本案例为通信bound场景,这里的配平,即将短块的通信时间和长块的计算时间匹配相等:将t1作为长块的计算执行时间,带入t1 = CostMM(m1)公式,计算得到m1,即为长块的长度。

4、根据短块长度m0、长块长度m1、原始M轴长度M,得出长块的切块个数count = (M - m0) / m1。该式一般不能整除,此时需要做如下处理:

  • 将结果的小数部分舍弃,保留整数部分作为切块个数。
  • 由于舍弃了小数部分,M轴长度有剩余,因此需要调整长块长度m1 = (M - m0) / count。
  • 为保持计算亲和,将长块长度m1调整至128对齐,即向下取128倍数的整数,更新长块长度m1。
  • 由于调整m1后,M轴长度有剩余,因此调整短块长度m0 = M - (m1 * count)。
  • 最终得到短块长度m0,长块长度m1,长块个数count。

验证优化方案性能收益

制定切分策略并验证性能收益

本MatmulAllreduce案例中,给定的输入矩阵Shape为M=4096,K=3072,N=8192,数据类型为half,分核数为8,通过融合前msProf工具采集,得到该输入的Matmul计算执行时间为803us,AllReduce通信的执行时间为1071us,总耗时1874us,属于通信bound场景。按照上述的切分算法,本案例的具体切分情况如下:

1、根据经验值选定短块(bound场景中短块为头块,即切分后的第一个数据块)M方向长度m0为384,则通信数据量x为:384 * 8192 * 2 / 1024 / 1024 = 6MB。按通信拟合公式估算通信的执行时间为143us。考虑可能会发生内存带宽冲突,因此再乘以15的系数,得出通信执行时间164us。

2、计算长块M方向长度。根据短块通信时间,配平长块的计算执行时间同样为164us,按计算拟合公式估算出该长度m1为768。

3、根据M=4096、m0=384、m1=768,计算长块个数:(4096 - 384) / 768 = 4.83,向下取整为4。

4、根据短块长度m0=384,长块个数4,调整长块m1长度:(4096 - 384) / 4 = 928,向下按128对齐,调整m1为896。

5、根据长块长度m1=896,长块个数4,调整短块m0长度:4096 - 896 * 4 = 512。

6、最终得到将原始输入矩阵切分为5个数据块,长度分别为:{512,896,896,896,896}。

如下代码所示,在算子的Tiling代码中设置制定好的切分策略。按该切分策略测试,融合后该算子的执行时间为1262us,则融合算子的性能收益为(1874 - 1262) / 1874 = 32.7%。

MatmulAllReduceCustomTilingData *tiling = context->GetTilingData<MatmulAllReduceCustomTilingData>(); 
tiling->param.rankDim = 8; 
tiling->param.tileM = 512; // 短块大小 
tiling->param.tileNum = 1; // 短块个数 
tiling->param.tailM = 896; // 长块大小 
tiling->param.tailNum = 4; // 长块个数 
tiling->param.rankM = 4096; 
tiling->param.rankN = 8192; 
tiling->param.rankK = 4096; 
tiling->param.isTransposeA = 0; 
tiling->param.isTransposeB = 0; 
tiling->param.cToFloatLen = 0; 
tiling->param.nd2NzWorkLen = true; 
tiling->param.dataType = static_cast<uint8_t>(HCCL_DATA_TYPE_MAP.at(aType));

针对切分膨胀做调整

上文提到,切分数据会引起计算或通信执行时间的膨胀,使实测结果与理论值有偏差。比如切块数量较多时,执行时间的膨胀对性能影响较大,可能导致性能收益变小或者出现性能劣化,因此最终需要根据上述理论切分策略,结合实测,对切分策略做调整。

总结

MC2算子通过数据切分后计算和通信的并行执行,获得性能收益,但受数据切分后执行时间膨胀的影响。对MC2算子进行性能调优的主要方式是制定数据切分策略,开发人员需要根据理论推导找到理想切分策略,然后根据实测结果调整,最终找到最优切分策略。

更多学习资源

了解更多Ascend C算子性能优化手段和实践案例,请访问:https://www.hiascend.com/ascend-c

【声明】本内容来自华为云开发者社区博主,不代表华为云及华为云开发者社区的观点和立场。转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息,否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。