編譯 | 郭   思

編輯丨陳彩嫻

在實際應(yīng)用大模型的過程中，尤其是處理長文本的上下文信息時，如何高效靈活地調(diào)度計算資源成為一個學(xué)術(shù)界與工業(yè)界共同關(guān)注的問題。

大語言模型所能容納的上下文長度直接影響了諸如 ChatGPT 等高級應(yīng)用與用戶交互體驗的優(yōu)劣程度，這給云環(huán)境下的 LLM 服務(wù)系統(tǒng)提出了嚴峻挑戰(zhàn)：不合理的資源配置不僅可能導(dǎo)致性能瓶頸，還可能造成寶貴的計算資源浪費。

最近，上海交通大學(xué)攜手阿里研究團隊針對這個問題開展了一項研究。

他們提出一個名為 DistAttention 的新穎注意力機制以及一套名為 DistKV-LLM 的分布式 LLM 服務(wù)架構(gòu)，針對長文本語言任務(wù)處理中的計算難題提出了新解法，或是對行業(yè)的啟示。

論文鏈接：https://arxiv.org/pdf/2401.02669.pdf

1、長文本處理，眾久難熬

LLM云服務(wù)是指通過云計算平臺提供的，基于大型語言模型的各項服務(wù)。各家在LLM云服務(wù)之上也下足了馬力。目前市場上主要的 LLM 云服務(wù)提供商包括但不限于亞馬遜的 SageMaker、谷歌的 Cloud AI Platform、微軟的 Azure Machine Learning 以及國內(nèi)的阿里云等。這些平臺通常提供了從模型開發(fā)到部署的一站式服務(wù)，包括計算資源、數(shù)據(jù)存儲、模型訓(xùn)練和部署等。

上個月，一則關(guān)于國內(nèi)某LLM云服務(wù)出現(xiàn)故障的消息在技術(shù)人員之間傳播開來。然而，由于 LLM 云服務(wù)這一概念尚未普及至大眾認知層面，因此該事件在持續(xù)了一天后，便鮮有人再提起。

但這樣一個不起眼的故障事件，為我們帶來了一個新的思考，基于大型預(yù)訓(xùn)練語言模型的在線LLM云服務(wù)雖然擁有巨大的超能力，能夠為用戶提供高效、實時的語言理解和生成能力，但而與之而來的是其對于算力資源的巨大挑戰(zhàn)。

拿阿里云 OpenSearch-LLM 智能問答服務(wù)為例，公開資料顯示，該服務(wù)利用了先進的LLM技術(shù)，在云端提供強大的自然語言處理功能。由于模型運行所需的計算資源波動較大，特別是在處理長上下文場景時對內(nèi)存和計算力的需求激增導(dǎo)致的。這種情況下，傳統(tǒng)的資源分配策略可能無法有效應(yīng)對動態(tài)變化的需求，從而引發(fā)服務(wù)不穩(wěn)定甚至中斷。

可以看出即使強大如阿里云也會受制于 LLM 長文本處理的難題。在 Infinite-LLM 的研究中，揭示了這樣一個現(xiàn)象：LLM服務(wù)的核心運作往往倚賴于多張 GPU 卡的協(xié)同工作以完成復(fù)雜的文本任務(wù)，然而其內(nèi)在的動態(tài)自回歸生成機制猶如一道待解的計算難題。

在模型自回歸過程中，文本如同一幅逐步渲染的畫卷，每一步都根據(jù)前序內(nèi)容迭代地生成新的詞語或token，并將它們?nèi)诤系疆斍吧舷挛闹凶鳛楹罄m(xù)生成的基礎(chǔ)輸入。這種高度動態(tài)且連續(xù)的過程使得提前精確規(guī)劃資源分配成為一項不可預(yù)知的任務(wù)，從而對設(shè)計和優(yōu)化云端 LLM 服務(wù)架構(gòu)構(gòu)成了實質(zhì)性挑戰(zhàn)。

圖注：在采用傳統(tǒng)模型并行策略將網(wǎng)絡(luò)分布在更多GPU上時，這些不隨上下文擴展的層會被過度細粒度地分割

打個比方，這就類似于一位頂級廚師運用最先進的智能廚房設(shè)備，嘗試烹飪一道工序繁復(fù)且需靈活調(diào)整口味的創(chuàng)新菜品。每當加入一種新鮮食材時，都必須依據(jù)現(xiàn)有的風(fēng)味組合重新調(diào)配調(diào)料，而且無法預(yù)見究竟需要多少種類的食材才能成就這道完美的佳肴。

面對這一亟待解決的問題，業(yè)界各方表現(xiàn)出共同的關(guān)注與期待。眾多研發(fā)團隊積極投入研究，其中 PagedAttention等方案試圖通過改進 GPU 與 CPU 之間的數(shù)據(jù)交換策略，有效地管理和調(diào)度內(nèi)存資源，以期化解上述困擾LLM服務(wù)效率提升的棘手問題。

但這種方法存在幾個局限性：

首先，PagedAttention的內(nèi)存置換范圍局限于單一節(jié)點內(nèi)的GPU和CPU內(nèi)存，因此對極長上下文長度的支持能力受限；

其次，盡管其分頁策略旨在最小化內(nèi)存碎片，但基于請求級別整體交換KV緩存（Key-Value，鍵值緩存，是一種計算機存儲技術(shù)），錯失了在分布式云環(huán)境中實現(xiàn)更為靈活、精細化調(diào)度的機會；

此外，對于被置換出的請求所造成的計算中斷可能會導(dǎo)致運行任務(wù)性能抖動，從而可能違反對云服務(wù)至關(guān)重要的服務(wù)協(xié)議（SLAs）。

2、雙管齊下帶來的新突破

為了解決業(yè)界長期面臨的大規(guī)模語言模型（LLM）服務(wù)中的內(nèi)存管理與資源分配難題，阿里與上海交大的團隊提出了一種名為 DistAttention 的新型注意力算法。

DistAttention將 KV 緩存劃分為rBlocks——統(tǒng)一的子塊，以便為具有長上下文長度的LLM服務(wù)分布式計算和內(nèi)存管理注意力模塊。與主要利用單個節(jié)點內(nèi)的GPU或CPU內(nèi)存的傳統(tǒng)方法不同，DistAttention允許優(yōu)化分布式數(shù)據(jù)中心中所有可訪問的GPU或CPU內(nèi)存資源，特別是那些現(xiàn)在利用率不高的資源。這不僅支持更長的上下文長度，還避免了與數(shù)據(jù)交換或?qū)崟r遷移過程相關(guān)的性能波動。

這就像一位技藝高超的倉儲大師，巧妙地將一個不斷擴展的、宛如巨大食材倉庫的KV緩存分割成大小適中的rBlocks儲物箱，使得在面對一道配料繁多、制作復(fù)雜的超長菜單（相當于處理長上下文任務(wù)）時，每一種“食材”（數(shù)據(jù)）都能迅速而準確地送達各自的烹飪臺（分布式計算節(jié)點）。

與那些只在單一廚房（單個GPU或CPU內(nèi)存節(jié)點）內(nèi)調(diào)配食材的傳統(tǒng)方法相比，這位“倉儲大師”更擅長調(diào)動整個美食廣場（即數(shù)據(jù)中心內(nèi)的所有可用GPU和CPU內(nèi)存資源），特別是那些閑置或使用率低的空間，使制作超長菜單變得可行且高效，避免了因頻繁搬運食材造成的混亂和效率波動。

換言之，DistAttention能夠靈活調(diào)度跨數(shù)據(jù)中心的所有可訪問GPU或CPU內(nèi)存資源，特別是那些利用率較低的部分，從而不僅支持更長的上下文處理，還能有效降低由于數(shù)據(jù)交換或?qū)崟r遷移帶來的性能起伏。

圖注：展示了DistKV-LLM如何解決內(nèi)存碎片化問題

基于此，Infinite-LLM團隊進一步開發(fā)了集成 DistAttention 技術(shù)的 DistKV-LLM 分布式LLM服務(wù)引擎。

DistKV-LLM 是一個與 DistAttention 無縫集成的分布式LLM服務(wù)引擎。DistKV-LLM 擅長管理KV緩存，有效地在數(shù)據(jù)中心內(nèi)的分布式 GPU 和 CPU 之間協(xié)調(diào)內(nèi)存使用。當一個LLM服務(wù)實例由于 KV 緩存擴展而面臨內(nèi)存不足時，DistKV-LLM主動從負擔較輕的實例尋求補充內(nèi)存。

相比起 DistAttention，DistKV-LLM更像一位精明的協(xié)調(diào)員，在數(shù)據(jù)中心內(nèi)部妥善管理和優(yōu)化分布式GPU和CPU之間的KV緩存使用。當一個LLM服務(wù)實例因為KV緩存擴大而導(dǎo)致內(nèi)存不足時，它會主動從負載較小的實例借用額外內(nèi)存。

同時，DistKV-LLM還引入了一套精細的通信協(xié)議，促進云端運行的多個LLM服務(wù)實例之間進行高效、擴展性強且一致的互動協(xié)作。這套協(xié)議的核心目標是高效管理和平衡龐大的內(nèi)存資源，并優(yōu)先關(guān)注數(shù)據(jù)存儲位置就近性和通信效率提升，這對于解決與KV緩存分布存儲相關(guān)的性能瓶頸至關(guān)重要。

這意味著DistKV-LLM能夠更好地駕馭大型語言模型在眾多GPU和CPU上的并行運算。當LLM服務(wù)因需處理海量信息而面臨內(nèi)存壓力時，DistKV-LLM能智慧地從負載較輕的區(qū)域獲取額外內(nèi)存，并制定一套高級協(xié)同規(guī)則，確保不同云上LLM實例間能夠高效有序、步調(diào)一致地完成工作。這一系列設(shè)計的關(guān)鍵在于優(yōu)化內(nèi)存使用、確保數(shù)據(jù)快速存取以及減少各部分間的通信成本，即使面臨復(fù)雜的分布式存儲挑戰(zhàn)，也能保障系統(tǒng)的整體高性能表現(xiàn)。

DistAttention與DistKV-LLM雙管齊下，為分布式環(huán)境下LLM服務(wù)所面臨的資源分配與優(yōu)化挑戰(zhàn)提供了一個切實有效的解決方案。

在具體的實驗測評之中，DistAttention與DistKV-LLM在資源管理方面也有卓越的表現(xiàn)。

研究人員在一個包含4個節(jié)點和32塊GPU的集群上部署了DistKV-LLM系統(tǒng)。每個節(jié)點配備了8塊NVIDIA A100（80GB）GPU。模型方面則選擇了一個具有代表性的模型LLaMA2 進行評估（LLaMA2系列包含了三個不同規(guī)模的模型：7B、13B和70B。）。

團隊對分布式系統(tǒng)配置進行了廣泛測試，涵蓋了從2個到32個實例的多種規(guī)模。在評測過程中，采用了包含上下文長度最高達1,900K字節(jié)的18個代表性基準數(shù)據(jù)集進行嚴格檢驗。結(jié)果顯示，相較于當前最尖端的技術(shù)，系統(tǒng)成功實現(xiàn)了1.03至2.4倍的顯著端到端性能飛躍，并且在處理更長上下文信息方面表現(xiàn)出色，支持的上下文長度可擴展至原先的2至19倍之多。同時，在處理標準長度上下文任務(wù)場景下，系統(tǒng)的吞吐量也取得了令人矚目的提升，增長幅度在1.4至5.3倍之間。

雷峰網(wǎng) 雷峰網(wǎng) 雷峰網(wǎng)(公眾號：雷峰網(wǎng))

3、結(jié)語

隨著深度學(xué)習(xí)技術(shù)在自然語言處理領(lǐng)域的廣泛應(yīng)用與深化，端到端性能的研究受到了廣泛關(guān)注。在應(yīng)對長文本時，這種性能飛躍的重要性尤為凸顯，因為它直接影響了我們能否高效而準確地駕馭海量文本數(shù)據(jù)，并從中抽絲剝繭般提取出有價值的信息。

DistAttention與DistKV-LLM的結(jié)合，通過智能管理內(nèi)存資源、優(yōu)化分布式計算策略，成功解決了大規(guī)模語言模型服務(wù)在長上下文處理中的難題，使系統(tǒng)能夠從容應(yīng)對超長文本序列，同時保持端到端性能提升以及上下文長度擴展能力。未來云端自然語言處理應(yīng)用有望迎來全新的突破與變革。

本文作者長期關(guān)注大模型計算與框架技術(shù)發(fā)展與行業(yè)動態(tài)，歡迎添加作者微信（lionceau2046）交流，互通有無。

雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。