DeepSpeed Ulysses-卸载

DeepSpeed Ulysses-Offload 是一个建立在 ZeRO 和 DeepSpeed Ulysses 之上的分块和卸载长上下文变压器模型训练方案系统。它采用了全流水线分布式变压器（FPDT），使得在仅使用4个GPU的情况下，可以在8B模型上进行2M上下文大小的训练，而在使用32个GPU的情况下，可以在70B模型上进行4M上下文大小的训练。阅读我们的 Ulysses-Offload 博客 和 论文 以了解更多信息！

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

Ulysses-Offload的设计

Ulysses-Offload 是一种基于分块和卸载的变压器实现，它保留了普通变压器的全部精度，同时显著减少了长上下文模型训练期间所需的激活内存。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个标记的块，直到整个序列全部处理完毕。
  • 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，并且还支持旋转位置嵌入 (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教程。

更新: