DeepSpeed-Ulysses：支持極長序列 Transformer 模型訓練的系統優化方法

flyfish

名字 Ulysses

“Ulysses” 和 “奧德修斯（Odysseus）” 指的是同一人物，“Ulysses” 是 “Odysseus” 的拉丁化版本

《尤利西斯》（詹姆斯·喬伊斯著，1922年出版）

以古希臘英雄奧德修斯（拉丁名“尤利西斯”）的神話為框架，將故事置于現代（20世紀初）的都柏林，通過三位主角一天內的經歷，映射奧德修斯的十年漂泊。

奧德修斯（古希臘神話核心人物）

奧德修斯是古希臘城邦伊薩卡的國王，以智慧著稱（如“特洛伊木馬”計策的設計者）。特洛伊戰爭結束后，他因觸怒海神波塞冬，踏上長達十年的返鄉漂泊，故事集中于《奧德賽》（荷馬史詩）。

“DeepSpeed-Ulysses”的借用“奧德修斯/尤利西斯”的“長途探索與克服挑戰”意象，象征長序列模型訓練的技術突破。

長序列Transformer模型訓練是指針對輸入序列長度極長（如包含數萬至百萬級 tokens）的Transformer模型進行優化訓練的過程，其核心挑戰在于平衡計算效率、內存使用和通信開銷。

一、長序列Transformer模型的定義與應用場景

1. 模型特征

Transformer模型以自注意力機制為核心，傳統模型的序列長度通常限制在數千tokens（如512或2048），而長序列模型需處理遠超這一范圍的輸入。其計算復雜度隨序列長度呈二次增長（注意力計算需遍歷所有token對），導致訓練時內存占用和通信開銷劇增。

2. 關鍵應用領域

生成式AI：
- 對話系統：處理長對話歷史（需支持多輪交互的上下文）。
- 長文檔處理：對書籍、學術論文等進行摘要生成（序列長度可達數萬詞）。
- 多模態任務：視頻生成、語音識別等需處理時空維度的長序列輸入。
科學AI：
- 基因組學：分析人類基因組的64億堿基對序列。
- 氣候預測：處理長時間序列的氣象數據。
- 醫療診斷：基于患者全病程記錄的預測模型。

二、長序列訓練的核心挑戰

1. 內存與通信效率瓶頸

激活值內存爆炸：Transformer的中間激活值（如注意力計算中的QKV張量）隨序列長度平方增長，單GPU內存無法容納極長序列。
通信開銷線性增長：現有序列并行方法（如Megatron-LM）在注意力計算時需進行全聚集（all-gather）通信，通信量隨序列長度線性增加，導致訓練效率隨序列長度下降。

2. 現有并行技術的局限性

數據/張量/流水線并行不適用：這些并行方式分別針對批量大小、隱藏維度、模型層數優化，但未考慮序列維度的劃分，無法解決長序列特有的內存和通信問題。
傳統序列并行方法低效：如ColAI-SP采用環形通信，Megatron-LM依賴全聚集操作，均無法在擴展序列長度時保持通信效率，且需大量代碼重構。

DeepSpeed-Ulysses解決長序列Transformer模型訓練難題的方案

DeepSpeed-Ulysses針對長序列Transformer訓練的四大核心創新，可拆解為數據劃分策略、通信原語重構、內存分層優化、注意力模塊化設計四大技術模塊
DeepSpeed-Ulysses通過序列分片解耦計算與存儲、全對全通信重構帶寬效率、ZeRO-3擴展內存邊界、注意力抽象層屏蔽多樣性，形成了從硬件通信到算法邏輯的端到端優化。

一、序列維度數據劃分：動態分片與負載均衡

1. 分片粒度設計

均勻切分：將輸入序列按長度均分為P份（P為GPU數），每個GPU處理N/P tokens（如N=1M，P=64時，單卡處理15,625 tokens），確保計算負載均衡。
細粒度對齊：結合注意力頭數（H），將每個分片的QKV張量按頭維度切分，實現“序列分片×頭分片”的二維并行（如H=32，P=8時，每卡處理4個頭的完整序列分片），避免頭間通信冗余。

2. 分片生命周期管理

前向傳播：分片在嵌入層后立即生成，貫穿注意力計算全流程，減少中間復制開銷。
反向傳播：梯度按分片聚合，通過通信原語自動對齊，無需顯式分片重組（對比Megatron-LM需手動拼接梯度）。

二、通信原語重構：全對全替代全聚合

1. 兩次全對全（All-to-All）通信

第一次All-to-All（QKV收集）：
- 輸入：各GPU持有(N/P, b, d)的Q/K/V分片（b為微批量，d為隱藏維度）。
- 操作：通過高效的NVLink/IB網絡，交換分片形成(N, b, d/P)的全局QKV（每個GPU獲取完整序列的1/P頭數據）。
- 優勢：通信量為3Nh/P（h為頭維度），對比Megatron-LM的全聚合（3Nh）減少P倍。
第二次All-to-All（結果分發）：
- 輸入：注意力計算后的(N, b, d/P)輸出。
- 操作：重組為(N/P, b, d)分片，供后續MLP層處理，通信量Nh/P。
- 總通信量：4Nh/P，復雜度O(N/P)，當N與P同比增長時保持恒定（如N×2，P×2，通信量不變）。

2. 通信拓撲優化

節點內NVSwitch優先：利用GPU間高速互聯，將All-to-All拆分為節點內（低延遲）和節點間（高帶寬）兩階段，實測減少通信時間
異步通信隱藏延遲：在注意力計算前預啟動通信，與計算重疊（如FlashAttention的分塊計算與通信流水線）。

三、內存分層優化：ZeRO-3×序列并行

1. 參數分片策略

跨組劃分：ZeRO-3原僅在數據并行組（DP組）分片參數，Ulysses擴展至DP組×序列并行組（SP組），形成二維分片（如DP=8，SP=8時，每個參數被分為64片）。
動態聚合：僅在參數更新時，通過All-Gather收集分片梯度，平時僅保留本地分片，顯存占用降低DP×SP倍。

2. 激活內存壓縮

分片激活：每個GPU僅存儲N/P tokens的激活值，總激活內存從O(N2)降至O((N/P)2)，支持N=1M時單卡激活內存僅為全序列的1/64（P=64）。
與激活重計算協同：結合FlashAttention的分塊重計算，進一步將激活內存降至O(N/P)。

四、注意力模塊化設計：統一接口與動態適配

1. 注意力計算抽象層

輸入標準化：無論密集/稀疏，輸入均為(N, b, d/P)的分片QKV（N為全局序列長度）。
頭級并行：每個GPU負責H/P個頭的完整注意力計算，支持任意頭數分片（如H=96，P=8時每卡12頭）。
輸出重組：通過第二次All-to-All自動對齊頭維度，屏蔽底層注意力實現差異。

2. 稀疏注意力特化優化

塊對角通信：針對滑動窗口等稀疏模式，僅交換相鄰分片的K/V，通信量進一步降至O(N/P×W)（W為窗口大小）。
與Blockwise Attention協同：支持將序列劃分為塊，塊內全連接、塊間稀疏連接，通信量減少70%（對比全局密集）。

3. 兼容性驗證

FlashAttention v2集成：通過分片QKV直接輸入FlashAttention內核，實測長序列（N=512K）吞吐量提升1.8倍
稀疏模式擴展：支持BigBird的全局+局部混合注意力，無需修改通信邏輯，僅替換注意力計算函數

五、工程實現細節：易用性與擴展性

1. 代碼侵入性

零修改集成：僅需在模型定義中添加@sequence_parallel裝飾器，自動替換原生Attention為分布式版本
并行組抽象：自動管理DP/SP組的創建與銷毀，用戶無需顯式管理通信域。

2. 故障恢復機制

分片檢查點：僅保存本地分片參數，檢查點大小減少DP×SP倍，支持1024卡集群下30B模型的快速恢復

注意力集成

DeepSpeed-Ulysses通過優化多種注意力機制，系統性解決了長序列Transformer訓練的效率瓶頸。

一、自注意力（Self-Attention）

1. 密集自注意力（Dense Self-Attention）

數學公式：
$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$
其中：
- $\in \mathbb{R}^{N \times d}$ , $\in \mathbb{R}^{N \times d}$ , $\in \mathbb{R}^{N \times d}$
- $N$ 為序列長度， $d$ 為隱藏維度
解決的問題：
- 傳統實現瓶頸：計算復雜度 $O(N^2)$ 和內存占用 $O(N^2)$ ，難以處理超長序列（如 $N > 10 k$ ）
- DeepSpeed-Ulysses優化：
  - 通過序列并行（Sequence Parallelism）將 $N$ 劃分為 $N / P$ （ $P$ 為GPU數）
  - 兩次All-to-All通信將計算復雜度降至 $O (N / P)$ （

2. 稀疏自注意力（Sparse Self-Attention）

數學公式：
$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T \odot M}{\sqrt{d_k}}\right)V$
其中：
- $M$ 為稀疏掩碼矩陣，僅保留部分非零連接
解決的問題：
- 傳統密集注意力缺陷：長序列下計算和內存開銷爆炸
- DeepSpeed-Ulysses優化：
  - 支持滑動窗口（Sliding Window）、塊對角（Block Diagonal）等稀疏模式
  - 通信量降至 $\times W)$ （ $W$ 為窗口大小）

二、交叉注意力（Cross-Attention）

數學公式：

$\text{CrossAttention}(Q, K_{\text{enc}}, V_{\text{enc}}) = \text{Softmax}\left(\frac{QK_{\text{enc}}^T}{\sqrt{d_k}}\right)V_{\text{enc}}$
其中：

$Q$ 來自解碼器， $K_{\text{enc}}, V_{\text{enc}}$ 來自編碼器
解決的問題：
- 長序列跨模態挑戰：生成式任務中，解碼器需高效關注超長編碼序列
- DeepSpeed-Ulysses優化：
  - 通過序列分片和分布式通信，支持 $N_{\text{enc}} \times N_{\text{dec}}$ 規模的交叉注意力
  - 結合ZeRO-3減少參數內存占用，支持萬億級參數模型

三、因果注意力（Causal Attention）

數學公式：

$\text{CausalAttention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T + \text{Mask}}{\sqrt{d_k}}\right)V$
其中：

$\text{Mask}_{i,j} = \begin{cases} 0 & \text{if } i \geq j \\ -\infty & \text{otherwise} \end{cases}$
解決的問題：
- 自回歸生成的信息泄露：防止未來token影響當前預測
- DeepSpeed-Ulysses優化：
  - 在分布式環境下高效應用掩碼，避免跨分片的未來信息泄露
  - 支持超長上下文生成

四、Blockwise Attention（分塊注意力）

數學公式：

$\text{BlockwiseAttention}(Q, K, V) = \sum_{i=1}^{B} \text{Softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right)V_i + \sum_{i \neq j} \text{Softmax}\left(\frac{Q_i K_j^T}{\sqrt{d_k}}\right)V_j \cdot M_{i,j}$
其中：

$B$ 為塊數， $M_{i,j}$ 為塊間連接掩碼
解決的問題：
- 長序列計算碎片化：將序列劃分為塊，塊內密集計算，塊間稀疏交互
- DeepSpeed-Ulysses優化：
  - 塊內通信量 $O(N_{\text{block}}^2)$ ，塊間通信量 $O(N_{\text{block}} \times K)$ （ $K$ 為稀疏度）
  - 實驗顯示通信量減少70%，支持塊間全局連接（如BigBird）

五、FlashAttention優化

數學公式：

$\text{FlashAttention}(Q, K, V) = \sum_{i=1}^{T} \text{Softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right)V_i$
其中：

$T$ 為內存塊數，通過分塊計算減少峰值內存
解決的問題：
- 長序列內存瓶頸：傳統注意力的激活值內存占用 $O(N^2)$
- DeepSpeed-Ulysses優化：
  - 結合序列并行，將內存降至 $O (N / P)$
  - 分塊計算與通信重疊，提升計算效率

六、DeepSpeed-Ulysses的核心優化公式

1. 分布式注意力計算

$\text{Output} = \text{AllToAll}\left(\text{Softmax}\left(\frac{\text{AllToAll}(Q_i, K_i, V_i)}{\sqrt{d_k}}\right)\right)$
其中：

$Q_i, K_i, V_i$ 為第 $i$ 個GPU的序列分片
解決的問題：
- 傳統數據并行的通信瓶頸：全聚集（All-Gather）通信量 $O (N)$
- 優化后：兩次All-to-All通信量降至 $O (N / P)$

2. 內存優化公式

$\text{Memory} = \text{Params} + \text{Activations} + \text{Gradients}$
其中：

$\text{Params} = \frac{\text{TotalParams}}{P \times S}$ （ $S$ 為ZeRO-3分片因子）
$\text{Activations} = O\left(\frac{N}{P}\right)$ （通過序列分片和激活重計算）
解決的問題：
- 長序列訓練的內存墻：傳統方法無法支持 $N > 100 k$ 的序列
- 優化后：在7B模型上支持 $N = 512 k$ ，內存效率提升4倍

All-to-All通信

分布式計算中的All-to-All通信與DeepSpeed-Ulysses優化

DeepSpeed-Ulysses通過兩次All-to-All通信，實現了長序列Transformer訓練的三大突破：
將O(N)復雜度降為O(N/P)，打破長序列訓練的帶寬瓶頸；利用NVSwitch/IB拓撲和異步流水線，將GPU算力利用率提升至硬件峰值的54%；支持1M token序列與萬億參數模型的聯合訓練，為科學AI（如基因組分析）和生成式AI（如長文本生成）提供底層支撐。

一、All-to-All通信的核心概念

1. 通信機制本質
All-to-All是分布式系統中一種數據交換模式：網絡中的每個節點（如GPU）同時向其他所有節點發送自身數據分片，并接收所有節點發來的分片。例如，當有P個GPU時，每個GPU初始持有數據切片X?, X?, …, X?，通信后每個GPU都會獲得完整的[X?, X?, …, X?]，相當于每個節點都收集到全局所有數據分片，形成完整數據集。

2. 與其他通信方式的區別

All-Gather：所有節點僅收集同一數據的不同分片（如多個GPU共同組成完整模型參數），通信量隨數據規模線性增長。
All-to-All：節點間雙向交換各自分片，通信量隨節點數P成反比，適合數據并行場景下的高效同步。

二、DeepSpeed-Ulysses的兩次All-to-All通信設計

1. 第一次通信：QKV張量聚合（注意力計算前）

輸入狀態：每個GPU持有序列長度為N/P的Q/K/V張量（形狀為(N/P, b, d)，b為微批量，d為隱藏維度）。
通信過程：
- 每個GPU向其他P-1個GPU發送自己的Q/K/V分片；
- 每個GPU接收所有分片后，拼接成全局序列的Q/K/V張量（形狀變為(N, b, d/P)），其中每個GPU負責d/P維度的注意力頭計算。
核心作用：使每個GPU獲取完整序列的部分注意力頭數據，為并行計算注意力做準備。

2. 第二次通信：注意力結果重分配（計算后）

輸入狀態：注意力計算后的輸出張量形狀為(N, b, d/P)。
通信過程：
- 每個GPU將輸出按序列維度重新切分為N/P長度的分片；
- 通過All-to-All將分片發送至對應GPU，最終每個GPU持有(N/P, b, d)的張量，供后續MLP層處理。
核心作用：將注意力計算結果按序列維度重新分配，確保后續操作的并行性。

三、通信量優化的數學邏輯

1. 傳統方法（如Megatron-LM）的瓶頸
采用All-Gather+Reduce-Scatter組合：

每次注意力計算需2次All-Gather（收集QKV）和2次Reduce-Scatter（聚合結果），總通信量為4Nh（N為序列長度，h為隱藏維度），復雜度為O(N)。
問題：當N從16K增至1M時，通信量增長62.5倍，極易造成網絡擁塞。

2. DeepSpeed-Ulysses的優化公式
兩次All-to-All的通信量計算：

第一次通信（QKV收集）：3Nh/P（Q、K、V三個張量）；
第二次通信（結果分發）：Nh/P；
總通信量：4Nh/P，復雜度降為O(N/P)。
關鍵優勢：當N與P同比增長時（如N×2，P×2），通信量保持不變。例如：
- N=256K，P=64時，傳統方法通信量為1024Kh，DeepSpeed僅為16Kh（減少64倍）；
- N=1M，P=64時，傳統方法為4000Kh，DeepSpeed為62.5Kh（僅為1.56%）。