超越GPT-4！騰訊AI新研究打破長文本生成模型限制，序列并行技術再突破
引言：探索長文本生成模型的新趨勢
在人工智能領域，隨著生成模型中上下文長度的不斷增長，我們見證了一種趨勢的形成。從Claude在大型語言模型（LLMs）中將序列長度擴展到100K標記，到OpenAI的GPT-4將上下文長度擴展到128K標記，再到多模態模型的出現，如Gemini 1.5 Pro擁有高達1000萬標記的上下文長度，以及OpenAI的Sora模型支持至少100萬視覺標記。這些突破性進展強調了生成AI技術需要巧妙處理更大上下文長度的重要性。
序列并行（Sequence Parallelism，簡稱SP）作為一種劃分輸入序列的技術，已經成為訓練或推理更長序列的有希望的方法。通過兩年的初步探索，到2023年下半年，兩項標志性的工作——DeepSpeed-Ulysses和Ring-Attention——標志著SP技術的成熟。DeepSpeed-Ulysses在序列長度和計算設備成比例增加時保持恒定的通信量，而Ring-Attention通過計算和通信的重疊隱藏了SP引入的P2P通信成本。然而，挑戰仍然存在，例如DeepSpeed-Ulysses的SP并行度受到注意力頭數的限制，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個標記。此外，多模態模型的出現進一步推動了這一趨勢，例如Gemini 1.5 Pro擁有高達1000萬個標記的上下文長度。
SP技術的關鍵在于它能夠有效地分配和管理大規模輸入序列的計算任務，從而使得模型能夠處理更長的序列而不會受到單個設備內存限制的束縛。這種技術的發展經歷了初期的探索，并在2023年后期通過DeepSpeed-Ulysses和Ring-Attention兩個里程碑式的工作達到成熟。
深入解析 DeepSpeed-Ulysses 和 Ring-Attention
1. DeepSpeed-Ulysses（SP-Ulysses）
DeepSpeed-Ulysses是一種序列并行方法，它通過All2All通信操作處理分割后的Q（查詢）、K（鍵）、V（值）和O（輸出）張量。這種方法的特點是當序列長度和計算設備成比例增加時，通信量保持不變。在All2All操作后，這四個張量的分割從序列維度L轉移到注意力頭數維度hc。這樣，每個注意力頭的softmax(QK^T)V計算可以完整地進行，而不會因為張量的分割而中斷。
2. Ring-Attention（SP-Ring）
Ring-Attention是另一種序列并行方法，它可以看作是FlashAttention的分布式版本。在計算輸出張量O的各個塊時，如果所需的K和V塊不在本地可用，則使用點對點（P2P）通信從其他設備獲取。這種方法的通信可以以環形方式組織，每個設備同時發送和接收K、V塊，允許通信與計算重疊。
這兩種方法各有優勢和局限。DeepSpeed-Ulysses對注意力頭數敏感，其并行度受到注意力頭數的限制，而Ring-Attention在計算和通信效率上存在挑戰。盡管如此，這兩種技術提供了處理長序列的有效途徑，為未來的研究和應用奠定了基礎。
提出統一的序列并行方法
在探索序列并行技術（SP）的發展過程中，DeepSpeed-Ulysses和Ring-Attention作為兩種主要的實現方式各有優勢和局限。DeepSpeed-Ulysses在增加計算設備和序列長度時能保持恒定的通信量，而Ring-Attention通過計算和通信的重疊隱藏了由SP引入的點對點（P2P）通信成本。然而，DeepSpeed-Ulysses的并行度受到注意力頭數的限制，而Ring-Attention在分塊矩陣乘法中的計算效率較低。
為了克服這些限制，我們提出了一種統一的序列并行方法，將DeepSpeed-Ulysses和Ring-Attention的技術結合起來。這種統一方法通過在序列維度上的混合并行策略，允許在不同的計算設備上分別運行DeepSpeed-Ulysses和Ring-Attention，從而優化了通信模式和計算效率。
具體來說，統一的SP方法在一個2D網格的進程組中運行，其中DeepSpeed-Ulysses沿著網格的行進行操作，而Ring-Attention沿著列進行。這種設置不僅提高了模型架構和網絡硬件的適應性，還通過更有效的負載平衡和通信模式，提高了整體的計算性能。
與其他并行技術的比較
在并行計算技術中，數據并行（DP）、張量并行（TP）、ZeRO優化和流水線并行（PP）是常見的方法。每種技術都有其特定的優勢和應用場景，而序列并行（SP）作為一種新興的并行方法，其與這些現有技術的關系和綜合應用同樣重要。
1. 數據并行（DP）與序列并行（SP）數據并行主要通過在多個計算設備上復制模型來并行處理不同的數據批次。SP與DP在激活的通信成本上相當，但SP在注意力模塊中引入了額外的通信開銷。當使用Ring方法時，盡管額外的P2P通信可以與計算重疊，但仍可能引入性能問題。
2. 張量并行（TP）與序列并行（SP）張量并行通過在多個計算設備上分片模型參數來實現并行。與SP相比，TP在通信成本上具有優勢，尤其是在處理長輸入序列時。然而，使用MQA技術可以顯著降低SP的通信成本，使得在某些情況下SP可能比TP更有優勢。
3. ZeRO優化與序列并行（SP）ZeRO通過在多個設備上分片優化器狀態、梯度和參數來減少每個設備的存儲需求。SP可以與ZeRO一起在同一個進程組中運行，這使得SP在內存成本上可以與TP相媲美，特別是當使用ZeRO-3時。
4. 流水線并行（PP）與序列并行（SP）流水線并行通過在不同的計算設備上分配模型的不同層來實現并行，而SP則是在變壓器塊內部分割張量。因此，SP與PP是完全兼容的，可以在4D混合并行系統中共同使用。
通過這些比較，我們可以看到統一的序列并行方法不僅提高了模型的可擴展性和效率，還能與其他并行技術有效結合，為處理大規模和長上下文的生成模型提供了強大的支持。
實驗驗證和性能分析
1. 實驗設置和結果
在實驗中，我們使用了兩個8xA800 GPU節點來測試Unified SP方法在不同序列長度下的性能。根據表3和表4的數據，我們發現在使用8xA100-SXM4 NVLink節點時，當ulysses_degree設置為8時，無論序列長度是32K還是128K，都能達到最高的吞吐量。這驗證了SP-Ulysses在處理大規模序列時的優勢。然而，在8xL20 PCIe節點上，最佳的設置是ulysses_degree為4，ring_degree為2，這表明在不同的硬件配置下，最優的并行度設置可能會有所不同。
2. 性能對比
我們進一步將SP-Unified方法與Megatron-LM中的現有SP實現進行了比較。結果顯示，盡管Megatron-DeepSpeed的SP-Ulysses實現在性能上落后于Megatron-LM，但我們的Unified SP方法在集成了負載平衡策略后，顯示出了更優的性能。特別是在處理大規模序列任務時，Unified SP的表現優于單獨使用SP-Ulysses或SP-Ring的策略。
序列并行技術在大規模集群上的應用前景
1. 大規模應用潛力
隨著序列長度的不斷增加，傳統的并行技術已經無法有效地處理超大規模的序列數據。Unified SP方法通過結合SP-Ulysses和SP-Ring的優點，不僅提高了計算效率，還降低了通信成本，使其成為處理大規模序列的理想選擇。此外，該方法對網絡硬件的要求更為靈活，適用于各種不同的網絡拓撲結構，這對于在大規模集群上部署模型尤為重要。
2. 未來研究方向
盡管當前的實驗結果已經很有希望，但在將序列并行技術應用于超過10,000個GPU的大規模訓練任務時，仍有許多挑戰需要克服。例如，如何進一步優化通信策略以適應不同的網絡條件，以及如何處理更長的序列長度以提高模型的表現。此外，與ZeRO-3和MoE等其他并行技術的結合也是未來研究的重要方向，這可能會開辟新的可能性，以支持更復雜和更大規模的AI模型訓練。
總結
在本文中，我們提出了一種將DeepSpeed-Ulysses和Ring-Attention結合的序列并行方法。這種方法整合了兩種技術的優勢，擴展了其適用性，并在某些情況下提供了更優越的性能。我們系統地分析了序列并行與其他已建立的并行方法之間的相互作用，并得出了一系列最佳實踐結論。這些結論通過在兩個GPU節點上獲取的實驗數據得到了驗證，證實了我們的發現在實際應用中的影響。
通過深入探討序列并行技術，我們不僅解決了長序列輸入在分布式Transformer計算中的挑戰，還通過提出統一的序列并行方法，優化了通信和內存成本。此外，我們的研究強調了序列并行與數據并行、張量并行和管道并行等其他并行策略的兼容性和互補性，為設計混合4D并行系統提供了策略和見解。
實驗結果顯示，統一序列并行方法在處理大規模語言模型時，尤其是在長序列上，不僅提高了計算效率，還優化了網絡通信和內存使用。這一成果為未來在更大規模的集群上部署高效的大模型訓練提供了可能。
我們對序列并行技術的未來發展持續保持樂觀態度，并期待其在推動生成式AI模型處理更長上下文長度方面發揮更大作用。同時，我們也認識到，隨著模型和數據規模的不斷擴大，未來的研究需要進一步探索序列并行技術與其他先進并行技術如專家模型并行的結合，以實現更高效的訓練過程和更優越的模型性能。
本文轉載自AI論文解讀，作者： 柏企 ????
