3 GroupedMatmul NZ技术穿刺

在MoE大模型推理中，GroupedMatmul使能NZ格式是比较有效的优化点。在该模型的shape下，穿刺前向每层的GroupedMatmul耗时32ms，如果使能NZ，耗时会优化至30ms，整网预计会有800ms收益。

但是在训练的过程中，如果模型的weight是NZ，但是计算出来的梯度是ND，会涉及到非常多的算子要支持NZ，工作量很大。而FSDP2的机制是，allgather将所有卡的模型参数搜集起来，然后调用torch:split_with_sizes_copy将这部分的参数copy_out到内存中，实际参与计算的是copy_out之后的参数，这部分参数在完成计算后会进行unshard释放掉，原来的模型参数还是ND，并不会干扰模型参数更新的逻辑。因此，只需要实现以下两点：

1. copy_out的时候，将weight从ND转成NZ

2. torch_npu.npu_grouped_matmul前反向计算的时候支持NZ（输入输出都是ND）
   

1）GroupedMatmul计算支持NZ

torch_npu.npu_grouped_matmul已经支持NZ，只是没有对应的反向逻辑，GroupedMatmul反向也是GroupedMatmul，因此可以通过以下方式实现反向：

2）torch:split_with_sizes_copy支持NZ

常规小算子拼接方式将GroupedMatmul的权重转换成NZ格式需要将权重数据从HBM搬运至device侧，转换一层的权重需要耗时13ms，会将GMM使能NZ拿到的收益完全抵消甚至劣化所以只能与前一个weight操作的算子相融合，减轻HBM带宽压力。

在现场临时开发了一版sliceNZ算子，支持融合slice及转换NZ两个操作，原理是通过修改slice搬运出数据的偏移量，直接slice出GMM所需的NZ格式的权重，这样对比原本的slice流程没有新增HBM访问压力。

在联调的过程中，发现转NZ之后会导致slice_copy的性能劣化，经过定位发现转NZ之后，性能对地址512对齐非常敏感，非512对齐时劣化30%，参考第二章通信性能优化的解决方案1：修改fully_shard切分的策略，在使能NZ的时候，我们在每个layer的FSDPParamGroup中，只切分linear，这样可以做到参数量512对齐，规避之后slice_copy的劣化问题得到解决，在slice_copy时转NZ相比于只进行slice_copy劣化不到5%。现场穿刺整网收益500ms左右，关键算子的性能拆解如下：

4 工具诊断

针对大集群场景通信快慢卡诊断与性能瓶颈识别，MindStudio全流程工具链配备了专门针对大集群的的性能工具能力，旨在实现高效且精准的瓶颈定位。

在某攻关项目中，512dieA3集群FDSP2训练出现通信快慢卡问题，如上图所示，一次迭代过程中，最初的通信算子耗时为170ms，而最后一个通信算子则延长至406ms，性能劣化了2.3倍。基于这一实战案例，我们将详细介绍所使用的工具及其能力。

4.1 采集篇：

在大规模集群环境中，通常涉及上百张甚至更多显卡，所收集的性能数据量往往达到TB级别。这些数据的解析、存储及转换测试的成本相当高，显著降低了问题定位的效率。

MindStudio性能调优工具新增了db格式的数据导出选项，相比传统的text格式，db格式能够将存储空间占用降低70%，同时提升50%的显示效率和分析效率，显著加快问题定位速度。在使用pytorch Profiler和msprof命令行进行数据采集时，只需简单设置export_type参数即可选择导出db格式。更多详细操作步骤可参考Pytorch Profiler的使用手册（https://www.hiascend.com/document/detail/zh/mindstudio/81RC1/T&ITools/Profiling/atlasprofiling_16_0090.html#ZH-CN_TOPIC_0000002353635602__zh-cn_topic_0000002370275077_section1548285513313）。

以浦江48层32K模型的512-Die集群为例，原先使用text格式时，每张卡的数据量为320MB，总数据量达到160GB。而切换到db格式后，每张卡的数据量降至61MB，总数据量仅为36GB，数据量减少了77%。这不仅大幅节省了存储空间，还显著提高了数据处理和分析的效率。因此，大规模训练场景建议切换至db格式，以获得更高效的问题定位效率。

4.2 分析篇：

msprof-analzye（https://gitcode.com/Ascend/mstt/tree/master/profiler/msprof_analyze/）是针对profiling数据进行性能分析的工具，针对大集群场景储备了高效、好用的分析能力。接下来将列举几个典型的集群分析能力，全量分析能力参考网站介绍：https://gitcode.com/Ascend/mstt/tree/master/profiler/msprof_analyze/#%E5%88%86%E6%9E%90%E7%89%B9%E6%80%A7%E4%BB%8B%E7%BB%8D

4.2.1 slow_rank快慢卡分析：

slow_rank是基于快慢卡全局投票算法进行全局通信算子分析的慢卡分析功能。

slow_rank基于集群全局每个通信域的通信算子信息，分析通信耗时、关联、rank。由于大多数算子需要进行互相同步，执行时间接近，而快慢卡现象会打破该规律。基于异常值检测构筑投票算法，找到集群中的被害者、肇事者，针对全局通信操作进行逐一分析，如下图所示hcom_allReduce_909_233_1的通信算子，7卡增加嫌疑慢卡1票，其他不增加，然后遍历汇总得出嫌疑慢卡的票数。

执行如下命令，调用slow_rank分析能力，-d传入汇总的全量profiling数据。

msprof-analyze -m slow_rank -d ./ --export_type notebook

输出结果保存在-d目录下的cluster_analysis_output\SlowRankAnalysis文件夹的rank_stats.csv文件，也可通过notebook查看。

如下图所示，132Rank的慢卡投票96次，高居首位，明显是根因慢卡。

4.2.2 hccl_sum通信全局分析：

hccl_sum是全局分析通信算子的分析能力，会汇总TOP-N的通信算子耗时，同时给出最大最小的卡的信息。

执行如下命令，调用hccl_sum分析能力，-d传入汇总的全量profiling数据。

msprof-analyze -m hccl_sum -d ./ --export_type notebook

输出结果保存在-d目录下的cluster_analysis_output\HcclSum文件夹的top_op_stats.csv文件，也可通过notebook查看（python -m jupyter notebook命令）。

可以查看到TOP15的最大耗时的通信算子，132rank最慢的是14次，说明132是慢卡。

4.2.3 compute_op_sum计算算子全局分析：

compute_op_sum是全局分析计算类算子的分析能力，会汇总集群下计算算子的耗时与汇总信息。

执行如下命令，调用hccl_sum分析能力，-d传入汇总的全量profiling数据。

msprof-analyze -m compute_op_sum -d ./ --export_type notebook

输出结果保存在-d目录下的cluster_analysis_output\ ComputeOpSum文件夹的rank_stats_by_optype.csv文件，也可通过notebook查看（python -m jupyter notebook命令）。

分析耗时TOP算子的不同rank的分布图来看，FlashAttentionScoreGrad和FlashAttentionScore算子的卡间分布特别不均匀，132卡出现明显的高耗时。而且有且只有132才会有。如此锁定根因，是132卡的FA算子计算耗时长导致。

4.3 可视化篇：

4.3.1 并行策略分析

（1） 并行策略可视化呈现

MindStudio Insight支持根据加速库和并行策略参数展示集群并行排布关系（对应下方的方框）和通信域连接关系（对应下方的连线），前者表示哪些卡共同构成DP/PP/TP等，后者反应多卡卡间通信通道。

点击方框或连线，可以在下方展示其总体性能柱状图，辅助快速找出异常卡，及其异常原因（计算、通信、空闲时间异常）。另一方面，点击上方的连线，右键可以进一步进行通信详细分析（后面章节详细展开）。

（2） 性能数据热力图

MindStudio Insight除了能够展示卡排布关系和通信域连接关系外，还可以基于上述二维卡排布，进一步得加卡级别性能总体数据，如计算时间、通信时间、空闲时间、TP/PP/DP等通信域总时间，形成热力图。可以通过计算时间长，找出计算慢卡，可以通过某个通信域通信耗时短，找出通信异常卡。

当集群较大时，可以支持折叠TP、CP等并行域，此时界面上叠加计算或通信时间的最大或最小值，进一步定位上层CP或DP内是否存在异常的慢并行域。

4.3.2 集群通信分析

上文中基于并行策略的总体性能分析，可以帮助快速找到慢卡（卡级别），但其深入原因无法得知。不管是计算慢还是通信慢，会对通信过程产生影响，可以进一步通过本节中集群通信数据分析，定位表现异常的特定通信算子（相当于模型层级别）。MindStudio Insight支持通过通信矩阵分析和通信耗时分析两种方式，快速定位通信问题。

（1）通信矩阵分析

MindStudio Insight支持按照指定通信域展示卡与卡间的性能指标，包括通信总量、通信耗时、通信方式、等效带宽等。可以通过上述指标的异常，如等效带宽低定位慢链路问题。

（2）通信耗时分析

MindStudio Insight还支持通信过程的分析，将通信算子按照通信域提取出来，形成通信算子缩略图，如下图所示。缩略图下方会展示出，通信域内通信时间最短的3张卡，及每张卡上耗时最短的3个通信算子。在集合通信时，通信耗时端往往意味着其启动通信晚，是被等待的卡，需要重点关注。从通信算子缩略图中，也可以关注指定通信域内某个通信算子在不同卡的表现，对于集合通信算子，通信耗时短的算子最后开始。

当找到具体的通信问题算子，且导入了单卡数据，可以在上图中右键点击“跳转到Timeline”，即可结合Timeline，进一步深入定位上下文，找出根本原因。

另一方面，MindStudio Insight还支持展示通信过程数据包分布情况，此功能可快速看出因大量小包造成的看似通信慢的问题。

5 总结与展望

现场性能优化过程：

走过这一段过程后，某企业也很认可昇腾超节点的软硬件能力：

• FSDP2是非常适合A3超节点的训练框架，当通信发挥出符合规格的性能后，剩下的就是内存和计算的优化。
• Mindspeed有挺多很好的功能，例如async activation offload。期望以后有更多与模型代码解耦的此类特性，可以吸引客户迁移到自己的框架上。
• 经过两年的时间，我们的工具易用性有了较大的提升，512die的快慢卡问题定位1天就能搞定。
• 使能GMM NZ是十分重要的一个穿刺，这是第一次在大模型训练中找到办法使能了GMM的NZ。特战队在客户界面，有的时候可能就缺那么0.5s，才能持平或超越H800/H200，永远不嫌武器多。

参考：
XTuner V1 ：专为“超大模型而生”，新一代训练引擎 XTuner V1 开源！