引言：探索長文本生成模型的新趨勢

在人工智能領域，隨著生成模型中上下文長度的不斷增長，我們見證了一種趨勢的形成。從Claude在大型語言模型（LLMs）中將序列長度擴展到100K標記，到OpenAI的GPT-4將上下文長度擴展到128K標記，再到多模態(tài)模型的出現(xiàn)，如Gemini 1.5 Pro擁有高達1000萬標記的上下文長度，以及OpenAI的Sora模型支持至少100萬視覺標記。這些突破性進展強調了生成AI技術需要巧妙處理更大上下文長度的重要性。

序列并行（Sequence Parallelism，簡稱SP）作為一種劃分輸入序列的技術，已經(jīng)成為訓練或推理更長序列的有希望的方法。通過兩年的初步探索，到2023年下半年，兩項標志性的工作——DeepSpeed-Ulysses和Ring-Attention——標志著SP技術的成熟。DeepSpeed-Ulysses在序列長度和計算設備成比例增加時保持恒定的通信量，而Ring-Attention通過計算和通信的重疊隱藏了SP引入的P2P通信成本。然而，挑戰(zhàn)仍然存在，例如DeepSpeed-Ulysses的SP并行度受到注意力頭數(shù)的限制，Ring-Attention由于矩陣乘法的細分而降低了計算效率。這些限制目前阻礙了序列并行在分布式Transformer計算中的更廣泛應用。

論文標題、機構、論文鏈接和項目地址

論文標題: A Unified Sequence Parallelism Approach for Long-Context Generative AI

機構: Tencent

論文鏈接: https://arxiv.org/pdf/2405.07719.pdf

項目地址: https://github.com/feifeibear/long-context-attention

序列并行技術（SP）的概述

序列并行技術（Sequence Parallelism，簡稱SP）是一種將輸入張量的序列維度分割到多個計算設備上的技術，這種技術對于解鎖生成性AI模型的長上下文能力至關重要。SP通過在多個計算設備上分配輸入序列，使得每個設備處理輸入序列的一個子段，從而支持更長的序列處理能力。

近年來，隨著生成性AI模型上下文長度的不斷增長，SP技術的重要性日益凸顯。例如，Claude在大型語言模型（LLMs）中將序列長度擴展到了100K個標記，而OpenAI的GPT-4則將上下文長度擴展到了128K個標記。此外，多模態(tài)模型的出現(xiàn)進一步推動了這一趨勢，例如Gemini 1.5 Pro擁有高達1000萬個標記的上下文長度。

SP技術的關鍵在于它能夠有效地分配和管理大規(guī)模輸入序列的計算任務，從而使得模型能夠處理更長的序列而不會受到單個設備內(nèi)存限制的束縛。這種技術的發(fā)展經(jīng)歷了初期的探索，并在2023年后期通過DeepSpeed-Ulysses和Ring-Attention兩個里程碑式的工作達到成熟。

深入解析 DeepSpeed-Ulysses 和 Ring-Attention

1. DeepSpeed-Ulysses（SP-Ulysses）

DeepSpeed-Ulysses是一種序列并行方法，它通過All2All通信操作處理分割后的Q（查詢）、K（鍵）、V（值）和O（輸出）張量。這種方法的特點是當序列長度和計算設備成比例增加時，通信量保持不變。在All2All操作后，這四個張量的分割從序列維度L轉移到注意力頭數(shù)維度hc。這樣，每個注意力頭的softmax(QK^T)V計算可以完整地進行，而不會因為張量的分割而中斷。

2. Ring-Attention（SP-Ring）

Ring-Attention是另一種序列并行方法，它可以看作是FlashAttention的分布式版本。在計算輸出張量O的各個塊時，如果所需的K和V塊不在本地可用，則使用點對點（P2P）通信從其他設備獲取。這種方法的通信可以以環(huán)形方式組織，每個設備同時發(fā)送和接收K、V塊，允許通信與計算重疊。

這兩種方法各有優(yōu)勢和局限。DeepSpeed-Ulysses對注意力頭數(shù)敏感，其并行度受到注意力頭數(shù)的限制，而Ring-Attention在計算和通信效率上存在挑戰(zhàn)。盡管如此，這兩種技術提供了處理長序列的有效途徑，為未來的研究和應用奠定了基礎。

提出統(tǒng)一的序列并行方法

在探索序列并行技術（SP）的發(fā)展過程中，DeepSpeed-Ulysses和Ring-Attention作為兩種主要的實現(xiàn)方式各有優(yōu)勢和局限。DeepSpeed-Ulysses在增加計算設備和序列長度時能保持恒定的通信量，而Ring-Attention通過計算和通信的重疊隱藏了由SP引入的點對點（P2P）通信成本。然而，DeepSpeed-Ulysses的并行度受到注意力頭數(shù)的限制，而Ring-Attention在分塊矩陣乘法中的計算效率較低。

為了克服這些限制，我們提出了一種統(tǒng)一的序列并行方法，將DeepSpeed-Ulysses和Ring-Attention的技術結合起來。這種統(tǒng)一方法通過在序列維度上的混合并行策略，允許在不同的計算設備上分別運行DeepSpeed-Ulysses和Ring-Attention，從而優(yōu)化了通信模式和計算效率。

具體來說，統(tǒng)一的SP方法在一個2D網(wǎng)格的進程組中運行，其中DeepSpeed-Ulysses沿著網(wǎng)格的行進行操作，而Ring-Attention沿著列進行。這種設置不僅提高了模型架構和網(wǎng)絡硬件的適應性，還通過更有效的負載平衡和通信模式，提高了整體的計算性能。

與其他并行技術的比較

在并行計算技術中，數(shù)據(jù)并行（DP）、張量并行（TP）、ZeRO優(yōu)化和流水線并行（PP）是常見的方法。每種技術都有其特定的優(yōu)勢和應用場景，而序列并行（SP）作為一種新興的并行方法，其與這些現(xiàn)有技術的關系和綜合應用同樣重要。

1. 數(shù)據(jù)并行（DP）與序列并行（SP）數(shù)據(jù)并行主要通過在多個計算設備上復制模型來并行處理不同的數(shù)據(jù)批次。SP與DP在激活的通信成本上相當，但SP在注意力模塊中引入了額外的通信開銷。當使用Ring方法時，盡管額外的P2P通信可以與計算重疊，但仍可能引入性能問題。

2. 張量并行（TP）與序列并行（SP）張量并行通過在多個計算設備上分片模型參數(shù)來實現(xiàn)并行。與SP相比，TP在通信成本上具有優(yōu)勢，尤其是在處理長輸入序列時。然而，使用MQA技術可以顯著降低SP的通信成本，使得在某些情況下SP可能比TP更有優(yōu)勢。

3. ZeRO優(yōu)化與序列并行（SP）ZeRO通過在多個設備上分片優(yōu)化器狀態(tài)、梯度和參數(shù)來減少每個設備的存儲需求。SP可以與ZeRO一起在同一個進程組中運行，這使得SP在內(nèi)存成本上可以與TP相媲美，特別是當使用ZeRO-3時。

4. 流水線并行（PP）與序列并行（SP）流水線并行通過在不同的計算設備上分配模型的不同層來實現(xiàn)并行，而SP則是在變壓器塊內(nèi)部分割張量。因此，SP與PP是完全兼容的，可以在4D混合并行系統(tǒng)中共同使用。

通過這些比較，我們可以看到統(tǒng)一的序列并行方法不僅提高了模型的可擴展性和效率，還能與其他并行技術有效結合，為處理大規(guī)模和長上下文的生成模型提供了強大的支持。

實驗驗證和性能分析

1. 實驗設置和結果

在實驗中，我們使用了兩個8xA800 GPU節(jié)點來測試Unified SP方法在不同序列長度下的性能。根據(jù)表3和表4的數(shù)據(jù)，我們發(fā)現(xiàn)在使用8xA100-SXM4 NVLink節(jié)點時，當ulysses_degree設置為8時，無論序列長度是32K還是128K，都能達到最高的吞吐量。這驗證了SP-Ulysses在處理大規(guī)模序列時的優(yōu)勢。然而，在8xL20 PCIe節(jié)點上，最佳的設置是ulysses_degree為4，ring_degree為2，這表明在不同的硬件配置下，最優(yōu)的并行度設置可能會有所不同。

2. 性能對比

我們進一步將SP-Unified方法與Megatron-LM中的現(xiàn)有SP實現(xiàn)進行了比較。結果顯示，盡管Megatron-DeepSpeed的SP-Ulysses實現(xiàn)在性能上落后于Megatron-LM，但我們的Unified SP方法在集成了負載平衡策略后，顯示出了更優(yōu)的性能。特別是在處理大規(guī)模序列任務時，Unified SP的表現(xiàn)優(yōu)于單獨使用SP-Ulysses或SP-Ring的策略。

序列并行技術在大規(guī)模集群上的應用前景

1. 大規(guī)模應用潛力

隨著序列長度的不斷增加，傳統(tǒng)的并行技術已經(jīng)無法有效地處理超大規(guī)模的序列數(shù)據(jù)。Unified SP方法通過結合SP-Ulysses和SP-Ring的優(yōu)點，不僅提高了計算效率，還降低了通信成本，使其成為處理大規(guī)模序列的理想選擇。此外，該方法對網(wǎng)絡硬件的要求更為靈活，適用于各種不同的網(wǎng)絡拓撲結構，這對于在大規(guī)模集群上部署模型尤為重要。

2. 未來研究方向

盡管當前的實驗結果已經(jīng)很有希望，但在將序列并行技術應用于超過10,000個GPU的大規(guī)模訓練任務時，仍有許多挑戰(zhàn)需要克服。例如，如何進一步優(yōu)化通信策略以適應不同的網(wǎng)絡條件，以及如何處理更長的序列長度以提高模型的表現(xiàn)。此外，與ZeRO-3和MoE等其他并行技術的結合也是未來研究的重要方向，這可能會開辟新的可能性，以支持更復雜和更大規(guī)模的AI模型訓練。

總結

在本文中，我們提出了一種將DeepSpeed-Ulysses和Ring-Attention結合的序列并行方法。這種方法整合了兩種技術的優(yōu)勢，擴展了其適用性，并在某些情況下提供了更優(yōu)越的性能。我們系統(tǒng)地分析了序列并行與其他已建立的并行方法之間的相互作用，并得出了一系列最佳實踐結論。這些結論通過在兩個GPU節(jié)點上獲取的實驗數(shù)據(jù)得到了驗證，證實了我們的發(fā)現(xiàn)在實際應用中的影響。

通過深入探討序列并行技術，我們不僅解決了長序列輸入在分布式Transformer計算中的挑戰(zhàn)，還通過提出統(tǒng)一的序列并行方法，優(yōu)化了通信和內(nèi)存成本。此外，我們的研究強調了序列并行與數(shù)據(jù)并行、張量并行和管道并行等其他并行策略的兼容性和互補性，為設計混合4D并行系統(tǒng)提供了策略和見解。

實驗結果顯示，統(tǒng)一序列并行方法在處理大規(guī)模語言模型時，尤其是在長序列上，不僅提高了計算效率，還優(yōu)化了網(wǎng)絡通信和內(nèi)存使用。這一成果為未來在更大規(guī)模的集群上部署高效的大模型訓練提供了可能。

我們對序列并行技術的未來發(fā)展持續(xù)保持樂觀態(tài)度，并期待其在推動生成式AI模型處理更長上下文長度方面發(fā)揮更大作用。同時，我們也認識到，隨著模型和數(shù)據(jù)規(guī)模的不斷擴大，未來的研究需要進一步探索序列并行技術與其他先進并行技術如專家模型并行的結合，以實現(xiàn)更高效的訓練過程和更優(yōu)越的模型性能。

本文轉載自AI論文解讀，作者： 柏企 ????