如何配置 DeepSpeed Ulysses-Offload FPDT 训练 256K 长上下文 GPT？-开发者社区

如何配置 DeepSpeed Ulysses-Offload FPDT 训练 256K 长上下文 GPT？

【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed

如果你想在少数几块 GPU 上训练 256K token 上下文的 GPT 类语言模型，可以按本文在 Megatron-DeepSpeed 中启用 DeepSpeed Ulysses-Offload 的 FPDT（Fully Pipelined Distributed Transformer）训练方案。它把长序列切分成 chunk 并在 host 内存与 GPU 显存之间搬移，以降低长上下文训练的激活显存占用，同时保持原生 Transformer 的完整精度。官方教程 DeepSpeed Ulysses-Offload 的示例环境是 1 节点 4 块 NVIDIA A100 GPU，并要求安装 Flash Attention（CUDA）。

Ulysses-Offload 的工作方式

Ulysses-Offload 是构建在 ZeRO 与 DeepSpeed Ulysses 之上的长序列训练方案：

FPDT 分块流水线：把长序列输入切分成较小的 chunk，在 host 与 GPU 内存之间搬移，以双缓冲设计让 chunk 的取回/卸载与 attention 计算重叠；
chunk size 可配置：用户可以按预期显存预算设置 chunk 大小；
与 ZeRO 叠加：ZeRO 在 GPU 间切分模型和张量，与 FPDT 优化激活显存互补，进一步抬高长上下文训练的上限。

按 Ulysses-Offload 博客的说法，FPDT 使 8B 模型仅用 4 块 GPU 即可训练 2M 上下文，70B 模型用 32 块 GPU 训练 4M 上下文。

准备条件

教程建议先阅读 Getting Started、ZeRO 和 Megatron-DeepSpeed 三个教程，因为 Ulysses-Offload 的集成路径依赖 Megatron-DeepSpeed 代码仓库。

必须项：

Flash Attention（CUDA）是 FPDT 的硬性要求。按 DeepSpeed-Ulysses 教程的说明，FlashAttention 仅兼容 NVIDIA Turing、Ampere、Ada 或 Hopper GPU；该教程中还给出了 FlashAttention 与 Triton 的安装步骤，可在那里查看完整命令。
Megatron-DeepSpeed 代码：Ulysses-Offload 已完全集成进 Megatron-DeepSpeed，训练脚本基于它启动（启动方式参见 Megatron-DeepSpeed 教程）。

在训练脚本中设置 256K 上下文

教程示例把上下文大小放在脚本开头配置。按文档要求，seq_len必须是 2 的幂，示例使用 256K（262144）上下文和 mini batch 1：

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

添加 Ulysses-Offload 的 Megatron 参数

教程以 6.7B 参数 GPT 为例，先启用 ZeRO-3、Ulysses 序列并行和带 CPU offload 的 activation checkpointing，再加上以下三个 FPDT 参数：

ds_sequence_parallel_fpdt：布尔值，是否启用 FPDT，默认 false；
ds_sequence_parallel_fpdt_chunk_size：整数，FPDT 的 chunk size，默认 65536——无论序列多长，FPDT 都会以 65536 token 为一个 chunk 处理完整条序列；
ds_sequence_parallel_fpdt_offloading：布尔值，是否使用 host 内存卸载 chunk，默认 false。

三项功能全部开启的 megatron_options 片段如下：

megatron_options="\ --ds-sequence-parallel-fpdt \ --ds-sequence-parallel-fpdt-chunk-size 65536 \ --ds-sequence-parallel-fpdt-offloading \ --ds-sequence-parallel-size 4"

其中--ds-sequence-parallel-size 4指定序列并行度。按 DeepSpeed-Ulysses 教程的说明，模型的 attention head 数量必须能被该值整除。

FPDT 要求 Flash Attention，且支持 Rotary Position Embedding（RoPE），教程给出的对应参数为：

--use-flash-attn-v2 \ --use-rotary-position-embeddings \ --rotary-percent 0.25 \ --rotary-position-embeddings-theta 100000000

另外，教程通过 CPU checkpointing 进一步压缩激活显存：

if [ "${activation_checkpoint}" = "true" ]; then deepspeed_options="${deepspeed_options} \ --deepspeed-activation-checkpointing \ --checkpoint-in-cpu" fi

启动训练

完整的示例脚本位于 Megatron-DeepSpeed 仓库的examples_deepspeed/sequence_parallel/目录下，文件名为ds_pretrain_gpt_6.7B_fpdt_32k.sh，可直接对照其中megatron_options与deepspeed_options的组装方式修改自己的训练脚本。启动 Megatron-DeepSpeed 任务的方式，参见 Megatron-DeepSpeed 教程中的脚本运行说明。

配置核对与预期结果

训练启动前，用以下条件核对配置是否与文档一致：

seq_len为 2 的幂（本例 262144）；
--ds-sequence-parallel-size的值能整除模型的 attention head 数量；
已启用 Flash Attention（--use-flash-attn-v2），且 GPU 属于 NVIDIA Turing/Ampere/Ada/Hopper；
未同时启用 Megatron-LM 的 tensor 或 pipeline 并行——DeepSpeed-Ulysses 教程明确说明序列并行目前与这两种并行方式不兼容。

关于运行效果，Ulysses-Offload 博客的实验结果（文档示例数据，不是单次运行的固定预期）：在 GPT-2.7B 到 Llama-80B 的模型规模上，Ulysses-Offload 达到 55% 以上的 Model FLOPs Utilization（MFU）；8B 参数模型在 4 块 A100-40GB 上支持 2M token 序列长度，GPT-30B 用 16 块 GPU 支持 3M，Llama-70B 用 32 块 GPU 支持 4M。本文 256K + 6.7B GPT 的配置即在此能力范围内。

限制与下一步

FPDT 依赖 Flash Attention（CUDA），不支持不兼容 FlashAttention 的 GPU；
chunk size 可按显存预算调整（默认 65536），博客的实验章节给出了序列 chunk 与 MFU 之间权衡的进一步分析；
更完整的设计细节与实验对比，见 Ulysses-Offload 博客及其引用的 arXiv 技术报告；FPDT 的实现代码位于 deepspeed/sequence/fpdt_layer.py。

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考