CANNTorchTitan: DeepSeek-V3.2 32k长序列预训练昇腾优化实践【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train随着 DeepSeek 团队发布 DeepSeek-V3.2其创新的DeepSeek Sparse Attention (DSA)架构通过稀疏注意力机制显著提升了长序列计算效率迅速成为业界关注的焦点。为把握这一技术趋势众多厂商迫切希望快速完成 DSA 的适配与复现并将其融合到自有模型中探索性能收益这也对昇腾 NPU 上训练框架的快速适配提出了挑战。本文主要介绍基于昇腾 (Ascend) NPU 集群利用 PyTorch 官方原生大模型训练框架TorchTitan实现 DeepSeek-V3.2 全链路预训练适配的技术实践。我们不仅打通了复杂的混合并行策略更在业界率先攻克了 DSA 算子的反向传播难题结合 torch.compile 编译加速实现了从模型定义到算子编译的原生对齐。CANN团队复现的DeepSeek-V3.2长序列CPT训练在A3 64卡集群极致显存场景下性能达成29 TPS/卡。Highlights训练框架使用TorchTitan TorchTitan-npu插件化方案并行策略使能原生TP/PP/EP/FSDP2昇腾适配同时针对稀疏DSA结构提出Custom CP在A3 64卡完成32K长序列CPT训练涉及的训练框架代码已完全开源基于Ascend C实现NPULightning Indexer和Sparse Flash Attention正反向融合Kernel发挥稀疏计算潜力Ascend C Kernel技术文档和代码已开源基于FSDP2Fully Sharded Data Parallel 2的原生内存优势配合A3超节点的高通信带宽的硬件特性权重预取的通信能被计算完全掩盖同时凭借A3高速H2D带宽使能Swap Optimizer至 Host 内存实现了671B长序列全量微调的内存最小化部署方案TorchTitan-npu 无缝对接torch.compile Inductor图优化能力实现FX计算图pass优化在DeepSeek-V3.2预训练场景拿到性能收益是未来训练场景PyTorch编译优化的主战场Outline为什么选择TorchTitan分布式并行策略内存优化DSA融合算子FSDP2Torch.compileinductor加速BenchmarkFuture Plan为什么选择TorchTitan原生分布式编程模型TorchTitan 最根本的优势在于深度应用 DTensor将张量并行抽象为一等公民语义实现并行策略的乐高式组合。统一抽象DTensor 通过 PlacementShard/Replicate/Partial将参数切分与模型定义解耦开发者仅需配置子模块的 ParallelStyle 声明切分方式无需显式编写通信逻辑。多维组合基于携带完整设备网格的 DTensorTorchTitan 可以极低侵入成本实现 5D 并行FSDP2TPPPEPCP组合在支持丰富特性的同时保持代码库紧凑、模块化。无缝接入 PyTorch 生态TorchTitan 是 PyTorch 社区原生的大模型训练框架接入 TorchTitan有助于加速昇腾生态与 PyTorch 社区深度融合复用开源方案有效提升 NPU 训练场景的易用性。torch.compile 深度集成TorchTitan 原生支持 torch.compile借助 Inductor 编译器实现训练优化FX图变换通过 FX 捕获静态图执行通算调度、模式匹配等图级 Pass自动削减冗余操作生成精简中间表示。算子融合依赖 Triton 实现水平融合合并连续逐元素操作链与垂直融合将计算密集型算子与偏置、激活等合并降低内存访问与内核启动开销。后续 CANN 平台可通过 Ascend C 自动融合对接 Inductor进一步提升性能。PyTorch社区原生工具集成TorchTitan 将 PyTorch 生态中的工具系统集成提供开箱即用的稳定性与调试能力分布式检查点原生集成 DCP支持异步保存CheckpointManager 以低开销保存模型分片、优化器状态及加载器偏移量恢复时自动定位最新检查点并校验完整性。容错和维测TorchFT 可实现副本级检查点管理在 HSDP 场景下节点故障时可仅恢复故障节点显著提升平均无故障时间基于 FlightRecorder 可缓存通信类型、数据量、调用栈等信息辅助定位死锁与训练超时根因。低精度训练插件化集成 TorchAO复用其低精度数据类型抽象与量化 recipe社区已在 TorchTitan 中实现对 FP8 / MxFP8 数据类型的支持完成低精度 MoE 网络训练验证。TorchTitan-npu插件化支持方案TorchTitan 提供了ModelConverter机制作用在模型定义后、并行策略TP/FSDP 等应用之前可以在不修改原始模型定义的前提下通过注册机制对模型模块进行替换或重写详细介绍请参考官方文档 基于这套机制我们设计了NPU插件化适配方案预期目标是不修改上游模型代码方便跟踪上有更新支持优化项通过开关配置支持多种优化项自由组合当前我们主要使能融合算子优化量化、optimizer优化等特性在逐步扩展中这套机制也方便结合torch.compile等能力可以实现进一步的优化。使用示例[model] # 在model配置下指定需要使能的融合算子即可 converters [npu_dsa, npu_rms_norm, npu_permute, npu_gmm]分布式并行策略MoE的TP策略优化大EP场景下路由专家不做TP切分TP extend EPTorchTitan原生的MoE TP策略如下图所示以上策略在两个分支上都存在不合理之处路由专家分支接收在TP域AllGather之后的激活tensor一个TP域内每个rank对相同的tensor做了topk选择产生了计算和内存的冗余同时Reorder重排布模块又做了Sequence Parallel相当于前面TP域AllGather对路由专家分支来说是冗余操作。共享专家分支与传统的TP不同其输出没有立刻做ReduceScatter通信而是维持Partial这种不安全的状态并且占据内存较多。以上不合理之处与TP域的AllGather和ReduceScatter位置安排不合理有关可将这两个通信移动至共享专家分支而路由专家分支不需要TP域的通信。调整后MoE的TP策略如下图所示路由专家支路不再有冗余的计算和通信共享专家支路为普通的TP策略自定义CP策略torch原生CP介绍长序列场景下通过Context ParallelCP对输入数据的Sequence维度尤为重要可大大降低激活内存的占用。TorchTitan使用了torch原生提供的context parallel能力其原理如下图所示使用一个上下文环境包裹模型训练step上下文内部对模型输入做sequence轴切分同时将Attention算子替换为RingAttention或AllGatherKV实现以确保Attention计算正确。针对DSA的自定义CPDeepSeek-V3.2模型与传统稠密注意力模型不同其内部的Lightning Indexer和DSA稀疏注意力在CP场景下都需要对Key和Value进行处理以确保计算的正确性。由于DeepSeek-V3.2内部注意力结构更加复杂且Attention算子不再是常用的SDPAtorch原生的CP能力并不适用。在此背景下我们在遵循torch原生CP设计逻辑的基础上为DeepSeek-V3.2设计实现了自定义CP其原理如下图所示CP场景下需确保q_indexer与query分别和完整序列的k_indexer、key value进行过注意力计算一般有两种方式达成此目的RingAttention形式或直接对KV做序列维度AllGather。RingAttention形式实现复杂且对Indexer与稀疏注意力计算都需要实现综合来看开发复杂且不够灵活。对KV做序列维度AllGather实现简单灵活。此外考虑到weight absorb模式下KV的内存占用较小、Indexer中key也内存占用很小AllGather带来的内存代价完全可以接受。综合以上分析我们实现了CP版本的DSA forward函数仅仅是在原版代码上增加对KV的AllGather通信。以使用NPU融合算子的DSA实现为例CP场景下的计算逻辑如下图所示值得一提的是我们首先实现了 CustomContextParallelContext 基类其中只进行模型输入的序列切分用户可在子类中自定义模型前向计算的CP patch逻辑具有较好的可扩展性。DTensor计算的Sharding策略优化算子Sharding策略当使用DTensor进行计算算子下发时需要对算子的输入包含权重、输出生成sharding策略。对于一个DTensor若选定的sharding策略与当前状态不符就会触发DTensor的状态转变会引入Slice或通信等额外的操作会对性能甚至内存产生非常不利的影响。因此对算子生成合理的sharding策略是至关重要的本小节将介绍对两个NPU算子的sharding策略优化。Matmul Sharding策略优化TP场景下当前默认的matmul sharding策略生成并不合理。以q_lora矩阵乘为例该矩阵乘在TP域内不做任何切分因此直接计算matmul即可如下图右侧所示。但目前sharding策略生成并不合理如下图左侧所示生成的策略希望将简单的Replicate矩阵乘转为权重做TP行切的形式引入了很多额外操作大大影响性能输入X从Replicate变为Shard(2)需要Slice操作。权重W从Replicate变为Shard(0)需要Slice操作。输出Y由于TP策略中指定q_lora模块不做并行因此最终该matmul输出需要调整为Replicate状态从Partial到Replicate则需要AllReduce通信。实际上pytorch已经编写了较为完善的matmul sharding策略生成函数可以利用该函数获取更合理的sharding策略。针对以上示例可通过如下代码获得预期的合理策略Matmul Backward Sharding策略优化matmul_backward算子当前也有sharding策略不合理的问题。以模型末尾的output矩阵乘为例其正向是权重TP列切的形式反向时输入的切分状态和生成的默认切分策略如下图左侧所示策略的不合理引入了冗余操作输入X从Replicate变为Shard(2)需要Slice操作。梯度dY从Shard(2)变为Replicate需要AllGather通信且占据更多内存。权重W从Shard(1)变为Shard(0)需要AllToAll通信。虽然torch没有matmul反向的sharding策略生成函数但考虑到matmul反向实际上就是计算dX和dW的两个矩阵乘我们设计对两个矩阵乘分别调用 _mm_like_strategy 生成策略然后融合得到最终的matmul反向sharding策略。此外由于_mm_like_strategy 对一个matmul会生成多条策略还需要考虑如何选择更合理的策略。由于通常计算中不会再调整权重的切分状态我们从众多策略中分别选择不调整W切分状态的策略dX计算和dW与W切分状态一致的策略dW计算将两条策略融合后即可得到上图右侧的理想策略。可以看到理想的策略实际上就是对算子的输入tensor不再做切分状态的调整这样也就没有了额外操作的引入。内存优化最小化切分策略本次部署面临的核心挑战是在64卡A3小规模集群上部署参数量高达671B的DeepSeek-V3.2模型训练任务。考虑到TorchTitan当前使用FP32类型存储模型权重DeepSeek-V3.2权重及梯度即占据5368GB超过了32卡A3的总内存大小4096GB因而64卡已经是模型部署可能的最小规模。在此约束下部署策略需要围绕‘极限压缩单卡内存占用’这一主线展开。FSDP2流程与特点FSDP2作为PyTorch分布式训练的核心演进方案其核心流程是在每个计算步骤前向/反向中按需全量收集AllGather被切分到各设备的模型参数计算完成后立即释放非本设备参数从而降低模型参数及梯度相关的驻留内存。相较于FSDP1FSDP2在分片粒度、通信重叠和内存复用方面进行了深度优化支持参数、梯度与优化器状态的统一分片管理采用逐层layer-wise收集与释放机制实现通信与计算的相互掩盖的高效流水同时原生适配混合精度训练与TorchTitan框架的自动化切分配置。零冗余权重在 MoE模型常见的大EP小TP模式部署下尽管模型内占据主要参数量的MoE部分通过EP与流水线并行PP实现完全切分各专家权重被无冗余地散布Attention/共享专家部分的权重往往是在卡间进行冗余复制的。在此类混合并行策略导致局部参数冗余的场景下FSDP2的优势尤为突出通过将冗余部署的Attention权重与共享专家参数以DP域为单位进行二次分片FSDP2能够在保持计算逻辑不变的前提下显著削减跨DP域的参数副本数量为更大批量或更长序列训练释放宝贵资源。具体而言考虑PP2TP4DP16的场景单卡非MoE权重和梯度的总占用量约为(671B - 654B)/PP2/TP4 * 8Byte ~18GB考虑到其中部分权重不参与TP切分如qkv_a与indexer实际值将会进一步增加。DP16场景开启FSDP后权重和梯度在各DP域间做完全均分几乎可以完全消除相关的内存占用。最终我们对比了如下几种切分策略在考虑MoE极端负载不均衡的可能后选择了PP2TP4EP64为基础的切分方式。Swap Optimizer大模型训练中常用的AdamW优化器中需存储两个FP32的动量这会带来巨大的内存代价训练的前反向计算期间优化器不会被使用因此优化器会浪费巨量的device内存。针对这一问题MindSpeed开源训练框架设计了SwapOptimizer特性其优化内存的方式如下为节省device内存在前反向期间将优化器卸载至host内存。为了降低优化器更新时的内存峰值并尽可能提升性能以param为粒度将“优化器加载——优化器更新——优化器卸载”设计为切片流水并行的形式。本样例在TorchTitan框架下参考并实现了SwapOptimizer特性大大节省了训练时的device内存空间。与MindSpeed有所不同的是TorchTitan训练中优化器只管理两个FP32动量并且需要处理权重为DTensor的场景。DSA融合算子LI/SFALightningIndexer/SparseFlashAttention相关算子介绍可参考NPU DeepSeek-V3.2 Ascend C 融合算子优化SparseLightningIndexerGradKLLossSparseLightningIndexerGradKLLoss基于一系列操作得到Sparse Training Stage每一个 token 对应的 Top-$k$ 个位置的反向梯度计算并融合了Loss的计算。反向梯度计算的过程依赖于Sparse Flash Attention计算过程softmax的聚合与归一化结果与Lightning Indexer自身的输出分布考虑到显存与算力的平衡SparseLightningIndexerGradKLLoss算子在保留了一些重要正向计算结果的前提下进行了部分重计算。Lightning Indexer的正向输入Index Query $\tilde{Q}\in\R^{g\times d}$、Index Key $\tilde{K}\in\R^{S_{k}\times d}$Sparse Flash Attention的正向输入Query $Q\in\R^{g\times d}$、Key $K\in\R^{S_{k}\times d}$,$W\in\R^{g\times 1}$其中 $g$ 为 GQA 对应的 group size$d$ 为每一个头的维度$S_{k}$ 是上下文的长度SparseLightningIndexerGradKLLoss算子的具体计算流程如下Lightning Indexer正向 Top-$k$ value的计算公式$${t,:}W{t,:}ReLU(\tilde{q}_{t,:}\tilde{K}_{topk,:}^\top) $$其中$W_{t,:}$是$W$矩阵中第$t$个token对应的$weights$$\tilde{q}_{t,:}$是$\tilde{Q}$矩阵第$t$个token对应的$G$个query头合轴后的结果$\tilde{K}{topk,:}$为根据$\tilde{q}{t,:}$的Top-$k$索引从$\tilde{K}$矩阵中取出的token组成的矩阵。Sparse Flash Attention正向 Top-$k$ value的计算公式$${t,:} \text{Softmax}(q{t,:} K_{topk,:}^\top/\sqrt{d}) $$其中$p_{t,:}$是第$t$个token对应的Softmax结果$q_{t,:}$是$Q$矩阵第$t$个token对应的$G$个query头合轴后的结果${K}{topk,:}$为根据${q}{t,:}$的Top-$k$索引从$K$矩阵中取出的token组成的矩阵。在训练阶段需要将Lightning Indexer的输出与Sparse Flash Attention的softmax结果进行对齐计算loss对应的loss function为$$ oss{}\sum_tD_{KL}(p_{t,:}||Softmax(I_{t,:})) $$其中$p_{t,:}$作为target distribution需要在之前的计算结果上进行所有的head的求和然后把求和结果沿着上下文方向进行L1正则化得到。$D_{KL}$为KL散度其表达式为$$ D_{KL}(a||b){}\sum_ia_i\mathrm{log}{\left(\frac{a_i}{b_i}\right)} $$通过求导可得Loss的梯度表达式$$ I\mathop{{}}\nolimits_{{t,:}}Softmax \left( I\mathop{{}}\nolimits_{{t,:}} \left) -p\mathop{{}}\nolimits_{{t,:}}\right. \right. $$利用链式法则可以进行weightsquery和key矩阵的梯度计算$$ W\mathop{{}}\nolimits_{{t,:}}dI\mathop{{}}\nolimits_{{t,:}}\text{} \left( ReLU \left( S\mathop{{}}\nolimits_{{t,:}} \left) \left) \mathop{{}}\nolimits^{\top}\right. \right. \right. \right. $$$$ \mathop{{\tilde{q}}}\nolimits_{{t,:}}dS\mathop{{}}\nolimits_{{t,:}}\tilde{K}_{topk,:} $$$$ \tilde{K}{topk,:}\left(dS\mathop{{}}\nolimits{{t,:}} \left) \mathop{{}}\nolimits^{\top}\tilde{q}\mathop{{}}\nolimits_{{:t, :}}\right. \right. $$其中$S$为$\tilde{Q}$和$\tilde{K}$矩阵乘的结果$S_{t,:}$是$S$矩阵中第$t$行。$dW_{t,:}$是$dW$矩阵中第$t$行$dq_{t,:}$是$d\tilde{Q}$矩阵中第$t$行$d\tilde{K}_{topk,:}$是$d\tilde{K}$矩阵中Top-$k$索引对应行的累加值。实际开发中需要充分利用昇腾硬件的特征以实现更好的性能以下将详细介绍当前的解决方案及具体实现方式。SLIG算子Tiling设计算子核间 Tiling 设计如下SparseLightningIndexerGradKLLoss对B和S轴进行分核切分处理算子核内 Tiling 设计如下SparseLightningIndexerGradKLLoss一次计算的基本块大小为$(128,128)$。L1空间划分如下部分$Q$ 矩阵$128\times 576\times2\text{ Bytes}144\text{ KB}$即$Q$ 矩阵矩阵常驻于L1$Qindex$ 矩阵 使能Double Buffer$2\times64\times 128\times2\text{ Bytes}32\text{ KB}$即$Qindex$ 矩阵常驻于L1$Gather$ 矩阵 使能Double Buffer$2\times128\times 128\times2\text{ Bytes}64\text{ KB}$其中$Gather$矩阵每次的搬运大小为$(128,128)$$reluGrad$ 矩阵$64\times 2048\times2\text{ Bytes}288\text{ KB}$其中$reluGrad$矩阵每次的搬运大小为$(64,2048)$。L0AL0BL0C使能DoubleBuffer分别划分为$2\times32\text{ KB},2\times32\text{ KB},2\times64\text{ KB}$在 $C_1$ 阶段$MatmulP$一次搬入 L0A 和 L0B的矩阵块大小分别为 $(128,128) / (128,64),(128,128)$。在 $C_1$ 阶段$MatmulSy$一次搬入 L0A 和 L0B的矩阵块大小分别为 $(64,128),(128,128)$。在 $C_2$ 阶段$MatmulDk$一次搬入 L0A 和 L0B的矩阵块大小分别为 $(128,64),(64,128)$。在 $C_2$ 阶段$MatmulDq$一次搬入 L0A 和 L0B的矩阵块大小分别为 $(64,128),(128,128)$。SLIG Pipeline 设计SparseLightningIndexerGradKLLoss的计算流程有五个阶段即预处理Gather($V_0$)矩阵乘和ReLu操作($C_1$)以及后续$V_1$侧多个向量计算$C_1$侧矩阵乘和后处理ScatterAdd操作($V_2$)。相较于LightningIndexer和SparseFlashAttention算子而言该算子计算流程更为繁琐数据依赖更加复杂朴素的流水排布导致Cube和Vector计算过程中出现大量气泡无法充分发挥算力优势。为了防止多个阶段的串行执行需要提前下发 (Preload) 一次 $V_0$ 以实现流水的互相掩盖。具体流水效果图如下图所示离散访存优化SparseLightningIndexerGradKLLoss算子Gather($V_0$)和ScatterAdd($V_2$)阶段都需要采用离散访存。对于非 Sparse 场景的 Attention 算子数据访存一般是连续的因此单次数据搬运耗时相较于指令下发耗时更长可实现连续的多次访存达到较高的访存带宽。然而对于 Sparse Attention 的离散访存而言单次访存的耗时大幅下降甚至可能小于指令下发的耗时这就会导致严重的指令下发阻塞如下图所示。基于 A3 芯片 Cube 核和 Vector 核数比为1:2的特性,利用 Vector 核发射访存指令 将指令的下发效率提高一倍。另一方面单次访存的数据量本身也会影响实际达到的带宽 而离散访存单次的访存量较小这会导致数据搬运本身的耗时高于预期。 针对这一问题,可以将离散的数据点两两聚合,利用srcGap参数实现成对的数据访存,提升访存带宽,如图所示。合轴优化SparseLightningIndexerGradKLLoss算子由于不同QueryToken对应的KeyToken不同因此把N轴放到内层循环轴将B和S轴作为切分轴。SparseFlashAttentionGradSparseFlashAttention算子整体沿用FlashAttentionGrad的计算流程针对离散访存进行了指令缩减及搬运聚合的细致优化反向计算整体分为以下几个核心步骤整体计算逻辑如下根据Top-$k$索引从$K$矩阵中取出对应的的token组成$K_{topk}$矩阵。$$ _{topk} \text{Gather}(K, \text{Top-K index}) $$重计算Sparse Flash Attention正向softmax结果$$ \text{Softmax}(Q K_{topk}^\top/\sqrt{d}) $$通过 $dY$ 与转置后的 $V$矩阵进行矩阵乘法得到$dP$$$ dP dY V^T $$将$dP$ 与FlashSoftmaxGrad计算结果做减法再与$P$ 相乘得到 $dS$$$ dS P *(dP - \text{FlashSoftmaxGrad}(dY attention_in)) $$其中FlashSoftmaxGrad的输入为反向梯度dY和正向输出attention_in。利用链式法则可以进行query、key和value矩阵的梯度计算$$ dQ dS K_{topk} /\sqrt{d} $$$$ dK \text{ScatterAdd}(dS^\top Q /\sqrt{d},\text{Top-K index} ) $$$$ dV P^\top dY $$SFAG 算子 Tiling 设计算子的 Tiling 设计决定了储存空间是否被充分利用进而影响整体的算力利用率以及算子性能。针对SparseFlashAttentionScoreGrad本章分别阐述核间和核内的 Tiling 设计。核间 Tiling 设计采用了如下方案当前SparseFlashAttentionScoreGrad算子将Batch轴与 $S_{Q}$ 合轴后作为分核轴核数随分核轴连续排布。T B * $S_{Q}$CoreProcessNum CeilDiv(T, CoreNum)核内 Tiling 设计采用了如下方案核内空间划分左矩阵L1基本块右矩阵L1基本块$Q_{index}$ $K_{index}$128 * 576128 * 128 * 2(DB)$d_{y}$ V128 * 512128 * 128 * 2(DB)$d_{s}$ K128 * 128 * 2(DB)128 * 128 * 2(DB)$d_{s}$ Q128 * 128 * 2(DB)128 * 576$d_{y}$ P128 * 128 * 2(DB)128 * 512$Q_{index}$以基本块$128\times 256$的粒度载入到L1中并在单基本块完成计算流的过程中常驻L1以在$d_{s}$ $Q_{index}$计算中复用减少重复访存量同样的技巧也应用在$V_{index}$。如表中所示其他矩阵乘计算中受L1空间限制源数据搬运启用DoubleBuffer,实现搬运与计算并行。同时UB上的基本块采用$128\times 128$。SFAG Pipeline 设计SparseFlashAttentionScoreGrad的计算流程包括五个阶段$K_{index}$和$V_{index}$离散转连续的数据搬运V0$Q_{index}$$K_{index}$和$d_{y}$VC1计算P和$d_{s}$V1计算$d_{s}$K最后是$d_{k}$和$d_{V}$离散写出核外V2。严格按照计算图的基础流水排布如下图上半部分所示多个阶段串行执行会导致计算过程中出现大量的气泡。这里优化为在起步阶段连续发起两次V0和C1接触后续计算的数据依赖以实现流水的掩盖。效果如下图下半部分所示可以看到CV间并行执行大大减少了CV数据依赖损耗。离散访存优化相较于传统FAGSFAG增加了大量的离散访存且单次访存的数据量较小这样会导致访存指令本身的开销较大带宽也无法充分发挥针对此类问题可以通过Vector核对离散数据做非连续转连续排布以$128\times headDim$为基本块放置到workspace临时空间而Cube侧发起计算时可以连续搬运输入数据显著提高带宽利用率如下图所示。FSDP2Torch.compileinductor加速A3超节点下的FSDP2预取掩盖在本实践中随着模型参数量和序列长度的激增通信开销往往成为制约集群扩展效率的瓶颈。在 TorchTitan CANN 的原生适配中我们深度利用了FSDP2的核心特性——权重预取Weight Prefetching并结合 A3 超节点的高带宽优势实现了计算与通信的极致并发。1FSDP2 权重预取确定性的通信流水线FSDP2 基于 DTensor 构建相较于传统 FSDP 提供了更细粒度的通信控制能力。我们在 NPU 上实现了精确的层间流水线调度当计算流Compute Stream执行第 N 层 Transformer Block 的前向或反向计算时通信流Communication Stream已通过异步调度发起第 N1 层的 AllGather 操作提前聚合下一层所需的全量参数与此同时第 N 层计算完成后其参数会立即被重新切分reshard回各卡释放显存。这种设计带来两方面收益一是显著降低了模型参数的峰值显存占用二是通过预取使计算单元无需等待参数就位实现了通信与计算的真正并发Compute-Communication Overlap。但需要注意的是通算掩盖的前提是通信耗时短于计算耗时。仅靠软件层面的预取调度并不充分——若通信延迟远大于计算时间流水线仍会产生阻塞。在这一点上A3 超节点的高带宽互联架构发挥了关键作用其高吞吐、低延迟的通信能力有效压缩了 AllGather 的耗时为通算掩盖提供了硬件基础。2软硬协同下的通算全掩盖在上图实际 Profiling 中可以观察到权重预取的All_gather通信几乎被Device侧的计算完全覆盖DeepSeek-V3.2 的训练流水线呈现出极高的紧凑度性能得到有效提升。这一结果源于 FSDP2 预取机制与 A3 高带宽能力的协同配合软件侧FSDP2 通过异步预取将通信任务精确地调度到计算间隙中硬件侧A3 的大带宽互联确保通信任务能在计算窗口内完成使通信延迟完全隐藏于 NPU 的 Device 侧计算之中。Inductor 图优化加速即使不依赖算子融合与 CUDA GraphsInductor 依然通过Codegen Wrapper和静态化分析一定程度上缓解了 Host-bound 问题。它将原本由 Python 解释器逐行驱动的动态调度过程转化为预编译的高效计算图执行流通过在编译期完成算子选型、shape推导和内存规划开启 Inductor 编译能够成功绕过 PyTorch 运行时昂贵的 Dispatcher 分发机制和 Python 交互开销加快 host 侧算子下发速度。Inductor 编译将原本Eager模式下每个算子下发固定开销在编译期间解决运行时host侧只剩下最精简的Kernel Launch调度通过Profiling我们可以看出原本linear层包含aten::t和aten::matmul两个算子需要经过PythonDispatchMode下发具体算子整体的host侧下发耗时为409.5us开启compile后通过可执行python文件直接完成算子下发同样四个操作仅需要118.8us的下发时间。开启inductor后bypass自动融合在 Inductor 编译流程中除了通过 fx_passes 进行图优化外还包含以下几个关键步骤decompositions 将复杂算子分解为基础算子组合lowerings 字典将 FX 节点映射为对应的 IR 节点Scheduler 根据融合收益对算子进行拓扑排序与调度决策Scheduler 输出融合后的 IR 节点以及对应的执行计划最后由 Codegen 阶段将这些信息进一步转换为可执行的 Triton Python 代码。实验中我们发现当前 NPU 上由 Inductor 自动融合生成的 Triton 算子其编译时间和执行性能均表现较差。因此现阶段我们暂时绕过了 Inductor 流程中的 Triton Codegen 部分后续将接入Ascend C Codegen 以更好地支持在 NPU 上使用 Inductor 后端进行完整编译。通过在 compile 之前清空 lowerings 和 decompositions使得大部分Aten算子都不会被分解也不会下沉到Inductor IRfallback到算子原始实现此外由于在GraphLowering中会对没有显式指定内存布局的反向算子添加require_contiguous约束对于不连续的张量插入copyclone操作变为连续内存这一操作会在反向阶段将aten算子转为单独的triton算子实现我们对以上代码进行修改去除对于反向算子的布局操作所有算子均使用原始的布局可以完全屏蔽triton算子。Benchmark精度我们在A3集群上预训练DeepSeek-V3.2模型实践表明语言模型损失lm loss与梯度范数grad norm稳健收敛表明TorchTitan框架结合TorchTitan-npu适配在大规模模型训练中的稳定性与高效性。lm loss从训练初始step1的12.25稳步下降至step37的8.37全程保持持续递减趋势无明显震荡梯度同步从初始28.24平稳回落至16.03整体趋于收敛且无异常波动。与此同时对DeepSeek-V3.2网络所提出的Lighting Indexer也进行了训练。在整个预训练过程中indexer loss稳步下降且全程保持平稳递减。结合lm loss和grad norm变化曲线证明在大规模分布式场景下精度正常能够高效支撑DeepSeek-V3.2模型的预训练任务。性能当前TorchTitan-npu通过融合算子以及inductor的优化能在基于Torch原生小算子实现的基础上实现较为明显的性能提升未来当Inductor后端接入昇腾原生的Ascend C Codegen预期会有更加明显的性能收益。融合算子通过前文提到的轻量级插件化方式可以实现在不修改torchtitian代码的前提下实现对NPU基础融合算子的支持当前支持配置RMSNorm、Rope、GMM、permute/unpermute四个基础融合算子和DeepSeek-V3.2模型DSA融合算子LI/SFA/LIG可以通过在训练配置toml文件中converters简单配置生效。由于原始小算子实现存在较大的内存开销我们在8层裁剪模型1个dense层7个moe层每个moe层8个专家和4层裁剪模型1个dense层3个moe层每个moe层8个专家设置训练长度4k与小算子实现性能进行对比性能及内存情况如下裁剪模型实现方式内存 (Memory)性能 (Performance)8层裁剪模型原始小算子实现58.32G155TPS8层裁剪模型矩阵吸收 DSA融合算子50.54G400TPS4层裁剪模型原始小算子实现54.58G402TPS4层裁剪模型矩阵吸收 DSA融合算子38.71G568TPSDeepSeek-V3.2模型在使能DSA融合算子后不仅可以显著降低训练中的内存占用同时提升计算速度。进一步开启基础融合算子在完整DeepSeek-V3.2模型上测试单个step内计算时间减少5.16%实验采集了融合算子开启前后device侧的执行耗时算子类型融合前耗时融合后耗时时间节省性能提升比例RMSNorm138.7μs39.9μs98.8μs71.2%Permute803.9μs127.3μs676.6μs84.2%GMM3.1ms2.4ms0.7ms22.6%torch.compile我们参考原始TorchTitan代码中apply_compile对于逐个transformer block进行整图编译的方式将完整Attention部分拆分为MLAProloginner_attentionMLAEpilog三个部分由于inner_attention已经使用DSA融合算子实现host侧不需要过多的小算子下发因此是否compile对整体性能影响不大。对于MLAProlog和MLAEpilog两个部分以及MoE部分的router、gate等进行编译缓解小算子较多带来的host侧下发开销。开启compile后对于MLAProlog存在明显host bound的情况有明显优化原始host侧下发时间host侧下发时间device侧计算执行时间开启compile前11.0ms10.7ms开启compile后5.7ms4.5ms在完整DeepSeek-V3.2模型实测开启inductor compile后单层耗时减少5.1%。由于DeepSeek-V3.2模型计算耗时集中于SFA、LI等主要算子开启inductor对于host侧下发优化对整网性能提升效果有限。Future PlanTorchTitan-npu 当前主要是兼容社区的生态将PyTorch原生简单易用的TorchTitan生态迁移到NPU上使得社区可以快速将训练任务迁移到昇腾芯片上部署丰富CANN的训练生态未来将进一步在基于torch.compile的开箱性能上持续优化当前SFAG算子将在下一个CANN版本进行升级大幅提升性能此外我们还将在以下方向持续演进Ascend C Autofuse对接PyTorch框架在图模式下的Inductor组件为此提供了基础的算子融合与编译能力其设计兼容现有PyTorch生态具备良好的软件继承性。然而由于NPU与GPU在内存架构、计算单元及并行机制等方面存在本质差异直接沿用面向GPU的优化策略难以充分发挥NPU的硬件潜力因此现阶段通过 Bypass Triton 方式规避 Triton 后端在 NPU 上的性能不足。为在继承Inductor整体流程与接口的基础上实现针对NPU的高效融合需具备面向NPU的即时编译JIT算子生成能力。为此CANN框架继承的AutoFuse组件针对NPU架构特征进行了深度优化实现了从计算图到高效NPU算子代码的自动化映射与生成。该组件不仅保持与Inductor技术的协同与兼容更进一步通过架构感知的融合策略、内存访问优化与指令调度等手段显著提升了算子融合在NPU上的执行效率从而支撑复杂模型在NPU平台上的高性能部署因此将进一步提升inductor编译后的模型性能。DualPipeV支持DualPipeV是一种专为大规模 MoE 模型如 DeepSeek-V3设计的高效流水线并行Pipeline Parallelism, PP调度策略其脱胎自DeepSeek团队设计的DualPipe并借鉴了InterLeave1F1B。DualPipeV的核心设计理念是一个设备上可以同时承载一个前向计算batch和一个反向计算batch两者的通信和计算可以互相并行以此可以很大程度上用另一个batch的计算掩盖MoE EP并行带来的通信开销。TorchTitan社区已初步支持DualPipeV特性但在实践中发现无法在ZeRO3 DP切分的场景下使能DualPipeV原因是torch原生的dW计算实现约束了权重W不能在DP域维持切分状态。为了同时拿到内存和性能收益后续会对TorchTitan进行适配以实现ZeRO3与DualPipeV同时使能。TP_async支持TP 通过 AllGather 和 ReduceScatter 在层内完成跨卡同步通信延迟直接暴露在关键路径上。Async TP 的核心思想是将通信与计算算子同时进行分解——AllGather 拆分为多个 Send/RecvMatmul 拆分为对应的子 Matmul——使得一个子 Matmul 的计算与下一个 chunk 的数据传输并发执行从而隐藏通信延迟。在 TorchTitan-npu 中Inductor 编译器在 FX 计算图的阶段可以识别这种all_gather mutmal或者matmul reduce-scatter这种pattern并自动完成这一分解与重排序无需修改模型代码。配合 A3 节点内的高带宽卡间互联层内 TP 通信可被有效压缩并隐藏在计算窗口之中。【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考