CANN/torchtitan-npu融合算子适配

NPU 融合算子适配【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-nputorchtitan_npu 在torchtitan_npu/converters/kernels下定义了多个 torchtitan ModelConverter 。在启动模型训练任务时它们会根据用户配置自动将模型中的原始模块替换为基于 NPU 融合算子的实现从而实现模型在 NPU 平台上的亲和适配。如何配置所有融合算子均通过在训练配置 TOML 文件 例如torchtitan_npu/models/deepseek_v32/train_configs/deepseek_v32_671b_debug.toml或实际启动训练时--job.config_file所指向的路径的[model]节中的converters列表中添加对应的配置项来启用。 示例[model] name deepseek_v32 flavor debugmodel hf_assets_path ./assets/hf/DeepSeek-V3.2 converters [npu_dsa, npu_rms_norm, npu_permute, npu_gmm]当前版本支持以下 ModelConverters 前往对应章节查看功能介绍及启用方式DSAGMMPermuteRMSNormRope关于本仓库适配的各融合算子的详细说明请查看对应的 NPU 融合算子开发者文档。DSA (DeepSeek Sparse Attention)DSA 是DeepSeek-V3.2中引入的一种特殊注意力机制主要由图中的两个模块构成Lightning Indexer筛选出少量高价值token的索引这些索引被用于高效的稀疏 Attention 计算(图中 Multi-Query Attention部分)。针对 DeepSeek V3.2 模型的 Attention 模块将以上两种核心组件替换为对应的 NPU 融合算子。具体对应关系如下DeepSeek V3.2 Attention 组件NPU融合算子Lightning Indexer 前向计算npu_lightning_indexerLightning Indexer 反向计算梯度 Lossnpu_sparse_lightning_indexer_grad_kl_loss稀疏注意力计算npu_sparse_flash_attention配置示例[model] converters [... npu_dsa, ...] # 添加 npu_dsa 配置项ModelConverter 源码路径torchtitan_npu/converters/kernels/dsa.py相关 NPU 融合算子开发者文档npu_lightning_indexernpu_sparse_lightning_indexer_grad_kl_lossnpu_sparse_flash_attentionGMMGrouped MatMul在 MoE 模块中每个专家执行前馈网络FFN运算如 Swiglu FFN输入先经过升维变换w1再通过激活函数最后经过降维变换w2得到输出。由于各专家执行结构相同的矩阵乘法为了将同类矩阵运算合并为一次算子调用提升计算效率本ModelConverter 引入分组矩阵乘法GMM算子npu_grouped_matmul。该算子接收 Permute 模块输出的重排后 token 及对应的专家索引在一次调用中并行计算所有专家的同一线性层如所有专家的w1。配置示例[model] converters [... npu_gmm, ...] # 添加 npu_gmm 配置项ModelConverter 源码路径torchtitan_npu/converters/kernels/gmm.py相关 NPU 融合算子开发者文档npu_grouped_matmulPermuteMoE 前向计算中为了利用 GMM 提升计算效率token 需要根据 MoE Router 为每个 token 分配的专家以特定顺序重排输出重排列后的 token 及其对应的专家索引计算完成后再将结果恢复至原始 token 顺序。本 ModelConverter 将“重排”和“恢复”操作替换为基于npu_moe_token_permute和npu_moe_token_unpermute算子的实现。配置示例[model] converters [... npu_permute, ...] # 添加 npu_permute 配置项ModelConverter 源码路径torchtitan_npu/converters/kernels/permute.pyRMSNormRMSNorm 通过计算输入张量每个样本的平方均值的平方根来稳定深层网络的训练。本 ModelConverter 将模型中的 RmsNorm 操作替换为基于npu_rms_norm融合算子的实现。TOML 配置项npu_rms_norm配置示例[model] converters [... npu_rms_norm, ...] # 添加 npu_rms_norm 配置项ModelConverter 源码路径torchtitan_npu/converters/kernels/rms_norm.py相关 NPU 融合算子开发者文档npu_grouped_matmulRoPERoPE 将 token 位置相关的旋转矩阵应用于自注意力机制中的 Query 和 Key 向量使每对 token 之间的相对位置信息在 Attention 计算中自然包含 Query 和 Key 的乘积。在模型实现中通常预先生成每个位置的旋转角度在 Attention 计算时即时对 Query 和 Key 进行旋转变换。本 ModelConverter 将这一旋转变换操作替换为基于npu_rotary_mul融合算子的实现。配置示例[model] converters [... npu_rope, ...] # 添加 npu_rope 配置项ModelConverter 源码路径torchtitan_npu/converters/kernels/rope.py相关 NPU 融合算子开发者文档npu_rope【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考