千亿参数，百万序列 | 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 只能优化模型参数和优化器状态占用的显存，超长训列训练过程中的显存开销主要来自激活值，而非模型参数和优化器状态。

图1 不同序列长度时，使用 ZeRO-3 训练 128k 上下文 yi-34B 模型的训练情况

为解决长序列训练过程中的显存问题，Megatron-LM 团队和 DeepSpeed 团队分别提出了两种序列并行算法，通过对长序列进行切分的方法来降低单 GPU 上计算的序列长度，如图 2 所示。

图 2 使用 8 块 GPU 进行分布式训练，序列并行度为 4 时的训练架构示意图

XTuner 中的序列并行设计思路参考了 DeepSpeed 的工作 DeepSpeed Ulysses，并加以优化，以实现一键开启序列并行策略。三者的对比如下：

	Attention 通信量 （序列长度 N，并行度 P）	代码侵入
Megatron-LM	O(N)	高，不支持 HuggingFace 模型直接训练
DeepSpeed Ulysses	O(N / P)	较高，需要改动 HuggingFace 部分源码
XTuner	O(N / P)	无，HuggingFace 模型直接启动训练

图 3 展示了在 RedPajama-Data-1T-Sample 数据集上训练 32k 上下文的 Llama2-7B ，在序列并行度设为 1 和 2 时的收敛性。XTuner 的序列并行是一种纯系统优化技术，用于实现超长序列的 Transformer 模型训练，而对于模型的训练精度上没有负面影响。

图 3 不同序列并行度的收敛性评估

Benchmark

超长序列

我们首先对比了 LLaMA-Factory 和 XTuner 在训练不同上下文长度的 Llama2-7B 模型时的训练效率。从图 4 的对比结果中可以看出：

• 由于集成了多个训练加速算子，XTuner 在序列长度较短时具有一定的速度优势

• 当序列长度逐渐变长时，Decoder Layer 中的激活值可能会占用大量的显存资源，进而导致 GPU out of memory。得益于序列并行策略，XTuner 可以将上下文长度提升至 1M (1048576) tokens。

• 由于与 Megatron-LM 等其他序列并行策略相比， XTuner 的序列并行策略显著减少训练中的通信量，因此随着序列长度增长，XTuner 的计算吞吐量没有明显下滑。

图 4 XTuner 与 LLaMA-Factory 在 Llama2-7B 模型上的训练效率对比

超长序列 + 超大模型

由于尺寸更大的模型在训练过程中模型状态会占用大量显存，因此大模型的训练对显存优化具有更高的要求。XTuner 支持将 Llama2-70B 模型的上下文长度扩展至一百万，同时计算吞吐量没有严重的下滑。

图 5 展示了 LLaMA-Factory 和 XTuner 在训练不同上下文长度的 Yi-34B 和 Llama2-70B 模型时的训练效率。从图中可知，序列并行策略对于大模型长序列序列是必要的。

图 5 XTuner 与 LLaMA-Factory 在 Yi-34B 和 Llama2-70B 模型上的训练效率对比

另外，XTuner 还快速支持了 Chatbot Arena 榜单上超越 GPT-4 的 Command R+ （千亿参数量）模型（千亿参数量）, 128k 长上下文训练效率如图 6 所示。

图 6 XTuner 训练 128k 上下文的 Command R+ 模型的训练效率

最佳实践

环境准备

1. 创建虚拟环境

   conda create --name xtuner-env python=3.10 -y
   conda activate xtuner-env

2. XTuner 安装

   # 方案一，PyPI安装（推荐）
   pip install -U 'xtuner[deepspeed]'
   # 方案二，源码安装
   git clone https://github.com/InternLM/xtuner.git
   cd xtuner
   pip install -e '.[deepspeed]'

3. flash_attn 安装

使用序列并行策略需要安装 flash attn 库（参考 flash attn 安装 ，安装过程需要 cuda）。

修改配置文件

可以通过运行以下命令查看 XTuner 提供的训练不同模型的配置文件：

xtuner list-cfg
xtuner copy-cfg ${TARGET_CONFIG} ./

修改被拷贝的 config 文件中sequence_parallel_size即可使用序列并行策略：

# parallel
- sequence_parallel_size = 1
+ sequence_parallel_size = 4  # take `sequence_parallel_size = 4`` as an example

• 需要保证所使用的 GPU 总数以及注意力头数 (num_attention_heads 而非 num_key_value_heads) 可以被 sequence_parallel_size 整除。

开始训练

需要使用 DeepSpeed 进行训练：

(DIST) NPROC_PER_NODE=${GPU_NUM} xtuner train ${CONFIG_PATH} --deepspeed deepspeed_zero2
(SLURM) srun ${SRUN_ARGS} xtuner train ${CONFIG_PATH} --launcher slurm --deepspeed deepspeed_zero2

• ${CONFIG_PATH} 为 Step 1 中修改得到的 config 文件路径

• 可根据实际情况选择使用不同的 zero 策略

相关链接

1. XTuner 序列并行训练文档 https://xtuner.readthedocs.io/zh-cn/docs/acceleration/train_extreme_long_sequence.html