CANN/torchtitan-npu MTP特性

多Token预测特性(Multi Token Prediction, MTP)【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu在大规模语言模型的训练与推理优化中MTP 通过单次前向传播同时预测多个连续目标 Token大幅提升模型训练效率与数据利用率。传统单 Token 预测仅能学习逐词依赖关系MTP 则通过扩展预测长度、引入辅助损失函数显著加速模型收敛速度尤其在长文本、代码、多轮对话等任务上效果显著。 torchtitan_npu在deepseek_v32模型的基础上进一步适配了MTP训练特性实现了可配置长度的MTP训练同时支持FSDP2/EP/TP等分布式训练。实现原理参考Deepseek-V3的技术报告,我们在Deepseek_V32模型代码中引入了MTPModule类的定义其继承于原有的TransformerBlock_V32类并在此基础上新增了MTP模块所需的额外结构与参数。此外为最大化复用原有 Transformer 层成熟的分布式训练实现我们在模型顶层定义中将标准 Transformer Layer 与 MTP Layer 统一封装至model.layers列表中实现与原有FSDP2、EP、TP等分布式并行逻辑的无缝兼容。相关代码见torchtitan_npu/models/deepseek_v32/model/model.py为了实现 MTP 模块的有效训练我们新增了适配 MTP 模块的训练损失函数。具体而言每个MTPModule都会独立计算对应的交叉熵损失在此基础上模型总训练损失被定义为主损失与 MTP 损失的加权和。相关代码实现见torchtitan-npu/patches/torchtitan/loss.py配置选项在训练任务的 TOML 配置文件例如torchtitan_npu/models/deepseek_v32/train_configs/deepseek_v32_671b_debug.toml或实际启动训练时--job.config_file所指向的路径中找到对应的 [training] 节并添加以下配置以启用 MTP训练配置项类型默认值说明num_mtp_modulesint0(不开MTP)MTP预测的token个数mtp_loss_weightfloat0.3MTP训练损失的权重total_loss main_loss mtp_loss_weight* mtp_loss配置示例[training] local_batch_size 4 seq_len 2048 num_mtp_modules 1 mtp_loss_weight 0.3【免费下载链接】torchtitan-npuAscend Extension for torchtitan项目地址: https://gitcode.com/cann/torchtitan-npu创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考