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

InternLM

已于 2024-04-16 16:54:07 修改

阅读量912

点赞数 16

文章标签：人工智能 ai 语言模型

于 2024-04-16 15:54:35 首次发布

本文链接：https://blog.csdn.net/2402_82411485/article/details/137828282

版权

在生成式 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+ 模型的训练效率

最佳实践

环境准备

创建虚拟环境

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

XTuner 安装

# 方案一，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 提供的训练不同模型的配置文件：

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