聊聊序列並行 Sequence parallelis

作者丨無惡不作

來源丨 https://zhuanlan.zhihu.com/p/653067104

編輯丨 GiantPandaCV

隨着大語言模型的發展，大模型的相關技術日益細化，在生成式 AI、蛋白質結構預測等多個領域開花結果。由此產生的技術痛點也越發明顯，特別是在蛋白質預測、多模態基礎模型、語音圖像處理模型等需要前後文推理的模型，對長序列訓練有明顯的需求。與此同時，現有的分佈式訓練技術，主要分爲：數據並行、張量並行、流水線並行等在數據維度正交切分的分佈式訓練技術。如下圖 0 所示，分別對應對於 batch size，hidden_size 維度進行切分。對於長序列模型訓練，由於 transformer 類模型 self-attention 部分序列長度和顯存消耗呈現平方關係，上述三種均不能提供較爲高效、穩定的訓練支持。而對於長序列模型採用序列並行可以很好的解決顯存溢出等問題，由此大模型訓練技術在序列並行技術逐漸開始迭代，下面針對主流的三種不同架構的序列並行進行介紹。

ColossalAI

ColossalAI 在 2022 年將序列並行技術集成到自己的代碼庫，細節可以參考：Sequence Parallelism: Long Sequence Training from System Perspective 論文。其主要聚焦於 transformer 結構的 self-attention 和 mlp 兩個結構，通過將序列切分到多個計算設備，並行計算，解決長序列訓練的問題。在 self-attention 部分，猶豫 Q 、 K 、 V 三者是互相交互的關係，如下圖 1 所示。

所以在實現上 ColossalA 借鑑了 Ring-Allreduce 算法實現，實現了 RingQK 和 RingAV 兩種計算方式，如下圖 2 所示。

這裏爲了方便理解我貼一下論文中的計算推導，方便理解爲什麼需要做 ring 操作。如下圖 3 所示。（其實就是線性代數中的分塊矩陣乘，這部分可以推理下）

(1) B: batch size; (2) L: sequence length; (3) H: hidden size of linear layers; (4) A: attention head size; (5) Z: number of attention heads; (6) N: number of GPUs.

RingQK：

RingAV：

上述實現也可以使用 AllGather 和 AllReduce 實現，或者 AlltoAll 實現，只要保證計算加集合通信的分解結果保持唯一性，有太多種解法，類似因式分解。當然也要考慮通信和計算的代價是否合理。

MLP 部分貼一下公式，很好理解。如下圖 4。

Megatron-LM

Megatron-LM 也在 2022 年的第三遍論文中實現了序列並行，但是這個序列並行和我之前介紹的序列維度切分不同，我的理解是在 張量並行維度的再優化。具體可以參考論文：Reducing Activation Recomputation in Large Transformer Models 。

其主要的想法比較好理解，如下圖 5 所示。

張量並行後，需要做相應的 AllReduce 操作得到完整中間結果，如紅色箭頭所示。但 Dropout（部分隨機）和 LayerNorm（隱層維度求和） 不依賴完整結果。故 AllReduce 可以被拆解爲 reduce-scatter 和 AllGather，如下圖 6 所示。

此部分從通信角度拆解完全等價，優點是 Dropout 和 LayerNorm 可以減少一半的運算，配合 checkpoint-activation 使用可以很好的降低顯存使用。

DeepSpeed Ulysses

DeepSpeed-Ulysses 是 DeepSpeed 最近提出的針對長序列模型訓練的解決方案。官方博客：deepspeed-ulysses。DeepSpeed-Ulysses 將各個樣本在序列維度上分割給序列並行需要的的 GPU 上。然後，在 self-attention 計算之前，它對已分割的 查詢 (Q) 、鍵 (K) 和值 (V) 執行 all-to-all 通信操作，以使每個 GPU 接收完整的序列，即通過 all-to-all 通信操作將 Q 、 K 、 V 矩陣進行轉置，把序列切分轉移到轉移頭的均分。這使得參與的 GPU 可以並行計算不同的注意力頭，等價轉換計算序列切分。如下圖 7 所示。

如上圖 7 綠框部分所示，在序列並行組中，原本 （ N/p , d ） 維度在 all-to-all 通信後被轉化爲（ N, d/p ）。由於計算分解性的轉移，原本處理序列並行需要類似 ColossalAI 一樣需要通過 ring 的方式保持結果唯一性，現在只需要分別計算注意力頭就可以了。同時對於一個注意力頭而言 序列並行與原計算無異，其他優秀的第三方優化庫依舊可以使用，原文中也提到了可以與 FlashAttention v2 一起使用。這一點真的是最得我意，在大模型技術蓬勃發展的今天，開源、合作纔是茁壯成長的要素。如果你的優化好，但是不能和別人玩到一起去，真的是隻能孤芳自賞，易用度和成長性會受到制約。

引用下原文的一些其他優點：

缺點：注意力頭的個數需要能整除序列並行數。

這裏再介紹下與 ZeRO-3 的結合（這裏沒有特別細緻看，有問題可以私聊我），DeepSpeed-Ulysses 這裏使用了正交通信，參考 Megatron-DeepSpeed 中通信組建立的的代碼：

num_sequence_data_parallel_groups: int = world_size // sequence_parallel_size // data_parallel_size

即將序列並行和數據並行做成正交的兩個通信組。如下圖 8 所示。

seq_data_group 就是數據序列並行通信組，可以看到 ZeRO-3 是在全局維度做切分，seq_group 是在 seq_data_group 其中的部分 rank 組成。在 seq_group 組內，數據是相同的，只不過是取了序列的不同維度，在數據並行組（ZeRO-3）數據是不同的。

這裏再補充下，DeepSpeed 從 DeepSpeed powers 8x larger MoE model training with high performance 開始，許多通信的重點都在向 all-to-all 靠攏，包括最近發的 ZeRO++ 也是在用層次 all-to-all 去代替 AllReduce。谷歌提出的 Switch Transformer 可能在 MoE 中不是最優的解決方案，但是由此大規模使用的 all-to-all 通信，確實給分佈式訓練帶來了更多的優化可能。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/ylScQOpJ1-ufyPK7X6VUjw

猜你喜歡