Ulysses(DeepSpeed 提出)专门解决超长序列训练的显存与通信瓶颈。两个核心创新:环形切分拉平因果掩码下的负载不均;All-to-All 通信让单卡通信量与序列长度无关。能在数十张卡上轻松扛百万级 token。

为什么需要 Ulysses

序列越长,注意力激活值占用呈平方级增长。单张 80 GB H100 处理几万 token 就吃紧,百万级 token 完全无能为力。

传统并行方式都解决不了:

策略短板
数据并行只切 batch,激活值单卡仍要存完整
张量并行切隐藏维度,对序列长度无效
传统序列并行All-Gather 通信量随序列线性增长;因果掩码导致负载不均

两大创新

环形切分(Striped Partitioning)拉平负载

传统按顺序切:第 8 张 GPU 拿到序列尾部,因果掩码下要关注前面所有 token,计算量数倍于第 1 张 GPU。

Ulysses 按 token_id mod N 切分(每 N 个 token 取一个),每张 GPU 都均匀持有”靠前”和”靠后”的 token,计算量自动拉平。

All-to-All 通信替代 All-Gather

传统 All-Gather:每张 GPU 把完整 K/V 汇总到一起,通信量随序列线性增长。

Ulysses 用 All-to-All:每张 GPU 只交换”自己负责的注意力头” × “全局序列”那块数据。当序列长度与 GPU 数量按比例增长时,单卡通信量保持常数

四步工作流

以 8 张 GPU、100 万 token 为例:

  1. 序列分区:环形切分,每卡 12.5 万 token
  2. 局部 QKV 投影:每卡只对自己的 token 做 Q/K/V 线性变换
  3. All-to-All 重排:从「按序列切」转为「按注意力头切」,每卡拿到全局 K/V,但只算自己的几个注意力头
  4. 结果回传:再来一次 All-to-All 把结果转回「按序列切」,拼接为完整输出

三大优势

  • 显存省:激活值降到 1/N,64 张 A100 可处理 100 万 token
  • 速度快:与 FlashAttention、ZeRO-3 等可叠加,吞吐相比传统序列并行 ≥ 2×
  • 好上手:DeepSpeed 中替换注意力模块 + 配通信组即可,兼容 dense / sparse 注意力

适用场景

  • 长对话、书籍级摘要、长文档检索增强
  • 基因组、气候等 AI for Science 的长时序数据
  • 多模态长视频/长音频
  • 单卡显存不足但有多卡资源(哪怕只有 2-4 张)

短序列任务无需启用,普通 DDP/FSDP 更直接高效。