企业官网建设流程全解析

DeepSpeed-Ulysses 深度解析：用 All-to-All 序列并行撑起百万 Token 超长序列训练

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

DeepSpeed-Ulysses 是 DeepSpeed 推出的一种**以注意力计算为中心、通过 All-to-All 通信实现序列并行（Sequence Parallelism）**的训练方案，目标是让 Transformer 大模型在极长序列（最长已实测到100 万 Token级别）下依然高效、可扩展。读完本文，你将掌握 Ulysses 的核心切分/通信设计、通信量 O(N/P) 的推导依据、与 ZeRO-3 / FlashAttention 的协同方式，以及如何在 Megatron-DeepSpeed 和 HF Transformers 代码中把原有自注意力替换成DistributedAttention从而开启超长序列训练。

本文的主体内容源自仓库 docs/_posts/2023-08-24-ulysses.md 所指向的 blogs/deepspeed-ulysses/README.md（同主题另有 中文版 与 日文版），并结合 deepspeed/sequence/layer.py、docs/_tutorials/ds-sequence.md 等仓库源码与教程做纵深印证。

为什么超长序列训练会成为系统瓶颈

从生成式 AI 到 AI for Science，超长上下文的需求正在全面爆发。会话式 AI、长文档摘要、视频生成需要在空间与时间维度上对长上下文进行推理；多模态基础模型要同时处理语音、图像与波形等高维输入。基因序列语言模型需要在 64 亿个碱基字母上学习基因组进化规律，基于完整病历的诊断预测模型同样依赖超长上下文。

然而，经典的数据并行、张量并行、流水并行都存在两个致命短板：其一，它们都无法沿序列维度扩展；其二，既有的序列并行方案（如 Megatron-LM）存在显著的内存-通信效率问题，且往往要求侵入式、易出错的重构代码。DeepSpeed-Ulysses（名字取自同名超长篇小说Ulysses）正是为突破这两个短板而设计的：一个简单、可移植、高效的超长序列 LLM 训练方法论。

核心设计：沿序列切分 + 两次 All-to-All

Ulysses 的设计并不复杂，如上图所示，其思想可以概括为“先做序列并行，再做注意力头并行，最后切回序列并行”：

1. 序列维度切分：将每个样本的输入序列 N 均匀切分到 P 个参与设备上，每个设备持有长度为 N/P 的本地序列；
2. 本地投影 QKV：每个设备对自己持有的 N/P 段做投影，得到局部的 Query（Q）、Key（K）、Value（V）；
3. All-to-All 聚合：通过高度优化的 All-to-All 集合通信，把各设备局部的 Q、K、V 汇聚为全局 QKV——通信完成之后，每个设备拥有完整序列、但只拥有全部注意力头的一个互不重叠子集；
4. 逐头注意力计算：每个设备对自己分到的注意力头做标准注意力运算：Output context = Softmax(Q·Kᵀ / √d)·V此时序列已完整，注意力本身的计算不跨设备，因此可以与任意注意力实现自由替换；
5. 反向 All-to-All：注意力计算完成后，再一次 All-to-All 把输出张量沿着注意力头方向收集起来，同时重新沿序列维度切回 N/P，供 Transformer 层后续算子（MLP MatMul、LayerNorm 等）继续使用。

从源码实现看，这一过程对应的正是 deepspeed/sequence/layer.py 中的DistributedAttention：它把可微的 All-to-All 封装成自定义torch.autograd.Function（_SeqAllToAll），在forward中先用scatter_idx维度切分再在gather_idx维度聚合；backward则把 scatter/gather 对调并复用同一套 All-to-All，从而保证梯度在分布式切分下正确回传。仓库中还封装了_DimZeroAllToAll，它只负责在第 0 维上的切分与聚合，是更简化的基础算子。

通信复杂度：O(N/P) 对 O(N)，拉开数量级差距

Ulysses 区别于其他长序列方案的最核心优势在于更小的单链路通信量与随并行度更好的可扩展性。现代集群（节点内 NVSwitch、节点间 Fat-Tree IB 拓扑）中，规模为 M 的聚合消息在 P 个 GPU 间做 All-to-All 时，每条链路传输的数据量是 M/P。

对于一个隐藏维度为 h、序列长度为 N、并行度为 P 的 Transformer 模型：

• Ulysses 在注意力计算前对 QKV 做一次总规模为3Nh的 All-to-All，在注意力后对输出 context 再做一次规模为Nh的 All-to-All；
• 因此每层每条链路的聚合通信量为4Nh/P，即通信复杂度O(N/P)；
• 当 N 与 P按比例同时增大时，该通信量保持不变——这正是 Ulysses 能把序列长度随 GPU 数线性推高的数学基础。

作为对照，Megatron-LM 式序列并行每层执行两次消息量为 Nh 的all-gather与两次消息量为 Nh 的reduce-scatter。当 P≫1 时，每个all-gather/reduce-scatter的代价仍是 M 而不是 M/P，因此其单链路通信量为4Nh，通信复杂度为O(N)——恰好比 Ulysses 大P 倍。这意味着并行度越高，Ulysses 的相对优势越明显。

两大关键特性：Attention Agnostic 与 ZeRO-3 集成

对任意注意力实现一视同仁

Ulysses 的核心是“注意力中心”的模块化设计，这让分布式注意力模块具有完全通用的性质：在注意力计算之前是 N/P 的序列并行，而注意力计算本身只是更少注意力头数上的完整逐头注意力。因此DistributedAttention内的local_attention可以被替换为任意注意力实现——稠密自注意力、交叉注意力、因果注意力，以及各类在局部注意力层支持长序列的高效内核（如不同版本的 FlashAttention）。这正是仓库代码中DistributedAttention把原始注意力模块作为第一个参数传入的原因：它不关心内部实现，只负责外部的切分与聚合。

与 ZeRO-3 结合，模型与序列“同时长大”

序列并行削减的是激活值内存，却不会减少模型状态（参数、梯度、优化器状态）占用的内存。因此要支撑大模型+超长序列，Ulysses 必须与 ZeRO-3 协同：

• 训练数据天然可以按batch（样本）维度与序列维度两个方向切分，将两者对应的并行组组合，即可构成更大的 ZeRO 并行组；
• Ulysses 把 ZeRO-3 的模型状态切分扩展到“数据并行 + 序列并行”联合组上——模型参数在数据与序列两个维度上同时被分片，需要时按 rank 做 allgather 取回；
• 梯度同样跨数据并行与序列并行两组 rank 归约后再更新参数。

这样内存节省同时发生在序列维与数据维，使规模扩展不仅限于序列长度，也覆盖模型体量。这也解释了官方评测中 Ulysses 总能比对照方法塞下更多样本、更长序列的原因。

官方评测回顾：从百万 Token 到收敛性

以下数据与结论均出自该 Ulysses 发布文档及其配套实验（GPT 模型、最多 64 张 40GB A100 GPU），用于说明其在典型稠密/稀疏注意力训练上的表现。

序列长度强扩展：最高 100 万 Token

在 1.2B 参数 GPT 模型上，Ulysses 实现了序列长度随 GPU 数量线性增长的强扩展实验，如图 2 所示：在合适的 GPU 数量下，不同序列长度维持了相近的计算吞吐——这正是其 O(N/P) 通信复杂度带来的直接结果。

稠密注意力：对 Megatron-LM 与 ColAI-SP 的系统性优势

官方在 7B 与 30B 参数稠密注意力 GPT 模型上，分别用 32 与 64 张 A100 与 Megatron-LM 序列并行（Megatron-LM）和 ColossalAI 序列并行（ColAI-SP）对比，并各自选取最优的（序列并行度 × micro-batch）配置；Ulysses 在 7B/30B 下分别使用 ZeRO-3 并行度 32/64。结果显示：在两者都能运行的序列长度区间内，Ulysses 吞吐一致占优，且能跑到对方跑不了的更长序列，优势来自两点——(1) ZeRO-3 在数据+序列联合组上的参数分片省下内存、可容纳更多样本；(2) 高效的all-to-all优于 Megatron-LM/ColAI-SP 的all-gather+reduce-scatter与环形 P2P 通信。文档同时如实指出：稠密注意力在长序列下吞吐主要由注意力本地计算的二次复杂度决定，因此与对照方法的差距会随序列变长而收窄。

稀疏注意力：吞吐 2 倍以上、序列长度 4 倍于对手

在稀疏（block sparse）注意力下，官方仅与 Megatron-LM 对比（ColAI-SP 无公开的 block sparse 注意力实现）：Ulysses 吞吐达到 Megatron-LM 的2 倍以上，并借助 ZeRO-3 的内存优势把可训练序列长度扩展到对手的4 倍。文档也指出，当前 Ulysses 的稀疏注意力吞吐受限于本地稀疏注意力实现，序列越长吞吐越低，并预期未来改进本地稀疏内核后差距会进一步拉大。

收敛性：纯系统优化、无损模型质量

最后，官方在 8 张 A100 上对 1.3B GPT 以 32K 序列长度（序列并行度 4）对比了 DeepSpeed 与 Megatron-LM 序列并行的收敛曲线，并分别测试了 DeepSpeed 在不同 ZeRO 阶段下的表现（相关曲线见 blogs/deepspeed-ulysses/media/convgZ.png）。结论是：Ulysses 是纯粹的系统级优化手段，对训练出的模型质量没有（负向）影响。

怎么用：极简接入与通信组初始化

两行代码替换注意力

原文档给出的接入方式极其精简：只要把原有的自注意力模块attn包装进DistributedAttention即可：

from deepspeed.sequence.layer import DistributedAttention # Replace the original self-attention (attn) with DeepSpeed-Ulysses's self-attention dist_attn = DistributedAttention(attn, get_sequence_parallel_group())

相比 Megatron-LM 等库，Ulysses不要求重构整个模型，这也是其“可移植、易用”的卖点。该 API 目前仍位于 deepspeed/sequence/layer.py，仓库中的签名已进一步演化为：

DistributedAttention( local_attention: Module, # 任意本地注意力实现 sequence_process_group: ProcessGroup, # 序列并行通信组 scatter_idx: int = 2, # all-to-all 切分维度 gather_idx: int = 0, # all-to-all 聚合维度 sp_stream=None, # 传入流时开启通信/计算重叠 num_kv_heads: int = None, # GQA 下 KV 头数（不整除时传入） )

其中scatter_idx/gather_idx的语义是：对输入的 Q/K/V 做序列切分→头聚合，对注意力输出做头切分→序列聚合。仓库代码还支持通过sp_stream开启 Q/K/V 前向与反向 All-to-All 与矩阵运算的通信重叠路径。

构建序列并行通信组

DistributedAttention需要拿到序列并行通信组。参考 docs/_tutorials/ds-sequence.md，可以在模型初始化阶段构建：

def initialize_model_parallel(..., sequence_parallel_size, ...): num_sequence_parallel_groups = world_size // sequence_parallel_size num_sequence_data_parallel_groups = (world_size // sequence_parallel_size // data_parallel_size) global _SEQUENCE_PARALLEL_GROUP for i in range(num_sequence_parallel_groups): ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size) group = torch.distributed.new_group(ranks) if rank in ranks: _SEQUENCE_PARALLEL_GROUP = group def get_sequence_parallel_group(): """Get the sequence parallel group the caller rank belongs to.""" return _SEQUENCE_PARALLEL_GROUP

需要注意的是：注意力头数必须能被序列并行度整除。仓库中这一约束由运行时校验强制保证——例如 deepspeed/runtime/sequence_parallel/parallel_state_sp.py 要求world_size % sequence_parallel_size == 0，且 head 数（以及 KV head 数）必须与 SP 规模整除或互为倍数关系。

在 Megatron-DeepSpeed 中启用

Ulysses 已完整集成并测试于 Megatron-DeepSpeed 仓库。具体做法是：把原来直接调用本地CoreAttention的前向改为调用DistributedAttention包装后的dist_attn（Q/K/V 与 attention mask 的传入方式不变）；然后用--ds-sequence-parallel-size参数设定序列并行度（同样要求注意力头数可被该值整除），仓库准备了类似ds_pretrain_gpt_1.3B_seq_parallel_32k.sh、ds_pretrain_gpt_30B_seq_parallel_32k.sh的示例脚本可供直接运行。需注意：该序列并行特性目前与 Megatron-LM 的张量并行、流水并行不兼容。

叠加 FlashAttention / Triton 内核

序列并行可以与多种注意力内核叠加以进一步改善显存与算力效率：

原教程建议为追求最佳性能使用 FlashAttention + Triton（当时实测于triton==2.0.0.dev20221202），并提示 FlashAttention 仅兼容 NVIDIA Turing/Ampere/Ada/Hopper 架构 GPU；为发挥最优性能，注意力头的 head size 应能被 8 整除（详见 FlashAttention 自身约束）。

从源码看 Ulysses 的当下形态与更多生态

在最新仓库中，Ulysses 的实现已经分化为两条互补的接入路径，读者可依自身框架选用：

• Megatron-DeepSpeed 风格的DistributedAttention：即上文 deepspeed/sequence/layer.py 中基于 Megatron 布局（[seq, batch, hidden]）的实现，同时支持非均匀头切分与 All-to-All 通信/计算重叠。
• 面向 HF Transformers 的UlyssesSPAttentionHF：位于 deepspeed/runtime/sequence_parallel/ulysses_sp.py，它以标准[seq, batch, head, dim]输入形式重新实现了DistributedAttention，并进一步适配了 HF Transformers 支持的多种注意力后端（sdpa、flash_attention_2/3、flex_attention 等）。其类方法register_with_transformers会调用 parallel_state_sp.py 的initialize_sequence_parallel建立 SP 组，并用 Ulysses 包装函数替换ALL_ATTENTION_FUNCTIONS中的对应实现——这样模型其它代码路径完全不动，只在真正调用注意力那一刻被“截获”。配套的UlyssesSPDataLoaderAdapter数据加载器适配器会把普通 DataLoader 产出的 batch 按序列维切分再分发给各 SP rank，并要求 batch 中携带input_ids、position_ids、labels（因果掩码依赖全局 position 重建）。对应的用法教程见 docs/_tutorials/ulysses-alst-sequence-parallelism.md。
• DeepSpeed Engine 原生支持：从 deepspeed/runtime/engine.py 与配置文件看，DeepSpeed 配置中可直接写入data_parallel_size与sequence_parallel_size（deepspeed/init.py 通过 mesh device 方式创建 data/sequence 两组），Engine 会据此暴露get_sequence_parallel_group()、seq_parallel_group等接口；ZeRO-1/2/3 各阶段（如 stage3.py）都会感知 SP 规模来正确折算参与梯度归约与参数分片的世界大小。
• 测试印证：仓库在 tests/unit/v1/sequence_parallelism/test_ulysses.py 中覆盖了非序列并行 checkpoint 在 SP 训练下的加载与再保存、mesh 设备下 SP/DP 组的构建，以及大批量（batch/sequence×head 数×head 维）参数化下 All-to-All 与原生参考实现的一致性校验；tests/unit/ulysses_alst/ 则覆盖 HF 版 SP 注意力与 Tile 化计算。

小结与使用前提

DeepSpeed-Ulysses 以“序列切分 + 注意力头并行 + 两次 All-to-All”的极简骨架，实现了通信量 O(N/P)的长序列扩展路径，并通过 attention-agnostic 设计兼容稠密/稀疏注意力与 FlashAttention，通过与 ZeRO-3 的联合分片同时支撑超大模型与超长序列。官方发布报告在 64 张 A100 上已展示4 倍于既有系统的序列长度、10 倍以上的通信削减、最高约 2.5 倍吞吐提升，以及超过 1M Token 的训练能力。

动手前请确认以下前提：DeepSpeed 版本不低于 v0.10.2（见 docs/_tutorials/ds-sequence.md）；注意力头数与 KV 头数需与序列并行度整除匹配；SP 规模需能整除总 world size；当前实现不与 Megatron-LM 的张量/流水并行混用。满足这些条件后，替换一个DistributedAttention包装器，即可把长上下文窗口训练推向下一个数量级。

延伸阅读：同主题更深入的源代码可在 deepspeed/sequence/layer.py、deepspeed/runtime/sequence_parallel/ulysses_sp.py 与 deepspeed/runtime/sequence_parallel/parallel_state_sp.py 中继续追踪；配套教程见 docs/_tutorials/ds-sequence.md 与 docs/_tutorials/ulysses-alst-sequence-parallelism.md。

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