千亿参数，百万序列 | XTuner 超长上下文训练方案

在生成式 AI 领域，长文档摘要和视频生成等任务都需要模型具有超长上下文的能力。

如何训练超长上下文的模型，既是生成式 AI 算法领域的研究热点，也是 AI Infra 领域的难点。

随着 AI 模型参数量的不断增大，为了能够训练超长上下文，通常需要使用一些复杂的并行策略，如 Nvidia Megatron, DeepSpeed Ulysses 等工作。这些工作虽然解决了超长上下文的训练问题，但需要开发者具有一定的 AI Infra 的知识，对生成式 AI 的研究人员很不友好。

为了让研究人员能够更加便捷地训练超长上下文模型，促进生成式 AI 领域的发展，XTuner 开发了一套超长上下文训练解决方案：

• 支持全量训练超过百万个 tokens 的序列

• 支持"千亿级参数" 模型，通过序列并行，扩展 ZeRO3 能力边界

• 开箱即用：可直接训练 HuggingFace模型

XTuner 项目链接：

https://github.com/InternLM/xtuner

（觉得好用欢迎点亮小星星）

并行策略简介

尽管开源模型支持的序列长度不断被刷新，但主流的显存优化策略（如 ZeRO 系列）却不足以解决大模型、长序列训练问题。如图 1 所示，使用 ZeRO-3 显存优化策略训练超长序列时，单纯增加 GPU 数量无法解决超长序列带来的 OOM 问题；这是因为，ZeRO-3 只能优化模型参数和优化器状态占用的显存，超长训列训练过程中的显存开销主要来自激活值，而非模型参数和优化器状态。