一念 LLM 大語言模型推理加速
一、大語言模型概要介紹
首先來看一下大語言模型的結(jié)構(gòu)。在 Transformer 結(jié)構(gòu)下的大語言模型推理的過程中,一個(gè) token 或者一個(gè)字的生成的過程大致上可以分成兩步:
- Step 1: 根據(jù)已有信息,也就是 input 的已知信息,估計(jì)下一個(gè) token 的概率分布。
- Step 2: 根據(jù)采樣的策略,從概率分布里面挑出最有可能的下一個(gè) token。
這個(gè)過程有可能是以概率最大的,偏 greedy 的方式來做,要考慮到后期生成的 token 的概率,從總體上去做采樣。這是跟傳統(tǒng)深度學(xué)習(xí)推理不太一樣的地方。這兩步是一個(gè)循環(huán)的過程,當(dāng)生成了下一個(gè) token 之后,這個(gè) token 會進(jìn)到下一步 Step 1 里面去再生成再下一個(gè) token,這就是推理的基本邏輯。
這里引出一個(gè)經(jīng)常提到的概念,KVCache。剛才提到的在 step 1 的時(shí)候,是根據(jù)已有信息,這里的已有信息包含兩個(gè)意思,一個(gè)是原始的輸入,另一個(gè)是之后生成的 token。如果我們把一個(gè)生成的過程拆開,前面部分是最原始的輸入,生成第一個(gè) token A。第二步從概率分布的邏輯上來說其實(shí)是要把前面的部分再加上 A 去估計(jì)下一個(gè) token,依此循環(huán)。這會導(dǎo)致一個(gè)問題,計(jì)算量是在不停增長的,而且會與前面已生成的部分和 input 部分成正比,可以想象到這樣的邏輯一定會越跑越慢。
在 Transformer 結(jié)構(gòu)里面,存在一個(gè)計(jì)算的特性,當(dāng)前 token 的結(jié)果只與前面的 token 有關(guān),可以把前面 token 的計(jì)算結(jié)果進(jìn)行緩存,形成兩個(gè)階段:
- Prefill 階段:輸入后走一遍全部的過程,這是全量的走模型的過程,走完之后,會產(chǎn)生一些中間結(jié)果。這些中間結(jié)果被緩存起來,放入到圖中標(biāo)紅的下一步的過程中,KVCache 在進(jìn)入 attention 之前,跟現(xiàn)有的新生成的 token 的結(jié)果做一個(gè) concat,然后再做計(jì)算。之后又是一個(gè) token 生成的過程。
- Decoding 階段:通過 KVCache 的優(yōu)化,decoding 階段的計(jì)算量和前面的 token 數(shù)就變得無關(guān)了。這里其實(shí)是一個(gè)近似的無關(guān)。因?yàn)樵谄渌饕牟糠侄际菬o關(guān)的,但是在 attention 計(jì)算的地方,是被恢復(fù)成了一個(gè)全長的 token,然后進(jìn)行 attention。
這就是 KVcache 的存在的意義,讓 decoding 階段的計(jì)算盡量復(fù)用以前已經(jīng)計(jì)算過的結(jié)果,這樣對前部分的數(shù)量就沒有依賴,從而提高整體的推理速度。
這里存在的問題是,首先,開始輸入很多 input,之后每次只會輸入一個(gè),比如這里輸進(jìn)去的是一本書,可能是成百上千萬的輸入,之后計(jì)算的數(shù)全部都是1。所以在推理的時(shí)候就會出現(xiàn)這樣的一種現(xiàn)象,比如一張 A100 的卡,它能夠并行推理的 token 數(shù)跟 GPU 的 TFLOPS 的關(guān)系是,當(dāng) token 數(shù)增長的時(shí)候,GPU 會被更充分地利用起來,到達(dá)一個(gè)閾值之后,可能跟算值本身的時(shí)限有關(guān),基本上到達(dá)了該 GPU 卡能夠提供的最大 TFLOPS。
我們發(fā)現(xiàn)在 prefill 的運(yùn)轉(zhuǎn)區(qū)間,可以把 GPU 壓滿。但是因?yàn)楹罄m(xù) decoding 的并行每一次只預(yù)測了一個(gè) token 數(shù),也就是并行度非常小,在生成過程中,GPU 都處于一種不飽和的工作狀態(tài)。對于不飽和,最簡單的處理方式就是做 batch,把 batch 加大。這里面存在一個(gè)問題,如何才能把 batch size 加大?由于 KVCache 的存在,而且由于在大模型的情況下,KV cache 占用顯存非常厲害,batch size 會受到顯存的限制。
我們需要看一看,顯存到底是被怎么消耗掉的。一個(gè)正常的執(zhí)行過程,包括了 prefill 和 decoding 兩個(gè)階段。一個(gè)模型加載之后會占用一部分顯存,之后第一把執(zhí)行 input token 的時(shí)候,顯存會有個(gè)快速的消耗,之后隨著 token 的逐步生成,顯存的消耗在慢慢變大,這和常規(guī)的深度學(xué)習(xí)的推理非常不一樣。因?yàn)樗?prefill 過程和后面的生成的過程。這個(gè)消耗過程其一是跟長度有關(guān)系,其二是跟我們加的輸入的長度也有關(guān)系,當(dāng)我們 batch size 擴(kuò)大的時(shí)候,這里的顯存就會成倍地上漲。
從顯存角度來看待的話,可以列個(gè)很簡單的公式,首先是模型占用了多少參數(shù),然后在模型的推理過程中有很多的中間變量其實(shí)也會占用一部分參數(shù),另外有多少 token 的 KVCache 緩存也會占一部分,最終是要小于顯存的大小。
正常而言,比如一個(gè) llama-13B 的模型,只要這種模型一上去,顯存基本就固定了。如果我們想優(yōu)化的是把 batch size 做大,就可以通過優(yōu)化 β,將顯存使用變小。這里主要涉及到 KVCache 的量化技術(shù)。
占用了 KV cache 的這些 token 的數(shù)量跟什么相關(guān)呢?可以這樣去列一個(gè)公式,實(shí)際上是 batch 內(nèi)平均的 token 的數(shù)量。比如現(xiàn)在已經(jīng)輸入了多少 token,以及生成了多少 token 的一個(gè)總數(shù)乘以 batch size。這里還存在一個(gè) γ 系數(shù),其實(shí)就是 batch 之中不同的請求之間的 token 可復(fù)用的 KVCache。
一念 LLM 框架的名字是取自一念三千,現(xiàn)在的大模型轉(zhuǎn)瞬之間就會生成各種不同的結(jié)果,而一念在佛教里面就是一剎那的意思。一念三千也代表我們的目標(biāo)是希望大模型在一瞬之間生成世間萬象。
因此這里會對 latency 和 throughput 兩個(gè)方向重點(diǎn)優(yōu)化。
二、一念 LLM 基本框架
一念 LLM 的基本框架如上圖所示。從上往下分別為:
最上層就是咱們經(jīng)常實(shí)現(xiàn)的如 Llama、Baichuan、QWen 等模型,與常規(guī)的深度學(xué)習(xí)模型的推理方案不一樣,我們采用的是手寫模型。我們拋棄了計(jì)算圖,計(jì)算圖以前都是用角度算值,拼成模型,這最大的好處是有很多的靈活性,便于算法人員去實(shí)現(xiàn)。而這樣會帶來另外一個(gè)問題,即深度優(yōu)化困難,以至于最后大家會走向一個(gè)方向,去做融合算子,然后把融合算子放到圖里面去,如果符合圖的 pattern,用融合算子去替換。
對于 Transformer 結(jié)構(gòu),因?yàn)榻Y(jié)構(gòu)非常簡單,并且現(xiàn)在結(jié)構(gòu)基本已開始收斂,大家不再卷模型結(jié)構(gòu),更多的是卷模型效果。也就是說,這一塊我們需要實(shí)現(xiàn)的模型的類型是比較少的,這就讓手寫模型和手寫算子變得有利可圖。
英偉達(dá)、Facebook 等各個(gè)大廠為大語言模型場景寫的算子都非常的大。甚至像 attention 這種,按現(xiàn)在基本就是一個(gè)標(biāo)準(zhǔn)的,用 Flash Attention 做一個(gè)大算子來做。
我們?yōu)槭裁匆獜氐讈G掉計(jì)算圖呢?因?yàn)槲覀冞€希望去優(yōu)化整個(gè)模型推理過程中的顯存,只要已經(jīng)規(guī)定好了模型,就可以盯著這個(gè)模型,仔仔細(xì)細(xì)地把里面顯存的使用全部去調(diào)好,使得整個(gè)顯存使用最小化。省出來的顯存都可以拿去做其他事,去做 KVCache 相關(guān)的事。
另外一個(gè)是高效調(diào)度以提高吞吐,后文中還會詳細(xì)介紹。
下面就是算子擇優(yōu),我們的底層算子,很多時(shí)候會是開源的封裝。因?yàn)槠鋵?shí)當(dāng)模型結(jié)構(gòu)相對固定之后,硬件廠商會專門針對這些大的算子進(jìn)行優(yōu)化,提高性能,最終目的是賣掉他們的卡?,F(xiàn)在各大廠商甚至為了不同的 tensor 的大小,去寫不同的算子,從而獲得收益的。
多硬件的支持方面,現(xiàn)在業(yè)界的主流框架,基本上都是英偉達(dá),當(dāng)然也有支持英特爾的 CPU。然而像支持 GPU 和華為昇騰等國產(chǎn)卡,還并不完善。從國內(nèi)廠商的角度來看都會面臨一個(gè)問題就是高性能的 GPU 卡進(jìn)不來。從業(yè)務(wù)安全的角度,我們也必須要去支持不同的硬件。
但是為什么不是按不同廠商使用不同框架,比如英偉達(dá)用一個(gè)框架,華為用華為的框架,兩邊的都能用到一個(gè)好的收益?實(shí)際上,當(dāng)你在做這一層優(yōu)化的時(shí)候,你會發(fā)現(xiàn)面臨不同的框架,需要重復(fù)做。另外各個(gè)框架本身對后期業(yè)務(wù)邏輯的適配也會不同。
我們通過上面統(tǒng)一框架,下面支持多種硬件,這樣就相當(dāng)于做到調(diào)度這層一次優(yōu)化在所有硬件上都可以用,這樣更有利于整個(gè)平臺的運(yùn)營。
三、一念 LLM 框架調(diào)度
第一個(gè)問題是 ContinuousBatching 和 PageAttention,從最原始的公式來看它其實(shí)就是在有效地優(yōu)化 batch size。
正常情況下,我們會把不同的請求打成一個(gè) batch 去做 GPU 的推理,這個(gè)過程往往是一個(gè) batch 一個(gè) batch 地進(jìn)去,要等到這個(gè) batch 里面的請求全部處理完才退出。從 GPU 的任務(wù)調(diào)度上來說這是最簡單的使用方式,但這種使用方式最大的問題在于它的有效 batch 會越來越低,因?yàn)槊總€(gè)請求的輸入長度不一樣,輸出長度不一樣。比如有可能第一個(gè)請求是一個(gè)摘要的任務(wù),會丟入一個(gè) 1,000 字的文章讓大模型用一句話總結(jié)出來;第二個(gè)請求可能是一個(gè)擴(kuò)寫的任務(wù),給了一小段話,讓模型把它擴(kuò)寫成一篇長文。也就說輸入輸出的不匹配,會導(dǎo)致有些請求很快就結(jié)束了,有些請求還要跑很久。這樣的話,要等到所有的請求都完成,這個(gè) batch 才會退出,這會導(dǎo)致有效的 batch 到后面越來越小。另外一個(gè)問題是,后面很多請求結(jié)束了,GPU 算力就會閑置。
因此很容易想到的是能否當(dāng)一個(gè)請求處理完成后,就把另外新的請求輸入進(jìn)去,這就是 ContinuousBatching 的基本思想。這個(gè)想法其實(shí)早在 22 年就已經(jīng)被微軟提出來,但是這么好的一個(gè)想法一直都沒爆發(fā),是因?yàn)榇嬖谝粋€(gè)問題,就是它的 KVCache 的顯存的操作成本比較高,比較麻煩。去年由伯克利提出的 PageAttention解決了這一問題,于是整個(gè) ContinuousBatching+PageAttention 的機(jī)制就迎來了一個(gè)爆發(fā)期,現(xiàn)在成為了大語言模型推理框架的一個(gè)標(biāo)配功能。
簡單講因?yàn)檎麎K操作成本高,所以將它切小,它采用了常規(guī)操作系統(tǒng)里面的內(nèi)存分頁的機(jī)制來管理顯存,把 KVCache 切成不同的塊,用指針去引用的方式來進(jìn)行計(jì)算。這樣就顯著降低了顯存操作的粒度。這樣 ContinuousBatching 的整體操作效率就得到了提升。
這里面還遺留了一個(gè)問題,就是 input 在請求與請求之間的共享問題,這在以前的深度學(xué)習(xí)的推理里面很少關(guān)注到,但是在推薦里面從很早的時(shí)候就已經(jīng)有類似的工作。比如一個(gè)用戶多個(gè) item 的 rank 操作,因?yàn)橛脩舳际且粯拥?,所以這一部分的推理只需要做一次。然后不同的 item 的推理多推多次,再融合到一起,再去做后面的推理工作。
請求是有共性的,這些共性的請求其實(shí)可以只算一次,然后把計(jì)算結(jié)果緩存,就可以把 KVCache 存起來,等到下一個(gè)請求,只需要處理后面的。Batch 之間也可以繼續(xù)復(fù)用,這樣整個(gè)后期請求的推理響應(yīng)就可以一下提上去。
KVCache 機(jī)制看起來非常美好,但也是有成本的,因?yàn)?KVCache 會占用顯存,而且一旦緩存大,就會面臨 cache 換入換出和命中問題。比如放兩個(gè)前序在顯存里面,就得占兩份顯存。最終在具體的執(zhí)行節(jié)點(diǎn)上面,命中率就決定了這個(gè)機(jī)制最終的收益。所以 prefix catching 更多就相當(dāng)于上面公式里面的 γ。
我們在服務(wù)節(jié)點(diǎn)的前置加了一個(gè)叫 prefix-token 的路由器。在這里,路由需要平衡兩件事情,命中率以及傳統(tǒng)路由的問題,包括負(fù)載平衡、容災(zāi)等問題。例如,如果前序都是一樣的,就直接打到一個(gè)節(jié)點(diǎn)上去,它一定會被命中。但是實(shí)際上還是會面臨負(fù)載平衡和容災(zāi)的事情怎么解決的問題。所以我們構(gòu)建了一個(gè)路由表,例如經(jīng)??吹降囊恍┙巧缪莸姆?wù),比如正在跟宋江聊,首先需要告訴模型你現(xiàn)在正在扮演的是宋江,宋江是一個(gè)什么樣的人、之前的生平、有什么樣的能力、他的性格等一些重要的事件,就像用戶簡歷一樣。然后再說你跟宋江之前聊過什么,下面請生成你準(zhǔn)備作為宋江回復(fù)給用戶的信息。這個(gè)過程里面有大量的信息其實(shí)是跟用戶 profile 一樣。還有另外一種,比如作為愛因斯坦或者牛頓這樣的角色,對不同的角色,在進(jìn)到模型之前,我們就會把前面的這一段做一個(gè) prefix-token。對相同的一段給他指定一個(gè)具體的路由表,在有限的機(jī)器里面去選中一個(gè)機(jī)器集合。
另外剛才也提到一個(gè)問題,我們需要解決不能有太多份的 cache,cache 份數(shù)多了,顯存里就全是 cache,沒有顯存去計(jì)算當(dāng)前要執(zhí)行的請求了。我們通過另一個(gè)維度的管理,對于單一的節(jié)點(diǎn),相當(dāng)于是一個(gè) server-set 其實(shí)是有限的。從這張圖可以感覺到,最后每個(gè)節(jié)點(diǎn)只命中了兩個(gè) set,從而達(dá)到對于單個(gè)節(jié)點(diǎn)而言 cache 份數(shù)是相對可控的,同時(shí)又能夠在 set 的維度完成負(fù)載平衡和容災(zāi)。
最后再講一下,CPU 跟 GPU 的混合推理,其實(shí)就是優(yōu)化 M。
前面提到把模型從 FP16 變成 INT8 或者 INT4,這種是量化操作,效果是有損的。從框架層面需要支持,待業(yè)務(wù)評估是否使用。這里講的是無損的一些方式,計(jì)算密度的問題。我們一般講的大語言模型是計(jì)算強(qiáng)密集的,通常指的是 Transformer 結(jié)構(gòu)部分。實(shí)際上它還有一個(gè) token embedding 部分,這部分就是查表,類似推薦里面的 sparse 操作。往往這部分又是放在 GPU 上做的,這個(gè)表一旦大了就會占顯存。而且可以看到一般在業(yè)務(wù)應(yīng)用的時(shí)候,通常都會擴(kuò) token,擴(kuò)詞表。因?yàn)樵嫉哪P屠锩娴哪切┰~并不能完全覆蓋到我們的業(yè)務(wù)里。業(yè)務(wù)里面有可能會有一些特殊場景的詞。開源模型如標(biāo)準(zhǔn)的 Llama-13B 在 v2 的時(shí)候是 3.2 萬的詞表。詞表會隨業(yè)務(wù)擴(kuò)展,Llama-13B 原生的詞表可能只占整個(gè)參數(shù)的百分之一。但是當(dāng)把它擴(kuò)到三十萬的詞表時(shí),它對整個(gè)模型參數(shù)量的占比就會到 11.8% 的顯存。而它又是個(gè) sparse 操作,很直接的一個(gè)想法,就是把它丟到 CPU 上去改。我們自己在測試 30 萬詞表的 Llama-13B 時(shí),能夠有 10% 的性能提升。但這 10% 也不一定是能拿到的收益,跟詞表大小有直接相關(guān)。因?yàn)檫@會涉及到 CPU 跟 GPU 的聯(lián)合推理的問題,意味著 CPU 執(zhí)行完的結(jié)果要拷貝到 GPU 上,多了一個(gè)成本。如果節(jié)省的顯存不足以去 cover 成本,收益就可能是負(fù)的。所以需要根據(jù)實(shí)際業(yè)務(wù)所用的模型去調(diào)整是否選用這個(gè)機(jī)制。
四、一念 LLM 在 GR 模型的應(yīng)用
目前大家都在考慮大語言模型的 Transformer,生成式的結(jié)構(gòu),能否用在推薦系統(tǒng)中。推薦系統(tǒng)的推理因?yàn)榱糠浅4螅臅r(shí)要求非常高,推理成本比常規(guī)的 AI 推理要高很多。
當(dāng)我們把常規(guī)的模型結(jié)構(gòu),變成一個(gè)大語言模型的結(jié)構(gòu),采用一個(gè)長序列輸入,計(jì)算量可能是成千上萬倍,成本也可能隨之增加。
生成式推薦其實(shí)是基于歷史序列去預(yù)測候選 item 的 action。這里單個(gè)用戶會有大量的 item 要預(yù)測,因?yàn)?rank 往往是千級別的。對于同一用戶的歷史序列,其實(shí)是一樣的,這是很好的符合 prefix catching 的場景。可以對這個(gè)用戶所有的 item 預(yù)測請求,做一次 prefix-cache,之后只做 item 部分的推理。這可以實(shí)現(xiàn)原來整個(gè)計(jì)算量是 prefix-token 加上一個(gè)待預(yù)測的 item,需要推理的 token 量乘以 item 的數(shù)量。相當(dāng)于每一個(gè) token 的計(jì)算量是乘的關(guān)系,所以會導(dǎo)致我們的計(jì)算量成千上萬倍的增長,因?yàn)?token 是成千上萬的。通過這個(gè)功能,就可以實(shí)現(xiàn) item 變成了 item 的數(shù)量加上 token 數(shù)量,也就是把乘變成了加。這樣最后的計(jì)算量就是跟 item 數(shù)量線性相關(guān),就跟現(xiàn)在正常推薦的推理相似。
這里只是解決了一個(gè)計(jì)算量的問題,還有另一個(gè)問題是 latency 的問題。如果是以萬級的 token 輸入,想要最后控制在 10 毫秒以下也是非常困難的,哪怕業(yè)務(wù)能夠接受更長時(shí)間的 latency,也不是將閾值從 10 放松到 50 毫秒這樣的一個(gè)狀態(tài),而是要放松到秒級。這里對于 item 的預(yù)測是可以分開的。在不知道 item 的情況下,就已經(jīng)知道了用戶的序列,可以提前計(jì)算。比如在用戶請求剛剛過來時(shí),就可以把序列發(fā)給用戶了,然后等到把 item 做了召回初排之后,再去執(zhí)行這一部分。這部分就只有一個(gè) item 的 token 耗時(shí),這也是我們傳統(tǒng)意義上講的 rank 的請求耗時(shí)。這個(gè)耗時(shí)就可以做到只跟 token 的最后 item 有關(guān)而跟前面的 prefix-token 的數(shù)量無關(guān)。這樣的話就可以把整個(gè)系統(tǒng)跟現(xiàn)有推薦的推理系統(tǒng)基本對齊。
無、未來規(guī)劃
未來的規(guī)劃就是圍繞整個(gè)架構(gòu)的幾層:
1. 對模型的支持
常用的大語言模型;業(yè)務(wù)的 GR 的推薦模型的支持。
2. 調(diào)度層面的優(yōu)化
計(jì)算/顯存的流水線,包括現(xiàn)在熱門的投機(jī)解碼等業(yè)界先進(jìn)技術(shù),會持續(xù)跟進(jìn)。
3. 硬件的支持
不只是在硬件上跑起來,更重要的是在硬件上定制化算子的開發(fā)。例如華為等國內(nèi)其他公司也在做加速類的硬件。要單獨(dú)提一下 CPU,因?yàn)楝F(xiàn)在最新的英特爾芯片,已經(jīng)在 CPU 盒里面加進(jìn)了矩陣計(jì)算的硬件單元。這種情況下,CPU 其實(shí)是可以去承擔(dān)一定程度的高密度的矩陣計(jì)算的,性能會好很多。
六、Q&A
Q1:之前在做 CTR 推理的時(shí)候做了很多類似顯存分配、動態(tài) batch、多流并行、kernel launch 等工作,未來有哪些 CTR 推理的能力和經(jīng)驗(yàn)是可以借鑒到 LLM 推理上的?CTR 和 LLM 之間的區(qū)別和優(yōu)化側(cè)重點(diǎn)有哪些共性和不同點(diǎn)?
A1:這里沒有提到傳統(tǒng)的深度學(xué)習(xí)推理里面的那些優(yōu)化,準(zhǔn)確來說這些優(yōu)化全部都有效。只是在大語言模型推理場景下面,因?yàn)殚L序列,序列輸入是以 token 方式去做并行這樣一個(gè)特殊性,引入了一些特殊的優(yōu)化方法,包括動態(tài) batch,其實(shí)很大程度上也會是跟剛才提到的 continuous batching 結(jié)構(gòu)類似,實(shí)際上 continuous batching 就是更細(xì)化的動態(tài) batch 操作。
多流的并形其實(shí)可以用在單個(gè)請求單個(gè) batch 前向推理的過程優(yōu)化里面。這部分相當(dāng)于已經(jīng)有一定的 batch 了,要去生成一個(gè) token,就要經(jīng)過圖的過程。所有以前用的優(yōu)化都可以繼續(xù)使用。
只是可能有一部分,比如手寫算子,圖優(yōu)化,因?yàn)闆]有圖了,所以也就不需要圖優(yōu)化了。
簡而言之,目前的 GR 模型其實(shí)并不是一個(gè)連續(xù)生成的模型,其實(shí)對 KVCache 連續(xù)生成這一點(diǎn)上的依賴沒那么重。就像在現(xiàn)有體系下是計(jì)算圖去實(shí)現(xiàn)。走一次前項(xiàng),再走一次前項(xiàng),只不過多了一個(gè) KVCache 的輸入。然后圖存在一些變化的部分,用傳統(tǒng)計(jì)算圖的一些優(yōu)化方式去實(shí)現(xiàn)也都是沒問題的。
這個(gè)理論可能會涉及到極致優(yōu)化的問題,因?yàn)槭謱懰阕訕O致優(yōu)化這件事情本身就是相當(dāng)于損失易用性來實(shí)現(xiàn)的。這張圖上的應(yīng)用,是跟咱們算法同學(xué)配合上線的。在大語言模型的場景下面,如果有一個(gè)新模型,就又得去支持它。從工程層面上來說,需要重新把從訓(xùn)練到推理的這一套,全部的給重新弄一遍,這其實(shí)就是取決于模型結(jié)構(gòu)本身的穩(wěn)定性的問題。因?yàn)樵诖笳Z言模型場景,結(jié)構(gòu)已經(jīng)非常穩(wěn)定,跟推薦場景比起來簡單太多了。推薦場景會有很多的 sparse,小算子來回拼來拼去的事情。在大語言模型這種情況下全都沒有了,只剩下大算子。
像 kernel launch 這種類型的優(yōu)化,在現(xiàn)有的大語言模型場景下?全部都可以用?,F(xiàn)在的主流的大語言模型推理框架基本上都不是用圖。除了 tensorRT LLM 算是用圖的,但其實(shí)它里面也是很多大算子,只是用圖大致串了一下幾個(gè)大算子。
Q2:一念現(xiàn)在用的是連續(xù)的 batch,在大 batch 中間,模型的 forward 部分和解碼采樣是會有一個(gè)串形的執(zhí)行流水線,這樣是不是可能會出現(xiàn) GPU 的空隙?
A2:實(shí)際上在解碼采樣這件事情上,目前主要的優(yōu)化手段還是把解碼采樣的時(shí)間段盡量放短,因?yàn)楝F(xiàn)在主要的方式仍然是 token 依次生成。當(dāng)前大語言模型,因?yàn)轱@存的問題,顯存占用太大了,比較難像以前一樣。這個(gè) batch 跟那個(gè) batch 交換著做一個(gè)流水線。
現(xiàn)在盡可能將一個(gè) batch 打大,打大之后顯存就已經(jīng)沒了,沒法再切到另一個(gè) batch 的顯存來做下一步的推理,所以現(xiàn)在這一塊基本上是串行?,F(xiàn)在主要的優(yōu)化是想辦法把解碼采樣這一環(huán)節(jié)壓縮,把它壓的更短,不要就那個(gè)地方的解碼采樣,現(xiàn)在很多都在最后全部采用 GPU 算子來做,而不是在 python 層面去寫,寫了再用。其實(shí)都是為了極致的壓縮,因?yàn)樗@整個(gè)過程現(xiàn)在就是個(gè)串行的,比以前的優(yōu)化的空間小很多,不能夠讓流水線讓下一個(gè) batch 的計(jì)算能夠進(jìn)來,還需等前一個(gè)解碼采樣算完,希望盡可能的用 GPU 去算解碼采樣。
現(xiàn)在主流的方案基本都是這樣的,因?yàn)橹虚g提到的 decoding 是有狀態(tài)的,其中的 KVCache 不可能在這么短的時(shí)間內(nèi)(一個(gè)采樣的時(shí)間)把它換出去,再換另外一個(gè) batch 進(jìn)來。這樣的操作會得不償失。
Q3:一念有沒有做過跟 TensorRT-llm 的對比,是否有跟 A800 或者 4090 做推理。它的性價(jià)比如何?
A3:A800 或者 A100 的有測,其他普通的非線上卡沒有測。跟 tensorrt-llm 的對比的話,在具體場景,有 10%-20% 的收益。這個(gè)收益主要源于我們在下面對開源算子進(jìn)行了封裝,包括 FastTransformer 的算子,vllm 的算子和 TensortRT 的算子。我們集成了開源的大語言模型推理框架,以及為業(yè)務(wù)定制的算子。這里面存在一個(gè)問題,因?yàn)樽隽舜罅康娜诤纤阕佣加?FP16 在推理,在這種情況下 FP16 有精度損失。大家都知道在業(yè)務(wù)實(shí)際應(yīng)用的時(shí)候,有些業(yè)務(wù)可能會有非常強(qiáng)的效果一致性的需求,有時(shí)候就需要把算子退化到比如 pytorch 的算子來跟訓(xùn)練做對齊。
對于 TensorLRT-LLM 來說,它其實(shí)完全用 FP16 的性能并不是那么好,只有開啟 INT8 和 FP16 的混合量化的機(jī)制之后性能才上來。
Q4:如果是對于 GR 生成推理的話是更適合做一個(gè) multi stream 的并發(fā),還是也像 llm 那樣做一個(gè)異步大 batch 之后,再去做一個(gè)串行執(zhí)行?
A4:在大語言模型情況下為什么沒有做這件事情,是因?yàn)轱@存不夠。但是在 GR 的場景,模型大小不是很大,像現(xiàn)在做的大語言模型,13B 的規(guī)模那就是 26G 的顯存沒了。但如果模型可能只有 G 這種級別,剩下的顯存就是足夠大的,就有空間去做多 batch。對于現(xiàn)在大語言模型的場景,很多問題都暴露在了顯存大小這件事情上。