新一代超大模型训练引擎XTuner V1开源

9月8日，上海人工智能实验室(上海AI实验室)开源书生大模型新一代训练引擎XTuner V1。

XTuner V1是伴随上海AI实验室“通专融合”技术路线的持续演进，以及书生大模型研发实践而成长起来的新一代训练引擎。相较于传统的3D并行训练引擎，XTuner V1不仅能应对更加复杂的训练场景，还具备更快的训练速度，尤其在超大规模稀疏混合专家(MoE，mixture of experts)模型训练中优势显著。

同时，为了进一步探究XTuner V1训练方案的上限，研究团队与昇腾团队在昇腾384超节点(Atlas 900 A3 SuperPoD)上进行联合优化，充分利用昇腾384超节点硬件特性，实现更高的模型算力利用率(MFU，model FLOPS utilization)。相比业界其他产品，昇腾384超节点的训练吞吐提升5%以上，MFU提升20%以上，该项研究成果技术报告也将于近期发布。

除了训练框架，书生大模型研发中使用的AIOps工具DeepTrace与ClusterX也将一并开源，为大规模分布式训练提供全方位保障。

XTuner V1:

https://github.com/InternLM/xtuner

DeepTrace:

https://github.com/DeepLink-org/DeepTrace

ClusterX:

https://github.com/InternLM/clusterx

目前开源社区主流的训练方案主要分为两类：

DeepSpeed/PyTorch FSDP(完全分片数据并行，Fully Shard Data Parallel)：通信量大但使用简单，尤其适合稠密型模型训练，开发者无需具备专业的AI Infra知识，也能开发出接近最优性能的训练系统;

3D并行：通信量小但使用复杂，开发者需要具备专业的AI Infra知识，针对不同硬件和训练场景进行针对性调优，尤其适用MoE模型训练。

为了同时兼顾易用性、高性能与超大模型训练，XTuner V1基于PyTorch FSDP进行开发，并针对FSDP通信量大的固有缺陷，进行了系列优化，可支持1T参数量级MoE模型训练，并首次在200B以上量级的混合专家模型上，实现训练吞吐超越传统的3D并行训练方案。

针对当前主流的MoE后训练需求，XTuner V1未通过序列并行方式，实现200B量级MoE模型单次forward-backward可处理64k序列长度，更适合当下流行的强化学习训练场景;对专家并行依赖小，长序列训练时受专家不均衡影响小，200B量级MoE无需专家并行，600B MoE只需节点内专家并行，更适合现代MoE Dropless训练模式;大规模长短序列混训场景提速2倍以上，数据并行负载均衡，大幅减小因需序列长度不均衡导致的计算空泡。

多维度技术优化

专为“超大模型”而生

XTuner V1之所以能在超大模型训练中展现出卓越的性能，核心在于它在显存、通信、负载等多个维度进行了系统性优化。这些优化协同作用，不仅带来了性能的跨越式提升，还兼顾了易用性、通用性与扩展性。

显存优化：Pytorch FSDP与3D并行最大的差异在于重计算。针对计算损失函数时的计算图，XTuner V1基于Liger-Kernel中的Chunk-wise Loss，扩展支持了更多种类的损失函数，能够支持昇腾NPU;针对重计算保留的激活值，XTuner V1借鉴了昇腾MindSpeed中的Async Checkpointing Swap。最终，无需借助序列并行技术，实现200B参数量级MoE模型训练64K长度序列。

通信掩盖：得益于极致的显存优化，XTuner V1可以让单次迭代的最大序列长度提升数倍，从而增加每层计算的耗时，掩盖参数聚合的通信耗时。针对因显存或通信带宽受限，无法实现通信掩盖的训练场景，XTuner V1通过Intra-Node Domino-EP来降低每一层聚合参数的通信量，同时掩盖因引入专家并行带来的额外通信开销。

DP负载均衡：由于XTuner V1中没有引入TP、PP等并行策略，相同卡数下，数据并行的维度会远大于3D并行。为了缓解变长注意力带来的计算空泡，并尽可能不影响数据的训练顺序，会对每n个step内的已拼接好的序列进行排序，让每次计算时，不同DP的最长子序列长度是接近的。

基于昇腾384超节点深度优化

训练效率超业界产品

为了进一步探究XTuner V1训练方案的上限，上海人工智能实验室XTuner团队联合昇腾技术团队在超节点上进行深度优化，充分利用昇腾384超节点硬件特性，实现了更高MFU。

昇腾384超节点通过高速总线连接多颗NPU，突破互联瓶颈，让超节点像一台计算机一样工作，更加适合FSDP训练：

更高的通信带宽：最大可实现384颗NPU点到点超大带宽互联，FSDP All Gather耗时仅为业界产品的1/4~1/3，更容易实现计算-通信掩盖

计算通信解耦：通过专用硬化调度和传输卸载，实现不占用计算核的高效数据通信，FSDP计算通信掩盖时不会影响计算速度

超节点高速互连：CPU和NPU通过高速总线互联，实现更大带宽，Checkpointing Swap的开销更小

除硬件固有优势外，昇腾还从通信、内存、计算、框架、工具等维度对基于昇腾384超节点的MoE训练进行了全方位的加持：

Cube调优：对于模型中集中了大量计算任务的GroupedMatmul算子进行分析，发现内部搬运带宽已经拥塞但cube利用率还有提升空间。针对此问题，联合研发团队重点优化GroupedMatmul算子分块逻辑，根据不同输入进行动态分块Tiling策略优化搬运效率。同时，根据场景的不同细化Cache策略，提高Cache命中率从而提升性能。

QoS调优：QoS(Quality of Service)即服务质量。在有限的带宽资源下，QoS为各种业务分配带宽，为业务提供端到端的服务质量保证。大规模训练过程中，计算流、通信流、swap流都会存在内存访问，并发的访问会导致内存带宽拥塞，从而影响整体性能。通过适当调低通信的内存访存优先级，可以减少计算的搬运时间，从而优化端到端性能。

跨流内存复用：在FSDP计算流和通信流异步重叠的场景中，Ascend Extension for PyTorch(PTA)中默认的跨流内存优化会导致显存不能及时释放，需要开启PTA中进阶版的跨流内存复用机制(MULTI_STREAM_MEMORY_REUSE=2)，可以显著降低显存峰值。

集群性能工具高效诊断：借助MindStudio全流程工具链中的msprof-analyze性能分析工具与MindStudio Insight可视化工具，开发者可以充分利用其强大的数据分析与可视化能力，在分钟级时间内精准识别训练过程中的“快慢卡”现象根因，快速定位出性能瓶颈，显著提升大集群调优效率。

书生大模型工具链研发团队现已将Xtuner V1的工作全部开源，希望为学术界与工业界提供高性能、低门槛、易扩展的大模型训练方案，丰富开源社区的训练工具生态，为超大模型研发和应用提供坚实易用的基础设施。

未来，在研究范式创新及模型能力提升的基础上，上海AI实验室将持续推进书生大模型及其全链条工具体系的开源，支持免费商用，同时提供线上开放服务，与各界共同拥抱更广阔的开源生态，共促大模型产业繁荣。