編譯 | 郭   思

編輯丨陳彩嫻

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

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

最近，上海交通大學(xué)攜手阿里研究團(tuán)隊(duì)針對(duì)這個(gè)問題開展了一項(xiàng)研究。

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

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

1、長(zhǎng)文本處理，眾久難熬

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

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

但這樣一個(gè)不起眼的故障事件，為我們帶來了一個(gè)新的思考，基于大型預(yù)訓(xùn)練語言模型的在線LLM云服務(wù)雖然擁有巨大的超能力，能夠?yàn)橛脩籼峁└咝?、?shí)時(shí)的語言理解和生成能力，但而與之而來的是其對(duì)于算力資源的巨大挑戰(zhàn)。

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

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

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

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

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

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

但這種方法存在幾個(gè)局限性：

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

3、結(jié)語

隨著深度學(xué)習(xí)技術(shù)在自然語言處理領(lǐng)域的廣泛應(yīng)用與深化，端到端性能的研究受到了廣泛關(guān)注。在應(yīng)對(duì)長(zhǎng)文本時(shí)，這種性能飛躍的重要性尤為凸顯，因?yàn)樗苯佑绊懥宋覀兡芊窀咝Ф鴾?zhǔn)確地駕馭海量文本數(shù)據(jù)，并從中抽絲剝繭般提取出有價(jià)值的信息。

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

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

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