阿里與上交大提出 LLM 長文本計算新解法：可處理文本長達 1900k 字節、效率提升 2.4 倍

本文作者：郭思

2024-01-09 12:12

導語：DistAttention與DistKV-LLM結合，帶來云端自然語言處理新變革。

編譯 | 郭思

編輯丨陳彩嫻

在實際應用大模型的過程中，尤其是處理長文本的上下文信息時，如何高效靈活地調度計算資源成為一個學術界與工業界共同關注的問題。

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

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

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

阿里與上交大提出 LLM 長文本計算新解法：可處理文本長達 1900k 字節、效率提升 2.4 倍

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

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

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

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

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

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

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

在模型自回歸過程中，文本如同一幅逐步渲染的畫卷，每一步都根據前序內容迭代地生成新的詞語或token，并將它們融合到當前上下文中作為后續生成的基礎輸入。這種高度動態且連續的過程使得提前精確規劃資源分配成為一項不可預知的任務，從而對設計和優化云端 LLM 服務架構構成了實質性挑戰。

阿里與上交大提出 LLM 長文本計算新解法：可處理文本長達 1900k 字節、效率提升 2.4 倍

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

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

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

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

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

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

此外，對于被置換出的請求所造成的計算中斷可能會導致運行任務性能抖動，從而可能違反對云服務至關重要的服務協議（SLAs）。

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

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

DistAttention將 KV 緩存劃分為rBlocks——統一的子塊，以便為具有長上下文長度的LLM服務分布式計算和內存管理注意力模塊。與主要利用單個節點內的GPU或CPU內存的傳統方法不同，DistAttention允許優化分布式數據中心中所有可訪問的GPU或CPU內存資源，特別是那些現在利用率不高的資源。這不僅支持更長的上下文長度，還避免了與數據交換或實時遷移過程相關的性能波動。

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

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

換言之，DistAttention能夠靈活調度跨數據中心的所有可訪問GPU或CPU內存資源，特別是那些利用率較低的部分，從而不僅支持更長的上下文處理，還能有效降低由于數據交換或實時遷移帶來的性能起伏。

阿里與上交大提出 LLM 長文本計算新解法：可處理文本長達 1900k 字節、效率提升 2.4 倍

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

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

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

相比起 DistAttention，DistKV-LLM更像一位精明的協調員，在數據中心內部妥善管理和優化分布式GPU和CPU之間的KV緩存使用。當一個LLM服務實例因為KV緩存擴大而導致內存不足時，它會主動從負載較小的實例借用額外內存。

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

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

DistAttention與DistKV-LLM雙管齊下，為分布式環境下LLM服務所面臨的資源分配與優化挑戰提供了一個切實有效的解決方案。

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

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

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

阿里與上交大提出 LLM 長文本計算新解法：可處理文本長達 1900k 字節、效率提升 2.4 倍