Ulysses 序列並行技術：實現百萬級 Token 上下文訓練

Ulysses 序列並行：百萬級 Token 上下文訓練

Ulysses 序列並行（Snowflake AI Research 提出的 Arctic 長序列訓練 (ALST) 協議的一部分）透過注意力頭並行（attention head parallelism）將注意力計算分佈在多個 GPU 上，提供了一個優雅的解決方案。在本篇文章中，我們將探討 Ulysses 的工作原理，以及它如何整合到 Hugging Face 生態系統中——從 Accelerate 到 Transformers Trainer 以及 TRL 的 SFTTrainer。

目錄

長序列訓練的挑戰

Transformer 中的注意力機制隨序列長度呈平方級增長。對於長度為 $n$ 的序列，標準注意力機制需要 $O(n^2)$ 的 FLOPs 以及 $O(n^2)$ 的記憶體來計算和儲存注意力評分矩陣。像 FlashAttention 這樣的優化實現透過分塊計算且不生成完整的注意力矩陣，將記憶體需求降低到 $O(n)$，但 $O(n^2)$ 的計算量仍然存在。對於極長序列（32k+ tokens），即使使用 FlashAttention，訓練仍然會挑戰單個 GPU 記憶體的極限。

考慮以下長上下文訓練至關重要的場景：

傳統的資料並行（Data Parallelism）對此無濟於事——每個 GPU 仍需在注意力模組內處理完整序列。我們需要一種將序列本身拆分到多個裝置上的方法。

Ulysses 的工作原理

DeepSpeed Ulysses 論文中引入的 Ulysses 序列並行 (SP) 採取了一種聰明的方法：除了在序列維度上進行拆分外，它還將注意力頭劃分到各個 GPU 上。

以下是它的運作方式：

1. 序列分片 (Sequence Sharding)：輸入序列沿序列維度拆分到 $P$ 個 GPU 上。每個 GPU $i$ 持有 token $[i \cdot n/P, (i+1) \cdot n/P)$。

2. QKV 投影 (QKV Projection)：每個 GPU 為其本地序列塊計算 Query、Key 和 Value 投影。

3. 全對全通訊 (All-to-All Communication)：一個全對全集體操作重新分佈資料，使得在投影之後，每個 GPU 持有所有序列位置，但僅持有注意力頭的一個子集。

4. 本地注意力 (Local Attention)：每個 GPU 使用標準注意力機制（FlashAttention 或 SDPA）為其分配的頭計算注意力。

5. 全對全通訊 (All-to-All Communication)：另一個全對全操作反向執行重新分佈，回到序列分片格式。

6. 輸出投影 (Output Projection)：每個 GPU 為其本地序列塊計算輸出投影。

核心洞察在於注意力頭是獨立的——每個頭可以單獨計算。透過用序列局部性交換頭局部性，Ulysses 以相對較低的通訊開銷實現了高效的並行化。

通訊複雜度

Ulysses 每個注意力層需要兩次全對全操作，每個 GPU 的總通訊量為 $O(n \cdot d / P)$，其中：

Ring Attention 每個 GPU 通訊量為 $O(n \cdot d)$ —— 是 Ulysses 的 $P$ 倍 —— 透過環狀結構進行 $P-1$ 次順序點對點傳輸。Ulysses 還受益於較低的延遲，因為全對全通訊可以在單個集體步驟中利用完整的對分頻寬（bisectional bandwidth），而 Ring Attention 則需在 $P-1$ 個跳轉（hops）上序列化。

與 Accelerate 整合

Accelerate 透過其 ParallelismConfig 類別和 DeepSpeed 整合，為 Ulysses 序列並行奠定了基礎。

配置

關鍵參數

使用 Accelerator

當你呼叫 accelerator.prepare() 時，Ulysses 會自動設定：

prepare() 呼叫會：

損失聚合 (Loss Aggregation)

使用 Ulysses 時，每個 GPU 在序列的不同部分計算損失。損失必須正確聚合，並根據每個 rank 的有效 token 數量進行加權。如果你使用的是 Transformers Trainer 或 TRL 的 SFTTrainer，這會自動處理——下面的程式碼僅在編寫自定義 Accelerate 訓練迴圈時才需要：

當 token 在各個 rank 之間分佈不均時（例如，某些 rank 僅包含填充或被遮蔽的 prompt token），加權損失聚合可確保梯度正確。

Ulysses 和 Ring Attention 在訓練期間都使用 position_ids 而非 attention_mask 進行因果遮蔽（causal masking）。在這種序列長度下，4D 注意力遮蔽將與注意力評分本身一樣令人望而卻步——在 128k tokens 時，那將是另一個約 1TB 的張量。Position IDs 以 $O(n)$ 記憶體而非 $O(n^2)$ 實現了相同的因果行為。在評估/推論期間，DeepSpeed 的 SP 注意力層可以完全繞過 SP 操作（透過 disable_in_eval）並回退到模型的預設注意力實現。

與 Transformers Trainer 整合

Transformers Trainer 透過 TrainingArguments.parallelism_config 提供無縫的 Ulysses 整合。它會自動處理所有 SP 相關細節——資料載入器封裝、序列分片和損失聚合——因此你不需要編寫上述任何自定義損失程式碼。

配置

只需將上述相同的 parallelism_config 傳遞給 TrainingArguments：

Trainer 自動處理的內容

1. 資料載入器封裝：模型準備後，Trainer 使用 UlyssesSPDataLoaderAdapter 封裝資料載入器。

2. 損失計算：compute_loss 方法會偵測 SP 模式並路由到專門的 _deepspeed_sp_compute_loss，後者處理：

3. 批次大小計算：有效資料並行世界大小（world size）會考慮 SP：

4. 資料載入器長度調整：訓練步數計算會根據 SP 對迭代次數的影響進行調整。

啟動指令

使用 accelerate 配置文件或命令列參數：

與 TRL SFTTrainer 整合

TRL 的 SFTTrainer 基於 Transformers Trainer，並針對長序列的有監督微調（SFT）添加了特定優化。