XTunerV1–上海AILab开源的新一代大模型训练引擎

XTuner V1是什么

XTuner V1 是上海人工智能实验室开源的新一代大模型训练引擎,基于 PyTorch FSDP 开发,针对超大模型训练中的显存、通信和负载问题进行系统性优化,支持 1T 参数量级 MoE 模型训练,能在 200B 量级模型上实现训练吞吐超越传统 3D 并行方案。XTuner V1 与华为昇腾技术团队合作,进一步提升训练效率,实现更高的模型浮点运算利用率。XTuner V1 为学术界和工业界提供高性能、低门槛、易扩展的大模型训练方案。

XTuner V1

XTuner V1的主要功能

  • 高效训练超大模型:支持1T参数量级的MoE模型训练,在200B以上量级的混合专家模型中,训练吞吐超越传统3D并行训练方案。
  • 优化显存使用:通过自动Chunk Loss机制和Async Checkpointing Swap技术,有效降低显存峰值,无需借助序列并行技术实现200B参数量级MoE模型训练64K长度序列。
  • 掩盖通信开销:用极致的显存优化提升单次迭代的最大序列长度,增加每层计算耗时以掩盖参数聚合的通信耗时;通过Intra-Node Domino-EP降低参数聚合通信量,掩盖专家并行带来的额外通信开销。
  • 实现DP负载均衡:对每n个step内的已拼接好的序列进行排序,使每次计算时不同DP的最长子序列长度接近,减少因变长注意力导致的计算空泡。
  • 适配多种硬件:与华为昇腾技术团队合作,在Ascend A3 NPU超节点上进行深度优化,充分用超节点硬件特性,在理论算力落后NVIDIA H800近20%的情况下,能实现训练吞吐反超H800近5%,MFU反超20%以上。

XTuner V1的技术原理

  • 基于PyTorch FSDP开发:PyTorch FSDP(Fully Shard Data Parallel)是数据并行策略,将模型参数均匀切分到每张卡上,通过提前聚合参数和重新切分参数节省显存。XTuner V1在FSDP的基础上进行深度优化,解决其通信量大的问题。
  • 显存优化
    • 自动Chunk Loss机制:针对计算损失函数时的计算图,开发支持多种训练场景和多种硬件的自动Chunk Loss机制,有效降低显存峰值。
    • Async Checkpointing Swap:在模型前向计算开始时,将重计算需要保存的激活值从Device搬运到Host,释放显存;在反向传播时,提前将激活值从Host搬运回Device,反向传播结束后释放显存,进一步优化显存使用。
  • 通信掩盖
    • 增加计算耗时掩盖通信耗时:通过极致的显存优化,提升单次迭代的最大序列长度,增加每层计算的耗时,使计算耗时大于通信耗时,掩盖通信开销,避免计算空泡。
    • Intra-Node Domino-EP:针对因显存或通信带宽受限的训练场景,通过Intra-Node Domino-EP降低每一层聚合参数的通信量,掩盖因引入专家并行带来的额外通信开销。
  • DP负载均衡:大模型训练时,将多条句子拼接至一个固定长度,计算时使用变长注意力机制。XTuner V1通过对每n个step内的已拼接好的序列进行排序,使每次计算时不同DP的最长子序列长度接近,减少因变长注意力导致的计算空泡,提高训练效率。

XTuner V1的项目地址

  • 项目官网:https://xtuner.readthedocs.io/zh-cn/latest/
  • GitHub仓库:https://github.com/InternLM/xtuner

XTuner V1的应用场景

  • 自然语言处理(NLP)领域:用在训练超大规模的语言模型,如GPT、BERT等,提升模型的语言理解和生成能力,应用于机器翻译、文本生成、情感分析等任务。
  • 计算机视觉(CV)领域:在图像识别、目标检测等任务中,训练大规模的视觉模型,如ResNet、Transformer等,提高模型的准确性和效率。
  • 多模态学习:结合语言和视觉信息,训练多模态模型,如CLIP等,用在图像描述生成、视觉问答等任务,提升模型对复杂场景的理解能力。
  • 强化学习:在长序列的强化学习任务中,如机器人控制、游戏AI等,XTuner V1能够高效处理长序列数据,加速模型训练,提高策略学习的效率。

相关文章