在生成式 AI 领域，长文档摘要和视频生成等任务都需要模型具有超长上下文的能力。
如何训练超长上下文的模型，既是生成式 AI 算法领域的研究热点，也是 AI Infra 领域的难点。
随着 AI 模型参数量的不断增大，为了能够训练超长上下文，通常需要使用一些复杂的并行策略，如 Nvidia Megatron, DeepSpeed Ulysses 等工作。这些工作虽然解决了超长上下文的训练问题，但需要开发者具有一定的 AI Infra 的知识，对生成式 AI 的研究人员很不友好。
为了让研究人员能够更加便捷地训练超长上下文模型，促进生成式 AI 领域的发展，XTuner 开发了一套超长上下文训练解决方案：
- 支持全量训练超过百万个 tokens 的序列
- 支持 "千亿级参数" 模型，通过序列并行，扩展 ZeRO3 能力边界
- 开箱即用：可直接训练 HuggingFace模型
并行策略简介
尽管开源模型支持的序列长度不断被刷新，但主流的显存优化策略（如 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+ 模型的训练效率
最佳实践
环境准备
-
创建虚拟环境
iniconda create --name xtuner-env python=3.10 -y conda activate xtuner-env -
XTuner 安装
bash# 方案一，PyPI安装（推荐） pip install -U 'xtuner[deepspeed]' # 方案二，源码安装 git clone https://github.com/InternLM/xtuner.git cd xtuner pip install -e '.[deepspeed]' -
flash_attn 安装
使用序列并行策略需要安装 flash attn 库（参考 flash attn 安装 ，安装过程需要 cuda）。
修改配置文件
可以通过运行以下命令查看 XTuner 提供的训练不同模型的配置文件：
bash
xtuner list-cfg
xtuner copy-cfg ${TARGET_CONFIG} ./修改被拷贝的 config 文件中sequence_parallel_size即可使用序列并行策略：
ini
# 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 进行训练：
scss
(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 策略
相关链接
- XTuner 序列并行训练文档 xtuner.readthedocs.io/zh-cn/docs/...