DeepSpeed Ulysses：面向极长序列 Transformer 模型训练的系统优化
机器学习
2023-10-05 v2 计算与语言
分布式、并行与集群计算
摘要
典型基于 Transformer 的大语言模型（LLM）的计算可由批大小、隐藏维度、层数与序列长度刻画。迄今，用于加速 LLM 训练的系统工作聚焦于前三个维度：批大小的数据并行、隐藏维度的张量并行与模型深度或层数的流水线并行。这些被广泛研究的并行形式并非针对长序列 Transformer 模型或为其优化。鉴于长序列 LLM 的实际应用需求，序列并行正重新受到关注。然而，现有序列并行工作受限于内存-通信低效，限制了其对长序列大模型的可扩展性。本工作中，我们提出 DeepSpeed-Ulysses，一种新颖、可移植且高效的方法，用于实现极长序列长度下高度高效且可扩展的 LLM 训练。DeepSpeed-Ulysses 的核心沿序列维度划分输入数据，并采用高效的全交换（all-to-all）集合通信进行注意力计算。理论通信分析表明，其他方法随序列长度增加产生通信开销，而 DeepSpeed-Ulysses 在序列长度与计算设备成比例增加时保持恒定通信量。此外，实验评估显示，相较于现有 SOTA 基线方法，DeepSpeed-Ulysses 在序列长度长 4 倍的情况下训练快 2.5 倍。
引用
@article{arxiv.2309.14509,
title = {DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models},
author = {Sam Ade Jacobs and Masahiro Tanaka and Chengming Zhang and Minjia Zhang and Shuaiwen Leon Song and Samyam Rajbhandari and Yuxiong He},
journal= {arXiv preprint arXiv:2309.14509},
year = {2023}
}