大模型正走向图文融合与更长上下文，复杂结构也让训练适配和多卡协同更具挑战。TorchTitan是基于PyTorch原生技术的训练框架，以模块化设计灵活组合模型与并行策略。

为了让开发者在昇腾上沿用这一体验，我们推出TorchTitan Ascend，复用上游训练流程，接入融合算子与MoonEP，支持Kimi K3.1 Pro训练，减少重复适配工作，并在Atlas 950 SuperPoD上完成减层训练验证。项目首次开源，欢迎开发者关注与贡献。

Kimi K3.1 Pro 模型结构

Kimi K3.1 Pro 采用原生多模态混合专家（MoE）架构，模型总参数约 2.8T，每个 Token 激活约 104B 参数，语言主干包含 93 层，模型上下文规格达到百万 Token。它延续 K3 的主体设计，将原有的 Gated MLA 全注意力层替换为 Gated DSA 稀疏注意力层，以 KDA 的状态递推配合 DSA 的关键位置选择，组织长上下文计算。

主干由 69 个 KDA 层和 24 个 Gated DSA 层组成：前 92 层按“3 个 KDA 层加 1 个 Gated DSA 层”重复 23 次，最后再接 1 个 Gated DSA 层。

多模态输入。 文本经过 Embedding 转换为主干表示；图像由约 401M 参数的 MoonViT-V2 视觉编码器处理，再经多模态投影模块映射到主干维度，与文本表示一同进入语言模型。

KDA 通过状态递推积累历史信息。通过门控矩阵状态递推积累历史信息，以状态读写替代对全部历史 KV 的访问，支持长序列建模。

Gated DSA 按需访问关键上下文。基于 MLA 潜在 KV，引入 Indexer 为 Query 选择最多 2048 个相关位置，再执行稀疏注意力，并通过输出门控融合结果。

LatentMoE 以稀疏激活组织专家计算。配置 896 个路由专家和 2 个共享专家，每个 Token 激活 16 个路由专家，通过低维潜空间计算扩展模型容量。

AttnRes 沿深度方向选择和融合信息。通过可学习注意力权重聚合不同深度的历史层表示，使当前计算能够选择性利用跨层信息。

图1 Kimi K3.1 Pro 整体结构。下方单独展示一次 AttnRes 融合，注意力前与前馈网络前分别计算。图中的视频关键帧表示视觉输入方式；本文介绍的昇腾适配范围为图文训练。

Kimi K3.1 Pro 昇腾训练适配

围绕 Kimi K3.1 Pro，适配与优化主要覆盖以下特性：

• 模型结构：适配 KDA 与 Gated DSA 混合注意力，复用 AttnRes、LatentMoE 等模块。

• 图文数据：衔接视觉编码、图文输入与序列打包（packing），传递文档边界和填充掩码。

• 训练流程：接入模型权重映射、Indexer 辅助损失及归一化处理，沿用统一训练器。

• 融合算子：覆盖 KDA、DSA、LIG-KL 和 AttnRes 等关键计算。

• 分布式并行：支持 FSDP2/HSDP、DDP 数据并行，以及张量与序列并行（TP/SP）、流水线并行（PP）、上下文并行（CP）和专家并行（EP），可按训练配方组合使用。

复用K3主干接入K3.1

TorchTitan Ascend 基于上游 K3 配置派生 K3.1，将 Gated MLA 层替换为 Gated DSA，继续复用 KDA、AttnRes、LatentMoE、图文输入处理和训练器。模型差异集中在 DSA 实现、训练接口、权重映射及对应训练配方中，使架构升级能够沿用既有训练流程。

DSA 适配为各层接入独立的 Indexer、稀疏注意力计算和辅助损失，并在训练接口中传递序列打包（packing）的边界信息、填充掩码与辅助损失的归一化分母。权重映射在继承 K3 规则的基础上，补充 Indexer 相关权重的映射，让新增结构与模型加载、数据处理和训练计算相衔接。

融合算子优化覆盖训练关键环节

混合架构减少了部分计算量，也带来了更复杂的数据依赖。围绕这些开销，TorchTitan Ascend 结合昇腾生态中的 Ascend C 与 Triton 算子实现，将优化接入注意力、索引器辅助损失和跨层聚合等训练环节，减少算子启动、中间张量读写与重复计算。

KDA前后处理融合：减少长序列训练的数据搬运

KDA 将长序列计算组织为块内处理与块间状态传递，归一化、门控和状态更新紧密关联，分散执行容易增加显存读写与调度开销。

针对这一特点，昇腾将 Q/K 归一化、β 参数变换及门控处理融入前向流程，使预处理结果直接衔接分块计算，减少中间数据搬运；反向复用前向保存的归一化统计、块内矩阵和状态，衔接状态梯度传播与 Q/K 梯度计算，减少重复归约和状态重算，降低前后处理成本。

DSA稀疏索引与注意力融合：降低数据搬运与显存占用

DSA通过选择关键 KV 减少计算，但不同 Query 的访问位置各异，索引筛选、Gather 与注意力计算之间的数据流转会削弱稀疏收益。

昇腾通过 Lightning Indexer（LI）在算子内部完成评分与 Top-K 选择，减少完整评分矩阵的中间处理；SFA／SFAG 则直接依据索引访问目标 KV，完成注意力前反向计算，避免预先展开庞大的 KV 张量，降低数据搬运与显存占用。

LIG-KL 损失与梯度融合：降低索引器辅助训练开销

DSA 索引器通过 KL 损失学习教师侧注意力分布，长序列下，教师信息展开与评分计算图保存会增加显存压力。

LIG-KL 融合算子复用 SFA 的归一化统计、压缩 Q/K 及稀疏索引，减少教师分布的显式展开；同时融合辅助损失与索引器梯度计算，使反向直接使用所需梯度，减少中间计算图保存与重复计算，在保持学习目标的同时降低辅助训练成本。

AttnRes 在线跨层聚合：减少历史特征的重复访存

AttnRes 动态加权不同深度的历史特征，分步执行归一化、评分和求和会反复读取特征、写回中间结果。

昇腾通过 Triton 融合归一化与评分，利用在线 Softmax 在同一次遍历中更新统计量和聚合结果，减少特征重读与中间存储；反向继续复用前向评分和归一化统计，进一步降低跨层聚合的重算开销，使跨层特征选择能够以更紧凑的数据流执行。

MoonEP动态负载均衡与专家并行优化

多卡MoE训练中，动态路由可能将大量 Token 集中到少数专家，造成热点设备繁忙、其他设备等待，整体速度受最繁忙设备限制。

TorchTitan Ascend 接入 MoonEP，通过在线规划与动态专家副本，将热点任务分配到更多设备。在保持 Token 所选专家不变的前提下，调整计算执行位置，让空闲算力参与热点任务，并结合权重预取与梯度归约，将负载均衡贯穿前反向训练。

图2 MoonEP 负载均衡原理示意。示例中，32 份任务由不均衡分布调整为 8 个 Rank 各承担 4 份；图中任务数用于说明调度机制，不代表实测性能。

在线负载规划：及时分担批次内的突发热点

即使路由在较长时间内整体趋于均衡，单个批次仍可能出现突发热点。MoonEP 根据当前路由结果统计各专家的任务量，结合设备剩余容量，在线确定热点专家的副本及计算分配方案。

规划完成后，Token 被分发给本地专家或对应的副本专家，让同一专家在多卡上处理不同 Token，保持专家身份不变，使空闲设备能够及时参与热点计算，减少单卡积压造成的整体等待。

图3 MoonEP 在线规划示意。根据当前负载识别热点与剩余容量，安排动态副本，再按计划分发任务。

权重预取与梯度归约：保障动态副本的完整训练

任务分配到其他设备后，目标设备既要获得相应专家的权重，也要将计算产生的梯度纳入统一更新。MoonEP 根据调度计划，在专家计算前将所需权重预取到目标设备的缓冲区，并通过统一计算视图，由 Grouped GEMM 组织本地与副本专家计算，为任务分担提供所需数据。

反向阶段，副本权重梯度写入归约缓冲，再汇总回专家所属设备，与本地梯度合并后统一更新。由此，专家可在多卡执行，参数仍统一管理，使负载均衡覆盖完整训练过程。

图4 MoonEP 权重与梯度流转示意。前向将权重预取到副本设备，反向将副本梯度归约回专家所属设备，汇总后统一更新。

昇腾通信适配：减少动态调度的传输与等待开销

动态专家计算需要协调 Token 传输、权重预取和梯度归约，阶段间的等待会影响负载均衡的实际收益。昇腾适配依托 MTE／UDMA 硬件数据传输能力，按设备互联条件与任务特点匹配传输路径和调度策略，减少数据搬运开销。

同步与异步完成机制为这些阶段提供依赖管理，使框架在数据就绪条件满足时安排其他任务，为通信计算重叠创造空间，降低额外等待。

Kimi K3.1 Pro 训练启动与验证

本示例在Atlas 950 SuperPoD上使用2个节点、共16张 NPU，以FSDP2并行策略启动训练。所用 kimi_k3_1_reduced_npu 配方保留全尺寸宽度，采用8层、64个路由专家，并包含视觉编码器。

两节点 16 卡启动

参考 TorchTitan Ascend 安装说明 准备依赖、算子、tokenizer 和训练数据，并加载环境变量。两台节点使用相同代码与训练配置，确认环境中包含 kimi_k3_1_reduced_npu 配方，并将该配方的 training.steps 设为 100。

在两台节点的 TorchTitan 源码根目录分别执行以下命令。将 RANK0_IP 替换为 Rank 0 节点的实际 IP；NODE_RANK 在两个节点分别设为 0 和 1，MASTER_IP 与端口保持一致。

MASTER_IP="RANK0_IP"NODE_RANK=0torchrun \--nnodes=2 \--nproc_per_node=8 \--node_rank="${NODE_RANK}" \--master_addr="${MASTER_IP}" \--master_port=29500 \-m torchtitan_ascend.train \--module torchtitan_ascend.recipes.kimi_k3_1 \--config kimi_k3_1_reduced_npu

详细指导见：

https://gitcode.com/Ascend/torchtitan-ascend/blob/main/docs/models/kimi_k3.md

100步训练结果

本次随机初始化实验完成 100 步训练，全局平均 Loss 从第 1 步的 12.3781 降至第 100 步的 11.4519，整体呈下降趋势。下图展示了该配置的短程训练结果。

图5 Kimi K3.1 Pro 示例配置的 100 步训练 Loss 曲线。实验使用 2 节点、16 卡、FSDP2，模型从随机初始化开始。

项目与社区

TorchTitan Ascend属于Ascend for PyTorch Distributed SIG，欢迎开发者参与模型适配、算子优化与训练验证，共建昇腾大模型训练生态。

TorchTitan Ascend 仓库：

https://gitcode.com/Ascend/torchtitan-ascend

TorchTitan 上游仓库：https://github.com/pytorch/torchtitan

Ascend for PyTorch Distributed SIG：https://gitcode.com/Ascend/community/blob/master/AscendForPyTorch/sigs/distributed/README.md

Ascend for PyTorch TC 与 SIG 组织：https://gitcode.com/Ascend/community/blob/master/AscendForPyTorch/README.md

SIG 会议日历：https://meeting.ascend.osinfra.cn/?sig=sig-distributed