DeepSpeed Ulysses-Offload

DeepSpeed Ulysses-Offload 是一个基于 ZeRO 和 DeepSpeed Ulysses 构建的，用于分块和卸载长上下文 Transformer 模型训练方案的系统。它采用了完全流水线分布式 Transformer (FPDT)，使得 8B 模型在仅 4 块 GPU 上即可训练 2M 上下文大小的模型，70B 模型在 32 块 GPU 上即可训练 4M 上下文大小的模型。阅读我们的 Ulysses-Offload 博客和论文了解更多信息！

我们建议您在开始本教程之前阅读 入门、ZeRO 和 Megatron-DeepSpeed 的教程。

Ulysses-Offload 的设计

Ulysses-Offload 是一种基于分块和卸载的 Transformer 实现，它保留了原始 Transformer 的全精度，同时显著减少了长上下文模型训练期间所需的激活内存。FPDT 将长序列输入分解为更小的块，并在主机内存和 GPU 内存之间移动它们，以实现卓越的内存效率，同时达到超过 50% 的 MFU。FPDT 采用双缓冲设计，将获取/卸载与注意力计算重叠。FPDT 还允许用户配置块大小以匹配预期的内存预算。

Ulysses-Offload 支持 ZeRO，它在 GPU 内存中分片模型和张量，进一步突破了长上下文模型训练在尖端硬件效率方面的限制。

训练环境

对于本教程，需要 Flash Attention (CUDA)。我们将使用 DeepSpeed Megatron-DeepSpeed 代码配置一个 80 亿参数的 LLaMA 模型。我们将使用 1 个节点，配备 4 块 NVIDIA Tesla A100-SXM4 Tensor Core GPU。

使用 Ulysses-Offload 训练 6.7B 参数的 GPT

用户可以在脚本开头设置上下文大小，对于此练习，我们将使用 256K 上下文和批处理大小为一的迷你批次。

### Main configs
seq_len=262144 # need to be power of 2

对于 6.7B 模型，我们将首先启用 ZeRO-3、Ulysses 和带 CPU 卸载的激活检查点，以达到可观的 GPU 内存效率，然后用户可以配置以下参数：

  • ds_sequence_parallel_fpdt: 布尔值，指示是否使用 FPDT，默认为 false。
  • ds_sequence_parallel_fpdt_chunk_size: 整数，指示 FPDT 中的块大小，默认为 65536，这意味着无论序列多长，FPDT 都将始终处理 65536 个 token 的块，直到整个序列处理完毕。
  • ds_sequence_parallel_fpdt_offloading: 布尔值，指示是否使用主机内存卸载块，默认为 false。

Megatron-DeepSpeed 配置更改

  1. 下面显示了一个启用所有 Ulysses-Offload 功能的 megatron-deepspeed 配置示例片段：
     megatron_options="\
 --ds-sequence-parallel-fpdt \
 --ds-sequence-parallel-fpdt-chunk-size 65536 \
 --ds-sequence-parallel-fpdt-offloading \
 --ds-sequence-parallel-size 4"
  2. FPDT 需要 Flash Attention，并支持 Rotary Position Embedding (RoPE)。
     --use-flash-attn-v2 \
 --use-rotary-position-embeddings \
 --rotary-percent 0.25 \
 --rotary-position-embeddings-theta 100000000 \
  3. 我们还启用了 CPU 检查点以减少激活内存占用。
     if [ "${activation_checkpoint}" = "true" ]; then
 deepspeed_options="${deepspeed_options} \
     --deepspeed-activation-checkpointing \
     --checkpoint-in-cpu"
 fi

您可以在此处找到完整的脚本。

有关如何启动 Megatron-DeepSpeed 任务的更多详细信息，请参阅 Megatron-DeepSpeed 教程示例。

恭喜！您已完成 Ulysses-Offload 教程。

更新日期：