Ulysses 序列平行化 用於 百萬標記 上下文訓練

Hugging Face 已將 Ulysses 序列平行化整合至 Accelerate、Transformers 與 TRL，使得透過在多個 GPU 上分散注意力計算，能夠訓練具備百萬標記上下文的 LLM。

長序列訓練的挑戰

標準的 transformer 注意力在 FLOPs 與記憶體上皆以序列長度 $n$ 的二次方 ($O(n^2)$) 成長。雖然 FlashAttention 可將記憶體使用量降低至 $O(n)$，但 $O(n^2)$ 的計算需求仍然存在。對於超過 32k token 的序列，訓練通常會超出單一 GPU 的記憶體上限，因而需要將序列本身跨多個裝置切分，而非僅依賴資料平行化。

Ulysses 序列平行化的運作方式

Ulysses 序列平行化 (SP) 透過在 GPU 之間分割序列維度與注意力頭來分散注意力計算。其流程如下：

  1. Sequence Sharding：輸入序列會在 $P$ 個 GPU 之間切分，每個 GPU 持有本地的標記區塊。
  2. QKV Projection：每個 GPU 為其本地區塊計算 query、key 與 value 的投影。
  3. All-to-All Communication：透過 all-to-all 集體操作重新分配資料，使每個 GPU 持有所有序列位置，但僅限於部分 attention heads。
  4. Local Attention：GPU 使用 FlashAttention 或 SDPA 計算其分配的 heads 的注意力。
  5. All-to-All Communication：第二次 all-to-all 操作將資料返回至序列切分的格式。
  6. Output Projection：每個 GPU 為其本地序列區塊計算 output projection。

通訊複雜度

Ulysses 於每個 attention layer 需要兩次 all-to-all 操作，每個 GPU 的通訊量為 $O(n · d / P)$（其中 $n$ 為序列長度，$d$ 為隱藏維度，$P$ 為平行度）。此效率高於 Ring Attention，其每個 GPU 的通訊量為 $O(n · d)$，且在 $P-1$ 個跳躍上序列化傳輸。

生態系統整合

Accelerate

Accelerate 透過 ParallelismConfig 類別與 DeepSpeed 整合來實作 Ulysses。關鍵參數包括 sp_size（用於序列平行化的 GPU 數量）與 sp_backend（必須設為 "deepspeed"）。當呼叫 accelerator.prepare() 時，系統會自動以 UlyssesSPAttentionHF 註冊模型，並以 UlyssesSPDataLoaderAdapter 包裝 dataloader。

Transformers Trainer

Transformers 的 Trainer 透過 TrainingArguments.parallelism_config 處理 Ulysses 整合。它自動化 dataloader 包裝、序列切分與加權 loss 聚合，確保即使 token 在不同 rank 上分配不均，梯度仍正確。

TRL SFTTrainer

TRL 的 SFTTrainer 為監督式微調加入最佳化，例如使用 packing 功能以減少 padding 浪費。它要求 pad_to_multiple_of 必須等於 sp_size 以確保序列可被整除。啟用 Ulysses 時，SFTTrainer 亦會自動管理預先移位的 labels。

Ulysses 與 Ring Attention 的比較

項目 Ulysses (DeepSpeed) Ring Attention (FSDP2)
平行化方法 Attention head partitioning Ring-based KV exchange
後端 DeepSpeed ZeRO PyTorch FSDP2
注意力支援 FlashAttention 2/3, SDPA SDPA only
通訊 Two all-to-alls per layer P2P ring communication
每個 GPU 的通訊量 $O(\text{total_seq} \times \text{hidden} / \text{sp_size})$ $O(\text{total_seq} \times \text{hidden})$
頭數限制 num_heads >= sp_size None

效能基準

Hugging Face 使用 Qwen3-4B 在 Gutenberg 英文資料集上，搭配 H100 80GB GPU，對 Ulysses SP 進行了效能基準測試。

記憶體減少

使用 SP=4 在相同序列長度下將每個 GPU 的記憶體使用量降低 3.3 倍。這使得可從 8K token（DP=4）的基線擴展至 96K token（SP=4），仍在 80GB 記憶體上限內。於 128K token 時，配置遭遇記憶體不足（OOM）狀態。

吞吐量

隨著序列長度增加，吞吐量提升，因為二次方的注意力計算主導了通訊開銷。於 64K token 時，SP=4 達到 13,396 token/秒，為 8K 基線的 3.7 倍。

實作最佳實踐

  • Sequence Divisibility：使用 pad_to_multiple_of 確保序列長度能被 sp_size 整除。
  • Attention Backends：對於 Ampere GPU 使用 FlashAttention 2，對於 Hopper GPU 使用 FlashAttention 3。
  • Memory Optimization：將 Ulysses 與 DeepSpeed ZeRO Stage 3 結合，並使用環境變數 PYTORCH_ALLOC_CONF=expandable_segments:True 以減少碎片化。
  • 2D Parallelism：根據 GPU 數量平衡 sp_size 與 dp_shard_size，以優化最大序列長度或更高吞吐量。
  • Additional Kernels：使用 Liger-Kernel 的 FusedLinearCrossEntropy 與 TiledMLP，進一步在 loss 計算與大型矩陣運算時減少工作記憶體。

Sources