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

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

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

論文標(biāo)題、機(jī)構(gòu)、論文鏈接和項(xiàng)目地址

論文標(biāo)題: A Unified Sequence Parallelism Approach for Long-Context Generative AI

機(jī)構(gòu): Tencent

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

項(xiàng)目地址: https://github.com/feifeibear/long-context-attention

序列并行技術(shù)（SP）的概述

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

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

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

深入解析 DeepSpeed-Ulysses 和 Ring-Attention

1. DeepSpeed-Ulysses（SP-Ulysses）

DeepSpeed-Ulysses是一種序列并行方法，它通過All2All通信操作處理分割后的Q（查詢）、K（鍵）、V（值）和O（輸出）張量。這種方法的特點(diǎn)是當(dāng)序列長度和計(jì)算設(shè)備成比例增加時(shí)，通信量保持不變。在All2All操作后，這四個(gè)張量的分割從序列維度L轉(zhuǎn)移到注意力頭數(shù)維度hc。這樣，每個(gè)注意力頭的softmax(QK^T)V計(jì)算可以完整地進(jìn)行，而不會(huì)因?yàn)閺埩康姆指疃袛唷?/p>

2. Ring-Attention（SP-Ring）

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

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

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

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

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

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

與其他并行技術(shù)的比較

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

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

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

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

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

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

實(shí)驗(yàn)驗(yàn)證和性能分析

1. 實(shí)驗(yàn)設(shè)置和結(jié)果

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

2. 性能對(duì)比

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

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

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

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

2. 未來研究方向

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

總結(jié)

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

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

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

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

本文轉(zhuǎn)載自AI論文解讀，作者： 柏企 ????