DeepSeek-V4 昇腾训练优化实践基于 CANN 的 TorchTitan-NPU 与 AutoFuse 极简训练方案【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train导读本文围绕 CANN 平台上 DeepSeek-V4-Flash 模型的训练优化实践展开介绍基于 TorchTitan-NPU 框架、采用纯 FSDP 大 EP极简并行策略、依托 torch.compile Ascend C AutoFuse 自动融合以及面向稀疏注意力结构的定制融合算子将 A3 集群 64 卡 4K 序列 BF16 训练吞吐从 397 提升至 1100 tokens/p/s 的完整技术路线。读者将掌握该方案的软件分层架构、并行切分与内存优化原理、AutoFuse 收益拆解方法以及基于 llm_pretrain/deepseekv4 目录下的部署脚本完成源码准备、权重转换、容器拉起与多机训练启动的实战路径。Highlights极简并行策略训练框架采用 TorchTitan TorchTitan-NPU 插件化方案以超节点亲和的大 EP 纯 FSDP的精简并行切分策略在极低适配成本和通信开销下实现内存占用最优取得易用性与性能的较好均衡。训练入图自动融合TorchTitan-NPU 深度适配 torch.compile 机制使能训练入图技术依托 Inductor AutoFuse基于 Ascend C 的 Codegen 后端实现端到端的 Vector 算子自动融合为整网带来高达 31.8% 的开箱即用性能收益。定制 NPU 融合算子针对稀疏注意力等复杂结构开发 SparseAttnSharedkv、LightningIndexer 等 4 个高效 NPU 融合算子从负载均衡分核计算、内存与计算均衡等维度协同优化充分释放芯片稀疏算力。可复现性能结果基于上述优化点CANN 已基于 TorchTitan 支持 DeepSeek-V4-Flash 模型训练A3 集群 BF16 精度 64 卡 4K 序列 MTP1 训练吞吐达 1100 tokens/p/s。引言为什么选择 TorchTitan-NPU AutoFuseTorchTitan 为业界带来了 PyTorch Native 大模型训练方案的全新选择其核心设计理念是模型算法与分布式并行、算子优化天然解耦。通过支持 torch.compile 训练入图优化告别手写融合 kernel大大提升了大规模分布式训练的易用性。在并行策略方面TorchTitan 深度集成 FSDP2实现参数、梯度和优化器状态的跨设备分片显著降低单卡内存占用。本实践的整体软件架构自顶向下由 TorchTitan、TorchTitan-NPU 插件、TorchInductor 及 CANN 层的 AutoFuse 组件构成各层分工明确、优势互补TorchTitan作为 PyTorch 原生分布式训练框架避免对 Megatron 等重型框架的依赖以精简代码库提供 FSDP 等并行能力的开箱即用支持。TorchTitan-NPU以插件形式提供可插拔安装针对昇腾集群实现通信与算子的硬件加速适配保持与上游 TorchTitan 的轻量集成。TorchInductor直接继承社区 torch.compile 的图捕获与融合优化能力无需重复造轮即可复用 PyTorch 编译生态的相关基础设施。AutoFuse面向昇腾 NPU定位为基于 Ascend C 的 Codegen Schedule 后端在 Inductor 基础上注入架构感知的融合、优化与展开策略将计算图优化转化为 NPU 高效指令序列。基于上述分层架构实践在 TorchTitan-NPU 加速插件的基础上完成了 DeepSeek-V4-Flash 模型续训练流程的适配与调优。整个方案围绕极简并行、自动融合、定制算子三条主线展开下文分别就基于 FSDPEP 的极简分布式并行、torch.compile 与 AutoFuse 自动融合协同、以及高性能融合算子优化加速进行详细阐述。极简分布式并行优化FSDPEP 极简切分优势在 MoE 大模型训练中EP 切分可有效分摊路由专家权重的存储压力但注意力模块、共享专家、词嵌入及词表投影层等非专家参数仍需额外切分机制加以应对。另一方面随着模型结构迎来又一个快速发展阶段传统 TP/PP 方案在性能开销与易用性层面均面临诸多挑战而 FSDP 在上述场景中展现出较为显著的优势。因此本实践基于 A3 超节点 64 卡环境在 DeepSeek-V4-Flash 模型训练中创新性地选择纯 FSDP EP 极简切分方案具体配置如下TPPPVPPEPFSDPMBSGBS集群1111281281102464卡即完全不启用张量并行TP1与流水线并行PP1EP 与 FSDP 域均覆盖全部 128 DIE64 卡 × 每卡 2 个逻辑 DIE微批大小MBS为 1、全局批大小GBS为 1024。这一配置在 llm_pretrain/deepseekv4 的部署指导中得到印证Flash 模型使用deepseek_v4_285b_43layers_4k_128die.toml配置在 8 机 64 卡 A3 集群上运行。FSDP 简介FSDPFully Sharded Data Parallel是 PyTorch 提供的数据并行技术通过将模型参数、梯度和优化器状态跨设备分片存储显著降低单卡内存占用。其原理是将 DDP 的 AllReduce 分解为前向的 AllGather 与后向的 ReduceScatter在计算过程中按需收集参数、用完即释以最小化常驻内存。此外FSDP 的通信操作可与计算重叠执行AllGather 提前预取下一层参数ReduceScatter 则在后向计算进行时同步传输有效将通信延迟掩盖在计算耗时内。内存收益尽管 MoE 模型的路由专家权重在全局参数量中占据绝对主导但经过 EP 维度切分后非路由专家权重在单卡内存中的占比将显著上升成为新的内存瓶颈。下表以 DeepSeek-V4-Flash 模型为例展示了 EP 切分前后瓶颈的转换其中主权重和模型梯度均按照混合精度训练场景下常见的 4 字节大小计算。FSDPDPEP总参数量非专家参数量非专家权重梯度专家参数量专家权重梯度非专家内存占比11281285B8B59.5GB277B2060GB2.8%1128128285B8B59.5GB~2B16GB80%1281128285B0.06B0.465GB~2B16GB2.8%上表给出了三种配置下的内存分布。当 FSDP1、EP1 时非专家权重梯度高达 59.5GB仅占总参数量的 2.8%但绝对值足以形成瓶颈引入 EP128 后非专家权重梯度仍为 59.5GB但其内存占比骤升至 80%成为新的内存瓶颈而进一步叠加 FSDP128 后非专家权重梯度被压缩至 0.465GB占比 2.8%内存压力被显著消除。考虑到 A3 单卡 64GB 的内存大小模型训练过程中不可避免地需要引入对非专家参数的切分。对比 TP/PPFSDP 的通信可被计算完全掩盖详见性能小节因而可以使用相对激进的内存切分配置而不受性能考量的制约。实测开启 FSDP128 后非专家权重被均分至全部 128 DIE几乎完全消除该部分内存占用——TorchTitan 完成模型初始化后单卡权重内存仅约 8~9GB可视作仅由 EP 切分后的路由专家权重构成。性能开销FSDP 通信数据不依赖当前层的中间计算结果框架得以将集合通信操作下发至独立的通信流上与模型前向/反向计算异步并发执行。**借助于 A3 超节点的高通信带宽实测 FSDP 的通信开销可被计算过程完全掩盖。**这一特性打破了性能与内存之间的折衷——采用 FSDP128 将通信域扩大到 8 机范围时依然不会因通信暴露而损害训练吞吐为内存侧的极致压缩提供了可行性保障。与之相对地TP/PP 中的通信与计算关键路径紧密耦合无法掩盖。即使借助 A3 芯片的 SIO 高带宽通信并采用 TP2 这类极小通信域配置张量并行依然会带来可观的通信开销另一方面PP 切分的流水线气泡率随通信域规模的扩大同步增加导致的设备空闲亦不可忽略。这些现实约束反过来制约了数据并行向 TP/PP 的转换实际部署中往往无法将数据并行完全切分为 TP/PP从而残留一定程度的内存冗余。易用性优势除内存与性能优势外FSDP 在工程复杂度上也显著优于 TP/PP对 TP 切分而言DeepSeek 系列模型中 MLA/DSA/C4A/C128A 结构的引入使得注意力机制的 TP 切分越来越复杂。DeepSeek-V4 的注意力机制在整体框架上继承了 DeepSeek-V3.2 引入的 DSA 结构其 Indexer 模块在计算前向结果和 Loss 值时均涉及模型 head 维度的累加操作如下图所示。该操作的语义决定了无法直接沿 head 维度进行 TP 切分——强行切分则需同时引入新的 AllGather 聚合通信以保证 ReduceSum 的正确性。对 PP 切分而言模型中非均匀的模块Embedding、LmHead、MTP 等极易导致 PP Stage 间的耗时不均MoE 模型中本身与数据分布相关的负载不均衡现象更进一步加重了 PP 性能调优过程中的困难。相较之下FSDP 对参数无差异分片并作异步 AllGather 的方式做到了模型算法和结构无感的权重切分——前反向计算时总可以获取完整的单层模型权重各个 DP rank 间执行的计算无任何差异且在适配上仅通过一行fully_shard函数调用即可较好地克服了前述场景中 TP/PP 的缺陷。双通信域 overlap 优化基于 FSDP 相对于 TP/PP 的上述优势方案选择了 FSDP128 EP128 作为模型的切分策略并在实际测试中优化了 TorchTitan 上游实现存在的一个通信域初始化问题由于前述切分策略中FSDP 域和 EP 域均与分布式并行组的全局默认通信域相同使得 PyTorch为 FSDP 和 EP 分配同一个 ProcessGroup导致两者的通信操作被迫串行执行无法并发。具体而言如上图所示EP 通信位于 MoE 模型执行的关键路径上任何阻塞都会直接延迟模型的整体计算。当 FSDP 与 EP 共享 ProcessGroup 时FSDP 的 AllGather/ReduceScatter 会与 EP 的 AllToAll 形成串行等待关系FSDP 的通信操作会阻塞 EP 通信的执行并最终进一步阻塞后续的计算算子。这意味着原本设计中被计算隐藏的 FSDP 通信暴露到了关键路径上造成整体吞吐下降。通过在初始化阶段对 EP 通信域进行独立标注FSDP 与 EP 得以持有各自的 ProcessGroup。两者的通信操作因此可在不同的 HCCL 流上并发执行从而恢复 FSDP 通信被模型主计算流程掩盖的设计初衷。修复后EP 的关键路径通信不再受 FSDP 阻塞实测结果显示修复使模型在 MoE 层的通信等待时长由最长 10ms 降低至约 3ms并带来了 1.1% 的模型训练性能提升。这一修复思路后续反馈至 PyTorch/TorchTitan 上游社区助力框架在后续版本中的优化。内存卸载优化在 FSDP2 原生方案中尽管参数、梯度和优化器状态均被切分但 AdamW 的两个 FP32 动量在前反向计算期间仍驻留设备内存造成静置浪费。作为 FSDP 的内存优化补充SwapOptimizer 将 FSDP 已切分至单卡的优化器状态进一步卸载至主机内存仅在权重更新阶段按需换入。实践参考 MindSpeed 的 swap-optimizer 特性思路在 TorchTitan 中实现了该能力针对 DTensor 权重场景适配并以参数为粒度设计切片流水将加载—更新—卸载串行过程按如下方式重叠执行降低优化器更新阶段的内存峰值。该机制将使用 AdamW 优化器场景下的常驻内存从 FSDP 的 16 字节/参数权重 4 梯度 4 优化器 8进一步压降至 8 字节/参数仅权重 梯度内存占用近乎减半。torch.compile AutoFuse随着 MoE 与多模态架构的普及动态、细粒度的小算子组合已成为主流设计模式。以 DeepSeek 提出的 mHC 为例其将 HC 过程分解为 HcRes、HcPre 与 HcPost 三部分分别通过双随机矩阵与 Sigmoid 门控系数保障数学性质但实现上依赖大量细碎的 PyTorch 操作。这类灵活结构在模型设计中日益普遍若延续传统手写融合算子的方式开发效率已难以匹配模型迭代速度算子自动化融合能力由此成为刚需。torch.compile 为此提供了基础能力通过捕获 FX 计算图将模型逻辑转换为算子图表示交由 Inductor 后端进行算子融合与代码生成。Inductor 通过垂直与水平融合策略将多个细碎操作合并为单一高效内核实现跨算子的全局优化。TorchTitan 对 torch.compile 提供了原生深度支持可充分复用 PyTorch 现有的编译生态基础设施与优化成果。然而NPU 与 GPU 在内存架构与计算范式上存在本质差异直接沿用 GPU 优化策略难以充分释放 NPU 潜力。为此**CANN 框架的 AutoFuse 组件在继承 Inductor 整体流程的基础上实现了目标为 Ascend C 的 Schedule Codegen 后端。**通过架构感知的融合策略与指令调度实现从计算图到高效 NPU 算子代码的自动化映射显著提升融合算子在昇腾平台上的执行效率。在本实践的训练场景中由于不开启 TP 切分使得各算子的计算量本身较大做算子融合后 Kernel Launch 开销占比极低不存在静态图的收益空间。因而采用动态图 Eager 模式执行编译后的算子图且 Eager 模式无需处理动态 Shape 约束更为简洁灵活已可获取融合的主要性能收益。AutoFuse 亮点使用简单PyTorch 前端增加一行代码即可使能性能提升依赖 NPU 亲和的算子生成技术本次模型融合收益约 31.8%亲和 NPU模板规约的 Schedule 技术基于硬件建模动态求解的算子 Tiling 技术基于 Ascend LoopIR 表达的Ascend C 算子 kernel 代码生成技术泛化与完备度当前已支持 152 个 LoopIR 表达中的 46 个未来将逐步补齐与 LoopIR 完整对等AutoFuse 收益分析使能 AutoFuse 后 DeepSeek-V4-Flash 模型整网吞吐收益为 31.8%具体性能提升数据拆解如下表所示收益主要来自 AIV 的算子融合以及 Host Bound 缓解带来的 Free Time 减少AIV 的收益主要来自各类算子融合融合后算子名称数量Count融合后总耗时 (s)融合前总耗时 (s)净收益时间 (s)优化核心autofused_mul_sum 系列10,2321.67.15.5带宽优化消除 Mul 结果写回内存再读回的操作直接在片上完成规约。基础逐元素融合 (Mul/Add/Div)~20万个7.812.85访存收敛将大量独立的乘、加、除操作合并显著减少了对内存的总访问次数。autofused_add_div_expand_mul_pow1,58412.61.6指令加速将极高延迟的独立 Pow 算子指令化通过计算掩盖访存延迟。BroadcastTo 隐式化 (聚合项)~3.3万个0.32.52.2隐式广播通过 Stride 逻辑变换实现广播彻底消除了 2s 多的物理内存拷贝耗时。autofused_npu_dtype_cast 系列14,8321.82.50.7调度简化吸收数万次类型转换操作大幅降低了 CPU 下发任务给 NPU 的频率。以下以其中一个片段为例说明三大核心收益来源消除访存墙Memory Bound将 12 个独立算子的读-算-写循环重构为 2 个高集成度融合内核中间张量驻留片上缓存SRAM/L1规避了频繁读写 HBM全局内存的巨大带宽开销掩盖高耗时指令Operator Hiding通过编译器对 Pow/RealDiv 等高周期算子进行指令级重排计算压力被掩盖在矢量计算流水线中实现非线性数学公式链的近乎零成本合并精简内核启动开销Launch Overhead大幅减少 Host-to-Device 握手次数清理了重量级算子如 MatMul之间的计算碎屑确保 NPU 核心持续处于高负载产出状态Host-Bound调度瓶颈的显著削减核心优化路径将原本离散间隔执行的小算子优化为融合算子的连续执行降低了硬件空泡保障 NPU 计算流水线连续性大幅提升硬件利用率。调度优化静态化分析 Codegen Wrapper → 预编译执行流取代动态 Dispatcher降低 Host 延迟算子级优化算子融合 → 减少 Launch 次数 → 缓解 Host 侧指令下发压力高性能融合算子优化针对 DeepSeek-V4-Flash 新模型的 Attention 核心结构设计并实现了 4 个在训练场景下使用的高性能 Ascend C 算子提升核心计算模块的性能并优化内存。LightningIndexerLightningIndexerLI算子基于一系列操作得到每个 token 对应的 Top-k 位置输出 Top-k 位置的索引供 SparseAttnSharedkv 作为输入完成计算具体计算公式如下$$ \operatorname{Top-}k \Bigl{ [1]{1 \times g} \bigl[ (W [1]{1 \times S_k}) \odot \operatorname{ReLU}( Q_{\mathrm{index}} K_{\mathrm{index}}^\top ) \bigr] \Bigr} $$LightningIndexer 的计算流程可分为 3 个阶段C0Cube 操作即 Q 和 K 的矩阵乘以及 ReLU 操作V1Vector 操作ReLU 后续的多个向量计算不包含 TopkV2Vector 操作待前置完整分数计算完毕后再通过二次分核完成每个 token 的 Topk 计算针对流水排布本算子设计了一次 Preload C 过程提升 C 和 V 计算间的流水并行度流水排布图示例如下SparseAttnSharedkvSparseAttnSharedkvSAS算子旨在完成以下公式描述的 Attention 计算根据输入 cmp_ratio 不同支持 3 种 Attention 计算分别为 Sliding Window AttentionSWA、Compressed AttentionCFA以及 Sparse Compressed AttentionSCFA。其中 Attention 部分为 Multi-Query Attention$$ S Q \tilde{K}^\top $$$$ m \max\bigl( \mathrm{sinks},; \max(S) \bigr) $$$$ \mathrm{Attention} \frac{ e^{S - m} \tilde{V} }{ \sum e^{S - m} e^{\mathrm{sinks} - m} } $$其中$\tilde{K}$ 和 $\tilde{V}$ 为基于 ori_kv原始的 KV、cmp_kv压缩后的 KV以及 cmp_ratio压缩率等入参控制的实际参与计算的 K 和 V。SAS 算子实现主要分为 5 个阶段分别为 V0/C1/V1/C2/V2其中 V 为 vector 计算C 为 Cube 计算对应计算公式如下$$ V_0 : \operatorname{Gather}\bigl( \mathrm{cmpkv},; \mathrm{topkIndices}[i] \bigr), \quad 0 \le i \mathrm{selectBlockCount} $$$$ C_1 : \mathrm{qk} Q K^\top $$$$ V_1 : P \operatorname{onlineSoftmax}( \mathrm{qk},; \mathrm{sinks} ) $$$$ C_2 : O P V $$$$ V_2 : O \operatorname{rescale}( O ) $$实现过程中流水排布时通过 Preload 一轮 V0C1 使得不同阶段间的依赖错开实现除头尾以外的 CV 流水并行。SparseAttnSharedkvGradSparseAttnSharedkvGradSASG是 SAS 的反向算子算子的计算流程如下SASG 算子主流程可分为 5 个阶段依次为 Gather、Cube12、Process、Cube345、Scatter如图中所示。算子流水排布实现如下通过 Preload 一次 GatherCube12 达成 CV 流水并行的效果。SparseLightningIndexerGradKLLoss由于 LI 模块进行 Loss 计算时存在巨大内存开销内存开销达到序列长度的平方级别因为需要计算 Main Attention score。SparseLightningIndexerGradKLLoss 算子将 Main Attention score 计算、LI 的反向以及 Loss 计算过程融合减少中间内存占用优化内存和性能。算子计算公式如下LI 中取 Top-k 的 value 的计算公式可以表示为$$ I_{t,:} W_{t,:} \operatorname{ReLU}\bigl( q_{t,:} (K_{:t,:})^\top \bigr) $$LI 单独训练时对应的 loss function 为$$ \mathcal{L}(I) \sum_t D_{\mathrm{KL}} \bigl( p_{t,:} ;|;\operatorname{Softmax}(I_{t,:}) \bigr) $$其中p 是 target distribution通过对 Main Attention score 在所有 head 维度上求和然后把求和结果沿着上下文方向进行 L1 正则化得到。其中$D_{\mathrm{KL}}$ 为 KL 散度其表达式为$$ D_{\mathrm{KL}}(a | b) \sum_i a_i \log \frac{a_i}{b_i} $$通过求导可得 Loss 的梯度表达式$$ \mathrm{d}I_{t,:} \operatorname{Softmax}(I_{t,:}) - p_{t,:} $$利用链式法则可进一步计算 weight、query 和 key 矩阵的梯度$$ \mathrm{d}W_{t,:} \mathrm{d}I_{t,:} \bigl( \operatorname{ReLU}( S_{t,:} ) \bigr)^\top $$$$ \mathrm{d}q_{t,:} \mathrm{d}S_{t,:} K_{:t,:} $$$$ \mathrm{d}K_{:t,:} ( \mathrm{d}S_{t,:} )^\top q_{:t,:} $$其中S 为 QK 矩阵 softmax 的结果计算过程可以拆分成 5 个阶段V0依据 Top-k 的索引从 Main Attention 的 K 和 LI 的 K 中提取有效数据C1完成 Main Attention 原始 Q 和 K 以及 LI 的 Q 和 K 矩阵运算V1完成 KLLoss、dW 计算C2完成 LI 反向的 dQ 计算V2通过 ScatterAdd 完成 LI 反向的 dK 计算在实现过程中针对 V0 进行两次 Preload并通过 PingPong 掩盖 C1 和 C2 的计算提升流水的并行度。性能结果与未来展望基于 A3 SuperPods 64 卡的性能结果依托 CANN 平台与 TorchTitan-NPU 插件实践在 A3 64 卡集群上快速完成了 DeepSeek-V4-Flash 模型的基础训练性能调优。方案采用大 EP 纯 FSDP 的极简并行切分策略集成针对稀疏注意力模块开发的融合算子并结合 Ascend C AutoFuse 自动融合机制实现了模型吞吐从初始 397 tokens/p/s 到 1100 tokens/p/s 的显著提升。其中定制融合算子和 AutoFuse 分别贡献了约 90% 和 30% 的吞吐提升两者存在收益叠加效应非简单相加。CUBEFAVECEPFSDPFREEMFU16.64s(27.89%)11.8s(19.77%)22.08s(37.01%)5.70s(9.56%)1.78s(2.98%)1.67s(2.79%)28.78%50.52s(84.68%)7.48s(12.53%)未来展望从 A3 集群现阶段的 Profiling 数据来看计算耗时占据绝对主导达到总时间的 84.68%其中 Vector 类算子更占到整网耗时的近 40%。因此进一步的性能优化可重点围绕以下方向展开采用更精细的按需重计算策略避免对无需保留激活值的 Vector 算子进行冗余重计算进一步扩大 AutoFuse 自动融合的覆盖范围以压缩计算时延当前 TorchTitan-NPU 版本尚未集成 MC2 或基于算子流水编排的 EP 域通信计算并行机制将在后续阶段进行能力补齐在功能拓展方面后续计划同步跟进 DeepSeek-V4 技术报告中的演进方向针对下一代 A5 平台支持 FP8 与 A8W4 低精度量化训练特性发挥 A5 代际的 MxFP8/MxFP4 低精度计算和通信能力在 TorchTitan-NPU 中集成 Muon 优化器功能支持以追求更快的模型收敛效果并配套提供 AutoFuse 与融合算子加速能力实战部署基于仓库脚本的 DeepSeek-V4-Flash 训练复现本仓库 llm_pretrain/deepseekv4 目录提供了一套可直接落地的部署配套与技术报告形成原理—实践闭环。以下操作步骤均以该目录下的脚本与说明为准。硬件与软件要求产品型号Atlas A3 系列DeepSeek-V4-Flash 最少 8 机 64 卡DeepSeek-V4-Pro 需要 24 机 192 卡操作系统Linux ARM驱动版本Ascend HDK 25.5.2如需支持虚拟优化器特性需更新固件/驱动至 25.5.2 并配套安装CANN 版本9.0.0镜像版本dsv4_train_torchtitan:cann9.0.0_v3.0建议先使用npu-smi info检查 Ascend NPU 固件与驱动是否正确安装且版本匹配。源码与脚本准备在 A3 环境执行如下命令拉取 TorchTitan-NPU 源码v0.2.2-dev分支并拷贝本仓库脚本至其scripts目录mkdir -p /home/code cd /home/code/ git clone -b v0.2.2-dev https://gitcode.com/cann/torchtitan-npu.git cd torchtitan-npucd ../ git clone https://gitcode.com/cann/cann-recipes-train.git cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_multinodes_dsv4_flash_pretrain.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_multinodes_dsv4_flash_perf.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_multinodes_dsv4_pro_pretrain.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_dsv4_flash_A5_BF16_pretrain.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_dsv4_flash_A5_MXFP8_pretrain.sh ./torchtitan-npu/scriptsDeepSeek-V4-Flash 模型训练使用配置文件./torchtitan_npu/models/deepseek_v4/train_configs/deepseek_v4_285b_43layers_4k_128die.tomlDeepSeek-V4-Pro 模型训练使用配置文件./torchtitan_npu/models/deepseek_v4/train_configs/deepseek_v4_pro_61layers_4k_384die.toml容器拉起与环境初始化加载 Docker 镜像并启动容器注意将 16 个 davinci 设备、davinci_manager、devmm_svm等设备透传进入容器并挂载驱动目录、/etc/hccn.conf、/usr/local/dcmi等关键路径gunzip -c dsv4_train_torchtitan_cann9.0.0_v3.0.tar.gz | docker load docker run -u root -itd --name dsv4_train_torchtitan_v3.0 --ulimit nproc65535:65535 --ipchost \ --device/dev/davinci0 ... --device/dev/davinci15 \ --device/dev/davinci_manager --device/dev/devmm_svm --device/dev/hisi_hdc \ -v /home:/home -v /data:/data -v /etc/localtime:/etc/localtime \ -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ -v /etc/ascend_install.info:/etc/ascend_install.info -v /var/log/npu/:/usr/slog \ -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /sys/fs/cgroup:/sys/fs/cgroup:ro \ -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/sbin:/usr/local/sbin \ -v /etc/hccn.conf:/etc/hccn.conf -v /root/.pip:/root/.pip -v /etc/hosts:/etc/hosts \ -v /usr/bin/hostname:/usr/bin/hostname \ --nethost --shm-size128g --privileged \ dsv4_train_torchtitan:cann9.0.0_v3.0 /bin/bash进入容器后初始化 CANN 环境与自定义算子环境docker exec -it dsv4_train_torchtitan_v3.0 /bin/bash source /usr/local/Ascend/cann/set_env.sh source /usr/local/Ascend/cann/opp/vendors/custom_transformer/bin/set_env.bash数据集与权重准备配置默认使用仓内样例数据集c4_testdataset_path ./tests/assets/c4_test若使用自定义数据集需提前准备数据目录并修改配置中的dataset与dataset_path。DeepSeek-V4-Flash 的原始权重FP8 格式需要先转换为 BF16 权重转换工具即本仓库 llm_pretrain/deepseekv4/utils/convert_model.py。该脚本读取model.safetensors.index.json与config.json对单字节FP8/INT8权重结合对应的.scale缩放因子执行反量化并同步更新权重索引与配置文件cd /home/code/cann-recipes-train/llm_pretrain/deepseekv4/utils python3 convert_model.py \ --input_fp8_hf_path /data/models/DeepSeek-V4-Flash \ --output_hf_path /data/models/DeepSeek-V4-Flash-bf16 \ --quant_type bfloat16convert_model.py的--quant_type参数还支持w8a8-int、w8a8-mx、w4a8-mx等量化输出见 convert_model.py其中 MX 格式的量化实现位于 mx_quantize.py量化配置生成逻辑位于 convert_config.py对应 A5 平台低精度训练的前置准备。训练配置与多机启动拉起训练前请重点确认训练配置文件中的路径与实际环境一致[model] hf_assets_path /data/models/DeepSeek-V4-Flash-bf16 [training] dataset c4_test dataset_path ./tests/assets/c4_test [checkpoint] initial_load_in_hf true initial_load_path /data/models/DeepSeek-V4-Flash-bf16多机启动脚本 run_train_multinodes_dsv4_flash_pretrain.sh 中需要按实际环境修改的部分包括网络接口Network_Interface、节点 IP 列表IPs、每节点 NPU 数NPUS_PER_NODE以及 HCCL 通信端口等。脚本内部通过torchrun --nnodes/--node_rank/--nproc_per_node拉起训练并支持NGPU、CONFIG_FILE、LOG_RANK等环境变量覆盖默认值。在所有参与训练的节点上同时执行CONFIG_FILE./torchtitan_npu/models/deepseek_v4/train_configs/deepseek_v4_285b_43layers_4k_128die.toml \ bash scripts/run_train_multinodes_dsv4_flash_pretrain.sh若需复现最优性能任务则使用deepseek_v4_285b_43layers_4k_128die_perf.toml配合 run_train_multinodes_dsv4_flash_perf.sh 启动。此外针对 A5 平台还提供了单机 BF16/MXFP8 预训练脚本run_train_dsv4_flash_A5_BF16_pretrain.sh、run_train_dsv4_flash_A5_MXFP8_pretrain.sh以及 HiF8 低精度预训练脚本run_train_dsv4_flash_A5_HiF8_single_node.sh后者基于torchtitan_npu/experiments/ao_npu/的 ParamSwap 参数级量化后端MatMul 与 Grouped MatMul 走 per-tensor 动态量化 npu_quant_matmul的真实低精度计算路径可用于打通 A5 低精度训练链路与性能观测。小结DeepSeek-V4 昇腾训练优化实践验证了一条极简开箱的 NPU 大模型训练技术路线以 TorchTitan-NPU 插件化方案承接 PyTorch Native 生态以纯 FSDP 大 EP的极简并行切分实现内存最优与通信隐藏以 torch.compile AutoFuse 使能 Vector 算子自动融合以定制 Ascend C 算子释放稀疏注意力结构的芯片算力最终在 A3 64 卡集群上将 DeepSeek-V4-Flash 的 BF16 训练吞吐提升至 1100 tokens/p/s。该方案为现有昇腾 A3 集群用户基于 DeepSeek 新模型架构快速开展续训练/SFT 及自研算法验证提供了可直接复用的参考实现。【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 DeepSeek-V4 昇腾训练优化实践：基于 CANN 的 TorchTitan-NPU 与 AutoFuse 极简训练方案
DeepSeek-V4 昇腾训练优化实践基于 CANN 的 TorchTitan-NPU 与 AutoFuse 极简训练方案【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train导读本文围绕 CANN 平台上 DeepSeek-V4-Flash 模型的训练优化实践展开介绍基于 TorchTitan-NPU 框架、采用纯 FSDP 大 EP极简并行策略、依托 torch.compile Ascend C AutoFuse 自动融合以及面向稀疏注意力结构的定制融合算子将 A3 集群 64 卡 4K 序列 BF16 训练吞吐从 397 提升至 1100 tokens/p/s 的完整技术路线。读者将掌握该方案的软件分层架构、并行切分与内存优化原理、AutoFuse 收益拆解方法以及基于 llm_pretrain/deepseekv4 目录下的部署脚本完成源码准备、权重转换、容器拉起与多机训练启动的实战路径。Highlights极简并行策略训练框架采用 TorchTitan TorchTitan-NPU 插件化方案以超节点亲和的大 EP 纯 FSDP的精简并行切分策略在极低适配成本和通信开销下实现内存占用最优取得易用性与性能的较好均衡。训练入图自动融合TorchTitan-NPU 深度适配 torch.compile 机制使能训练入图技术依托 Inductor AutoFuse基于 Ascend C 的 Codegen 后端实现端到端的 Vector 算子自动融合为整网带来高达 31.8% 的开箱即用性能收益。定制 NPU 融合算子针对稀疏注意力等复杂结构开发 SparseAttnSharedkv、LightningIndexer 等 4 个高效 NPU 融合算子从负载均衡分核计算、内存与计算均衡等维度协同优化充分释放芯片稀疏算力。可复现性能结果基于上述优化点CANN 已基于 TorchTitan 支持 DeepSeek-V4-Flash 模型训练A3 集群 BF16 精度 64 卡 4K 序列 MTP1 训练吞吐达 1100 tokens/p/s。引言为什么选择 TorchTitan-NPU AutoFuseTorchTitan 为业界带来了 PyTorch Native 大模型训练方案的全新选择其核心设计理念是模型算法与分布式并行、算子优化天然解耦。通过支持 torch.compile 训练入图优化告别手写融合 kernel大大提升了大规模分布式训练的易用性。在并行策略方面TorchTitan 深度集成 FSDP2实现参数、梯度和优化器状态的跨设备分片显著降低单卡内存占用。本实践的整体软件架构自顶向下由 TorchTitan、TorchTitan-NPU 插件、TorchInductor 及 CANN 层的 AutoFuse 组件构成各层分工明确、优势互补TorchTitan作为 PyTorch 原生分布式训练框架避免对 Megatron 等重型框架的依赖以精简代码库提供 FSDP 等并行能力的开箱即用支持。TorchTitan-NPU以插件形式提供可插拔安装针对昇腾集群实现通信与算子的硬件加速适配保持与上游 TorchTitan 的轻量集成。TorchInductor直接继承社区 torch.compile 的图捕获与融合优化能力无需重复造轮即可复用 PyTorch 编译生态的相关基础设施。AutoFuse面向昇腾 NPU定位为基于 Ascend C 的 Codegen Schedule 后端在 Inductor 基础上注入架构感知的融合、优化与展开策略将计算图优化转化为 NPU 高效指令序列。基于上述分层架构实践在 TorchTitan-NPU 加速插件的基础上完成了 DeepSeek-V4-Flash 模型续训练流程的适配与调优。整个方案围绕极简并行、自动融合、定制算子三条主线展开下文分别就基于 FSDPEP 的极简分布式并行、torch.compile 与 AutoFuse 自动融合协同、以及高性能融合算子优化加速进行详细阐述。极简分布式并行优化FSDPEP 极简切分优势在 MoE 大模型训练中EP 切分可有效分摊路由专家权重的存储压力但注意力模块、共享专家、词嵌入及词表投影层等非专家参数仍需额外切分机制加以应对。另一方面随着模型结构迎来又一个快速发展阶段传统 TP/PP 方案在性能开销与易用性层面均面临诸多挑战而 FSDP 在上述场景中展现出较为显著的优势。因此本实践基于 A3 超节点 64 卡环境在 DeepSeek-V4-Flash 模型训练中创新性地选择纯 FSDP EP 极简切分方案具体配置如下TPPPVPPEPFSDPMBSGBS集群1111281281102464卡即完全不启用张量并行TP1与流水线并行PP1EP 与 FSDP 域均覆盖全部 128 DIE64 卡 × 每卡 2 个逻辑 DIE微批大小MBS为 1、全局批大小GBS为 1024。这一配置在 llm_pretrain/deepseekv4 的部署指导中得到印证Flash 模型使用deepseek_v4_285b_43layers_4k_128die.toml配置在 8 机 64 卡 A3 集群上运行。FSDP 简介FSDPFully Sharded Data Parallel是 PyTorch 提供的数据并行技术通过将模型参数、梯度和优化器状态跨设备分片存储显著降低单卡内存占用。其原理是将 DDP 的 AllReduce 分解为前向的 AllGather 与后向的 ReduceScatter在计算过程中按需收集参数、用完即释以最小化常驻内存。此外FSDP 的通信操作可与计算重叠执行AllGather 提前预取下一层参数ReduceScatter 则在后向计算进行时同步传输有效将通信延迟掩盖在计算耗时内。内存收益尽管 MoE 模型的路由专家权重在全局参数量中占据绝对主导但经过 EP 维度切分后非路由专家权重在单卡内存中的占比将显著上升成为新的内存瓶颈。下表以 DeepSeek-V4-Flash 模型为例展示了 EP 切分前后瓶颈的转换其中主权重和模型梯度均按照混合精度训练场景下常见的 4 字节大小计算。FSDPDPEP总参数量非专家参数量非专家权重梯度专家参数量专家权重梯度非专家内存占比11281285B8B59.5GB277B2060GB2.8%1128128285B8B59.5GB~2B16GB80%1281128285B0.06B0.465GB~2B16GB2.8%上表给出了三种配置下的内存分布。当 FSDP1、EP1 时非专家权重梯度高达 59.5GB仅占总参数量的 2.8%但绝对值足以形成瓶颈引入 EP128 后非专家权重梯度仍为 59.5GB但其内存占比骤升至 80%成为新的内存瓶颈而进一步叠加 FSDP128 后非专家权重梯度被压缩至 0.465GB占比 2.8%内存压力被显著消除。考虑到 A3 单卡 64GB 的内存大小模型训练过程中不可避免地需要引入对非专家参数的切分。对比 TP/PPFSDP 的通信可被计算完全掩盖详见性能小节因而可以使用相对激进的内存切分配置而不受性能考量的制约。实测开启 FSDP128 后非专家权重被均分至全部 128 DIE几乎完全消除该部分内存占用——TorchTitan 完成模型初始化后单卡权重内存仅约 8~9GB可视作仅由 EP 切分后的路由专家权重构成。性能开销FSDP 通信数据不依赖当前层的中间计算结果框架得以将集合通信操作下发至独立的通信流上与模型前向/反向计算异步并发执行。**借助于 A3 超节点的高通信带宽实测 FSDP 的通信开销可被计算过程完全掩盖。**这一特性打破了性能与内存之间的折衷——采用 FSDP128 将通信域扩大到 8 机范围时依然不会因通信暴露而损害训练吞吐为内存侧的极致压缩提供了可行性保障。与之相对地TP/PP 中的通信与计算关键路径紧密耦合无法掩盖。即使借助 A3 芯片的 SIO 高带宽通信并采用 TP2 这类极小通信域配置张量并行依然会带来可观的通信开销另一方面PP 切分的流水线气泡率随通信域规模的扩大同步增加导致的设备空闲亦不可忽略。这些现实约束反过来制约了数据并行向 TP/PP 的转换实际部署中往往无法将数据并行完全切分为 TP/PP从而残留一定程度的内存冗余。易用性优势除内存与性能优势外FSDP 在工程复杂度上也显著优于 TP/PP对 TP 切分而言DeepSeek 系列模型中 MLA/DSA/C4A/C128A 结构的引入使得注意力机制的 TP 切分越来越复杂。DeepSeek-V4 的注意力机制在整体框架上继承了 DeepSeek-V3.2 引入的 DSA 结构其 Indexer 模块在计算前向结果和 Loss 值时均涉及模型 head 维度的累加操作如下图所示。该操作的语义决定了无法直接沿 head 维度进行 TP 切分——强行切分则需同时引入新的 AllGather 聚合通信以保证 ReduceSum 的正确性。对 PP 切分而言模型中非均匀的模块Embedding、LmHead、MTP 等极易导致 PP Stage 间的耗时不均MoE 模型中本身与数据分布相关的负载不均衡现象更进一步加重了 PP 性能调优过程中的困难。相较之下FSDP 对参数无差异分片并作异步 AllGather 的方式做到了模型算法和结构无感的权重切分——前反向计算时总可以获取完整的单层模型权重各个 DP rank 间执行的计算无任何差异且在适配上仅通过一行fully_shard函数调用即可较好地克服了前述场景中 TP/PP 的缺陷。双通信域 overlap 优化基于 FSDP 相对于 TP/PP 的上述优势方案选择了 FSDP128 EP128 作为模型的切分策略并在实际测试中优化了 TorchTitan 上游实现存在的一个通信域初始化问题由于前述切分策略中FSDP 域和 EP 域均与分布式并行组的全局默认通信域相同使得 PyTorch为 FSDP 和 EP 分配同一个 ProcessGroup导致两者的通信操作被迫串行执行无法并发。具体而言如上图所示EP 通信位于 MoE 模型执行的关键路径上任何阻塞都会直接延迟模型的整体计算。当 FSDP 与 EP 共享 ProcessGroup 时FSDP 的 AllGather/ReduceScatter 会与 EP 的 AllToAll 形成串行等待关系FSDP 的通信操作会阻塞 EP 通信的执行并最终进一步阻塞后续的计算算子。这意味着原本设计中被计算隐藏的 FSDP 通信暴露到了关键路径上造成整体吞吐下降。通过在初始化阶段对 EP 通信域进行独立标注FSDP 与 EP 得以持有各自的 ProcessGroup。两者的通信操作因此可在不同的 HCCL 流上并发执行从而恢复 FSDP 通信被模型主计算流程掩盖的设计初衷。修复后EP 的关键路径通信不再受 FSDP 阻塞实测结果显示修复使模型在 MoE 层的通信等待时长由最长 10ms 降低至约 3ms并带来了 1.1% 的模型训练性能提升。这一修复思路后续反馈至 PyTorch/TorchTitan 上游社区助力框架在后续版本中的优化。内存卸载优化在 FSDP2 原生方案中尽管参数、梯度和优化器状态均被切分但 AdamW 的两个 FP32 动量在前反向计算期间仍驻留设备内存造成静置浪费。作为 FSDP 的内存优化补充SwapOptimizer 将 FSDP 已切分至单卡的优化器状态进一步卸载至主机内存仅在权重更新阶段按需换入。实践参考 MindSpeed 的 swap-optimizer 特性思路在 TorchTitan 中实现了该能力针对 DTensor 权重场景适配并以参数为粒度设计切片流水将加载—更新—卸载串行过程按如下方式重叠执行降低优化器更新阶段的内存峰值。该机制将使用 AdamW 优化器场景下的常驻内存从 FSDP 的 16 字节/参数权重 4 梯度 4 优化器 8进一步压降至 8 字节/参数仅权重 梯度内存占用近乎减半。torch.compile AutoFuse随着 MoE 与多模态架构的普及动态、细粒度的小算子组合已成为主流设计模式。以 DeepSeek 提出的 mHC 为例其将 HC 过程分解为 HcRes、HcPre 与 HcPost 三部分分别通过双随机矩阵与 Sigmoid 门控系数保障数学性质但实现上依赖大量细碎的 PyTorch 操作。这类灵活结构在模型设计中日益普遍若延续传统手写融合算子的方式开发效率已难以匹配模型迭代速度算子自动化融合能力由此成为刚需。torch.compile 为此提供了基础能力通过捕获 FX 计算图将模型逻辑转换为算子图表示交由 Inductor 后端进行算子融合与代码生成。Inductor 通过垂直与水平融合策略将多个细碎操作合并为单一高效内核实现跨算子的全局优化。TorchTitan 对 torch.compile 提供了原生深度支持可充分复用 PyTorch 现有的编译生态基础设施与优化成果。然而NPU 与 GPU 在内存架构与计算范式上存在本质差异直接沿用 GPU 优化策略难以充分释放 NPU 潜力。为此**CANN 框架的 AutoFuse 组件在继承 Inductor 整体流程的基础上实现了目标为 Ascend C 的 Schedule Codegen 后端。**通过架构感知的融合策略与指令调度实现从计算图到高效 NPU 算子代码的自动化映射显著提升融合算子在昇腾平台上的执行效率。在本实践的训练场景中由于不开启 TP 切分使得各算子的计算量本身较大做算子融合后 Kernel Launch 开销占比极低不存在静态图的收益空间。因而采用动态图 Eager 模式执行编译后的算子图且 Eager 模式无需处理动态 Shape 约束更为简洁灵活已可获取融合的主要性能收益。AutoFuse 亮点使用简单PyTorch 前端增加一行代码即可使能性能提升依赖 NPU 亲和的算子生成技术本次模型融合收益约 31.8%亲和 NPU模板规约的 Schedule 技术基于硬件建模动态求解的算子 Tiling 技术基于 Ascend LoopIR 表达的Ascend C 算子 kernel 代码生成技术泛化与完备度当前已支持 152 个 LoopIR 表达中的 46 个未来将逐步补齐与 LoopIR 完整对等AutoFuse 收益分析使能 AutoFuse 后 DeepSeek-V4-Flash 模型整网吞吐收益为 31.8%具体性能提升数据拆解如下表所示收益主要来自 AIV 的算子融合以及 Host Bound 缓解带来的 Free Time 减少AIV 的收益主要来自各类算子融合融合后算子名称数量Count融合后总耗时 (s)融合前总耗时 (s)净收益时间 (s)优化核心autofused_mul_sum 系列10,2321.67.15.5带宽优化消除 Mul 结果写回内存再读回的操作直接在片上完成规约。基础逐元素融合 (Mul/Add/Div)~20万个7.812.85访存收敛将大量独立的乘、加、除操作合并显著减少了对内存的总访问次数。autofused_add_div_expand_mul_pow1,58412.61.6指令加速将极高延迟的独立 Pow 算子指令化通过计算掩盖访存延迟。BroadcastTo 隐式化 (聚合项)~3.3万个0.32.52.2隐式广播通过 Stride 逻辑变换实现广播彻底消除了 2s 多的物理内存拷贝耗时。autofused_npu_dtype_cast 系列14,8321.82.50.7调度简化吸收数万次类型转换操作大幅降低了 CPU 下发任务给 NPU 的频率。以下以其中一个片段为例说明三大核心收益来源消除访存墙Memory Bound将 12 个独立算子的读-算-写循环重构为 2 个高集成度融合内核中间张量驻留片上缓存SRAM/L1规避了频繁读写 HBM全局内存的巨大带宽开销掩盖高耗时指令Operator Hiding通过编译器对 Pow/RealDiv 等高周期算子进行指令级重排计算压力被掩盖在矢量计算流水线中实现非线性数学公式链的近乎零成本合并精简内核启动开销Launch Overhead大幅减少 Host-to-Device 握手次数清理了重量级算子如 MatMul之间的计算碎屑确保 NPU 核心持续处于高负载产出状态Host-Bound调度瓶颈的显著削减核心优化路径将原本离散间隔执行的小算子优化为融合算子的连续执行降低了硬件空泡保障 NPU 计算流水线连续性大幅提升硬件利用率。调度优化静态化分析 Codegen Wrapper → 预编译执行流取代动态 Dispatcher降低 Host 延迟算子级优化算子融合 → 减少 Launch 次数 → 缓解 Host 侧指令下发压力高性能融合算子优化针对 DeepSeek-V4-Flash 新模型的 Attention 核心结构设计并实现了 4 个在训练场景下使用的高性能 Ascend C 算子提升核心计算模块的性能并优化内存。LightningIndexerLightningIndexerLI算子基于一系列操作得到每个 token 对应的 Top-k 位置输出 Top-k 位置的索引供 SparseAttnSharedkv 作为输入完成计算具体计算公式如下$$ \operatorname{Top-}k \Bigl{ [1]{1 \times g} \bigl[ (W [1]{1 \times S_k}) \odot \operatorname{ReLU}( Q_{\mathrm{index}} K_{\mathrm{index}}^\top ) \bigr] \Bigr} $$LightningIndexer 的计算流程可分为 3 个阶段C0Cube 操作即 Q 和 K 的矩阵乘以及 ReLU 操作V1Vector 操作ReLU 后续的多个向量计算不包含 TopkV2Vector 操作待前置完整分数计算完毕后再通过二次分核完成每个 token 的 Topk 计算针对流水排布本算子设计了一次 Preload C 过程提升 C 和 V 计算间的流水并行度流水排布图示例如下SparseAttnSharedkvSparseAttnSharedkvSAS算子旨在完成以下公式描述的 Attention 计算根据输入 cmp_ratio 不同支持 3 种 Attention 计算分别为 Sliding Window AttentionSWA、Compressed AttentionCFA以及 Sparse Compressed AttentionSCFA。其中 Attention 部分为 Multi-Query Attention$$ S Q \tilde{K}^\top $$$$ m \max\bigl( \mathrm{sinks},; \max(S) \bigr) $$$$ \mathrm{Attention} \frac{ e^{S - m} \tilde{V} }{ \sum e^{S - m} e^{\mathrm{sinks} - m} } $$其中$\tilde{K}$ 和 $\tilde{V}$ 为基于 ori_kv原始的 KV、cmp_kv压缩后的 KV以及 cmp_ratio压缩率等入参控制的实际参与计算的 K 和 V。SAS 算子实现主要分为 5 个阶段分别为 V0/C1/V1/C2/V2其中 V 为 vector 计算C 为 Cube 计算对应计算公式如下$$ V_0 : \operatorname{Gather}\bigl( \mathrm{cmpkv},; \mathrm{topkIndices}[i] \bigr), \quad 0 \le i \mathrm{selectBlockCount} $$$$ C_1 : \mathrm{qk} Q K^\top $$$$ V_1 : P \operatorname{onlineSoftmax}( \mathrm{qk},; \mathrm{sinks} ) $$$$ C_2 : O P V $$$$ V_2 : O \operatorname{rescale}( O ) $$实现过程中流水排布时通过 Preload 一轮 V0C1 使得不同阶段间的依赖错开实现除头尾以外的 CV 流水并行。SparseAttnSharedkvGradSparseAttnSharedkvGradSASG是 SAS 的反向算子算子的计算流程如下SASG 算子主流程可分为 5 个阶段依次为 Gather、Cube12、Process、Cube345、Scatter如图中所示。算子流水排布实现如下通过 Preload 一次 GatherCube12 达成 CV 流水并行的效果。SparseLightningIndexerGradKLLoss由于 LI 模块进行 Loss 计算时存在巨大内存开销内存开销达到序列长度的平方级别因为需要计算 Main Attention score。SparseLightningIndexerGradKLLoss 算子将 Main Attention score 计算、LI 的反向以及 Loss 计算过程融合减少中间内存占用优化内存和性能。算子计算公式如下LI 中取 Top-k 的 value 的计算公式可以表示为$$ I_{t,:} W_{t,:} \operatorname{ReLU}\bigl( q_{t,:} (K_{:t,:})^\top \bigr) $$LI 单独训练时对应的 loss function 为$$ \mathcal{L}(I) \sum_t D_{\mathrm{KL}} \bigl( p_{t,:} ;|;\operatorname{Softmax}(I_{t,:}) \bigr) $$其中p 是 target distribution通过对 Main Attention score 在所有 head 维度上求和然后把求和结果沿着上下文方向进行 L1 正则化得到。其中$D_{\mathrm{KL}}$ 为 KL 散度其表达式为$$ D_{\mathrm{KL}}(a | b) \sum_i a_i \log \frac{a_i}{b_i} $$通过求导可得 Loss 的梯度表达式$$ \mathrm{d}I_{t,:} \operatorname{Softmax}(I_{t,:}) - p_{t,:} $$利用链式法则可进一步计算 weight、query 和 key 矩阵的梯度$$ \mathrm{d}W_{t,:} \mathrm{d}I_{t,:} \bigl( \operatorname{ReLU}( S_{t,:} ) \bigr)^\top $$$$ \mathrm{d}q_{t,:} \mathrm{d}S_{t,:} K_{:t,:} $$$$ \mathrm{d}K_{:t,:} ( \mathrm{d}S_{t,:} )^\top q_{:t,:} $$其中S 为 QK 矩阵 softmax 的结果计算过程可以拆分成 5 个阶段V0依据 Top-k 的索引从 Main Attention 的 K 和 LI 的 K 中提取有效数据C1完成 Main Attention 原始 Q 和 K 以及 LI 的 Q 和 K 矩阵运算V1完成 KLLoss、dW 计算C2完成 LI 反向的 dQ 计算V2通过 ScatterAdd 完成 LI 反向的 dK 计算在实现过程中针对 V0 进行两次 Preload并通过 PingPong 掩盖 C1 和 C2 的计算提升流水的并行度。性能结果与未来展望基于 A3 SuperPods 64 卡的性能结果依托 CANN 平台与 TorchTitan-NPU 插件实践在 A3 64 卡集群上快速完成了 DeepSeek-V4-Flash 模型的基础训练性能调优。方案采用大 EP 纯 FSDP 的极简并行切分策略集成针对稀疏注意力模块开发的融合算子并结合 Ascend C AutoFuse 自动融合机制实现了模型吞吐从初始 397 tokens/p/s 到 1100 tokens/p/s 的显著提升。其中定制融合算子和 AutoFuse 分别贡献了约 90% 和 30% 的吞吐提升两者存在收益叠加效应非简单相加。CUBEFAVECEPFSDPFREEMFU16.64s(27.89%)11.8s(19.77%)22.08s(37.01%)5.70s(9.56%)1.78s(2.98%)1.67s(2.79%)28.78%50.52s(84.68%)7.48s(12.53%)未来展望从 A3 集群现阶段的 Profiling 数据来看计算耗时占据绝对主导达到总时间的 84.68%其中 Vector 类算子更占到整网耗时的近 40%。因此进一步的性能优化可重点围绕以下方向展开采用更精细的按需重计算策略避免对无需保留激活值的 Vector 算子进行冗余重计算进一步扩大 AutoFuse 自动融合的覆盖范围以压缩计算时延当前 TorchTitan-NPU 版本尚未集成 MC2 或基于算子流水编排的 EP 域通信计算并行机制将在后续阶段进行能力补齐在功能拓展方面后续计划同步跟进 DeepSeek-V4 技术报告中的演进方向针对下一代 A5 平台支持 FP8 与 A8W4 低精度量化训练特性发挥 A5 代际的 MxFP8/MxFP4 低精度计算和通信能力在 TorchTitan-NPU 中集成 Muon 优化器功能支持以追求更快的模型收敛效果并配套提供 AutoFuse 与融合算子加速能力实战部署基于仓库脚本的 DeepSeek-V4-Flash 训练复现本仓库 llm_pretrain/deepseekv4 目录提供了一套可直接落地的部署配套与技术报告形成原理—实践闭环。以下操作步骤均以该目录下的脚本与说明为准。硬件与软件要求产品型号Atlas A3 系列DeepSeek-V4-Flash 最少 8 机 64 卡DeepSeek-V4-Pro 需要 24 机 192 卡操作系统Linux ARM驱动版本Ascend HDK 25.5.2如需支持虚拟优化器特性需更新固件/驱动至 25.5.2 并配套安装CANN 版本9.0.0镜像版本dsv4_train_torchtitan:cann9.0.0_v3.0建议先使用npu-smi info检查 Ascend NPU 固件与驱动是否正确安装且版本匹配。源码与脚本准备在 A3 环境执行如下命令拉取 TorchTitan-NPU 源码v0.2.2-dev分支并拷贝本仓库脚本至其scripts目录mkdir -p /home/code cd /home/code/ git clone -b v0.2.2-dev https://gitcode.com/cann/torchtitan-npu.git cd torchtitan-npucd ../ git clone https://gitcode.com/cann/cann-recipes-train.git cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_multinodes_dsv4_flash_pretrain.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_multinodes_dsv4_flash_perf.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_multinodes_dsv4_pro_pretrain.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_dsv4_flash_A5_BF16_pretrain.sh ./torchtitan-npu/scripts cp ./cann-recipes-train/llm_pretrain/deepseekv4/run_train_dsv4_flash_A5_MXFP8_pretrain.sh ./torchtitan-npu/scriptsDeepSeek-V4-Flash 模型训练使用配置文件./torchtitan_npu/models/deepseek_v4/train_configs/deepseek_v4_285b_43layers_4k_128die.tomlDeepSeek-V4-Pro 模型训练使用配置文件./torchtitan_npu/models/deepseek_v4/train_configs/deepseek_v4_pro_61layers_4k_384die.toml容器拉起与环境初始化加载 Docker 镜像并启动容器注意将 16 个 davinci 设备、davinci_manager、devmm_svm等设备透传进入容器并挂载驱动目录、/etc/hccn.conf、/usr/local/dcmi等关键路径gunzip -c dsv4_train_torchtitan_cann9.0.0_v3.0.tar.gz | docker load docker run -u root -itd --name dsv4_train_torchtitan_v3.0 --ulimit nproc65535:65535 --ipchost \ --device/dev/davinci0 ... --device/dev/davinci15 \ --device/dev/davinci_manager --device/dev/devmm_svm --device/dev/hisi_hdc \ -v /home:/home -v /data:/data -v /etc/localtime:/etc/localtime \ -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \ -v /etc/ascend_install.info:/etc/ascend_install.info -v /var/log/npu/:/usr/slog \ -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /sys/fs/cgroup:/sys/fs/cgroup:ro \ -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/sbin:/usr/local/sbin \ -v /etc/hccn.conf:/etc/hccn.conf -v /root/.pip:/root/.pip -v /etc/hosts:/etc/hosts \ -v /usr/bin/hostname:/usr/bin/hostname \ --nethost --shm-size128g --privileged \ dsv4_train_torchtitan:cann9.0.0_v3.0 /bin/bash进入容器后初始化 CANN 环境与自定义算子环境docker exec -it dsv4_train_torchtitan_v3.0 /bin/bash source /usr/local/Ascend/cann/set_env.sh source /usr/local/Ascend/cann/opp/vendors/custom_transformer/bin/set_env.bash数据集与权重准备配置默认使用仓内样例数据集c4_testdataset_path ./tests/assets/c4_test若使用自定义数据集需提前准备数据目录并修改配置中的dataset与dataset_path。DeepSeek-V4-Flash 的原始权重FP8 格式需要先转换为 BF16 权重转换工具即本仓库 llm_pretrain/deepseekv4/utils/convert_model.py。该脚本读取model.safetensors.index.json与config.json对单字节FP8/INT8权重结合对应的.scale缩放因子执行反量化并同步更新权重索引与配置文件cd /home/code/cann-recipes-train/llm_pretrain/deepseekv4/utils python3 convert_model.py \ --input_fp8_hf_path /data/models/DeepSeek-V4-Flash \ --output_hf_path /data/models/DeepSeek-V4-Flash-bf16 \ --quant_type bfloat16convert_model.py的--quant_type参数还支持w8a8-int、w8a8-mx、w4a8-mx等量化输出见 convert_model.py其中 MX 格式的量化实现位于 mx_quantize.py量化配置生成逻辑位于 convert_config.py对应 A5 平台低精度训练的前置准备。训练配置与多机启动拉起训练前请重点确认训练配置文件中的路径与实际环境一致[model] hf_assets_path /data/models/DeepSeek-V4-Flash-bf16 [training] dataset c4_test dataset_path ./tests/assets/c4_test [checkpoint] initial_load_in_hf true initial_load_path /data/models/DeepSeek-V4-Flash-bf16多机启动脚本 run_train_multinodes_dsv4_flash_pretrain.sh 中需要按实际环境修改的部分包括网络接口Network_Interface、节点 IP 列表IPs、每节点 NPU 数NPUS_PER_NODE以及 HCCL 通信端口等。脚本内部通过torchrun --nnodes/--node_rank/--nproc_per_node拉起训练并支持NGPU、CONFIG_FILE、LOG_RANK等环境变量覆盖默认值。在所有参与训练的节点上同时执行CONFIG_FILE./torchtitan_npu/models/deepseek_v4/train_configs/deepseek_v4_285b_43layers_4k_128die.toml \ bash scripts/run_train_multinodes_dsv4_flash_pretrain.sh若需复现最优性能任务则使用deepseek_v4_285b_43layers_4k_128die_perf.toml配合 run_train_multinodes_dsv4_flash_perf.sh 启动。此外针对 A5 平台还提供了单机 BF16/MXFP8 预训练脚本run_train_dsv4_flash_A5_BF16_pretrain.sh、run_train_dsv4_flash_A5_MXFP8_pretrain.sh以及 HiF8 低精度预训练脚本run_train_dsv4_flash_A5_HiF8_single_node.sh后者基于torchtitan_npu/experiments/ao_npu/的 ParamSwap 参数级量化后端MatMul 与 Grouped MatMul 走 per-tensor 动态量化 npu_quant_matmul的真实低精度计算路径可用于打通 A5 低精度训练链路与性能观测。小结DeepSeek-V4 昇腾训练优化实践验证了一条极简开箱的 NPU 大模型训练技术路线以 TorchTitan-NPU 插件化方案承接 PyTorch Native 生态以纯 FSDP 大 EP的极简并行切分实现内存最优与通信隐藏以 torch.compile AutoFuse 使能 Vector 算子自动融合以定制 Ascend C 算子释放稀疏注意力结构的芯片算力最终在 A3 64 卡集群上将 DeepSeek-V4-Flash 的 BF16 训练吞吐提升至 1100 tokens/p/s。该方案为现有昇腾 A3 集群用户基于 DeepSeek 新模型架构快速开展续训练/SFT 及自研算法验证提供了可直接复用的参考实现。【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考