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

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

最佳实践

环境准备

  1. 创建虚拟环境

    ini 复制代码
    conda create --name xtuner-env python=3.10 -y
    conda activate xtuner-env
  2. XTuner 安装

    bash 复制代码
    # 方案一,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 提供的训练不同模型的配置文件:

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 策略

相关链接

  1. XTuner 序列并行训练文档 xtuner.readthedocs.io/zh-cn/docs/...
相关推荐
视觉语言导航4 分钟前
复杂地形越野机器人导航新突破!VERTIFORMER:数据高效多任务Transformer助力越野机器人移动导航
人工智能·深度学习·机器人·transformer·具身智能
kebijuelun5 分钟前
OpenVLA:大语言模型用于机器人操控的经典开源作品
人工智能·语言模型·机器人
掘金安东尼13 分钟前
大模型+Python脚本,打造属于你的“批量生成文档”应用!
人工智能
vocal15 分钟前
谷歌第七版Prompt Engineering—第二部分
人工智能·后端
Blossom.11824 分钟前
量子计算在密码学中的应用与挑战:重塑信息安全的未来
人工智能·深度学习·物联网·算法·密码学·量子计算·量子安全
子燕若水25 分钟前
How do I install OpenCV with gpu support
人工智能·opencv·计算机视觉
明明跟你说过35 分钟前
深度学习常见框架:TensorFlow 与 PyTorch 简介与对比
人工智能·pytorch·python·深度学习·自然语言处理·tensorflow
搏博37 分钟前
专家系统的基本概念解析——基于《人工智能原理与方法》的深度拓展
人工智能·python·深度学习·算法·机器学习·概率论
yzx99101337 分钟前
决策树随机深林
人工智能·python·算法·决策树·机器学习
巷95540 分钟前
卷积神经网络迁移学习:原理与实践指南
人工智能·cnn·迁移学习