【元來如此】第二章——打破序列長度限制，讓無限Token成為可能！
正文共：1929字 6圖
預(yù)計(jì)閱讀時間：3分鐘
作者：思成
軟件生態(tài)中心·應(yīng)用平臺部
“
此前，我們在【元來如此】第一章——大模型技術(shù) · 起航&推理篇中對大模型推理技術(shù)做了一些基礎(chǔ)介紹。
還沒有關(guān)注的同學(xué)快來補(bǔ)課!
(點(diǎn)擊圖片，即可跳轉(zhuǎn)閱讀)
目前主流大模型推理框架比如vLLM、FasterTransformer+Triton Server、TGI（Text Generation Inference）等在Latency（時延）、Throughput（吞吐）、Utilization（利用率）、易用性等方面各有優(yōu)劣。但是當(dāng)前主流大模型只能生成有限長度的文本（比如LLaMA 2 7B模型的限制是2048個Tokens），剩余的部分會被截斷掉。
本文，針對以上長序列問題，我們將介紹Tecorigin在無限長序列推理生成上的一些探索和努力。
我們基于StreamingLLM[1]算法提出了inference-infinitely，并在Teco-Inference-Engine（推理引擎框架）上支持了這種算法，從而賦予了大模型無限長度生成的能力。我們在Baichuan2-7B-Chat上進(jìn)行了驗(yàn)證。
StreamingLLM
圖1 StreamingLLM和其他方法示意圖，引用自[1]
1(a)
如圖1(a)所示，是Transformer模型最基礎(chǔ)的Attention形式，這種Dense的方式在生成第T個Token的時候需要關(guān)注前面0到T-1個Token，做全局的注意力，這意味著0(T2)的時間復(fù)雜度和不斷增加的KV cache顯存占用。
1(b)
1(b) 其改進(jìn)版本，一種Window Attention，這種方式在生成第T個Token的時候只需要關(guān)注前面L個Token即可，這種方式獲得了0(TL)的時間復(fù)雜度，以及可控的KV cache顯存占用，但是也帶來了斷崖式的效果下降。
1(c)
1(c) 另一種改進(jìn)版本，稱作Sliding Window Attention，需要在每次計(jì)算的時候重新計(jì)算長度為L的KV cache，這種方法會有比較好的模型效果，但是也帶來了性能的損失。
1(d)
1(d) StreamingLLM的做法，理解起來很簡單，在1(b)的基礎(chǔ)上，保留一些初始的Token。通過這樣的方式即有0(TL)的時間復(fù)雜度，同時也取得了比較好的模型效果。
StreamingLLM的方法到底為什么生效呢？其實(shí)之前在Lost in the Middle [2] 中的研究已經(jīng)發(fā)現(xiàn)，Dense Attention的方式其注意力會集中在序列中的頭和尾兩個部分，對于中間Token的Attention會比較弱。
圖2 lost in the middle引用自[2]
這種現(xiàn)象在StreamingLLM中也得到了進(jìn)一步的驗(yàn)證：
圖3 除去0和1層的Head之外，其他層都重點(diǎn)關(guān)注頭部部分，引用自[1]
論文中也對此進(jìn)行了一定解釋：要么是開始的Token在模型推理過程中扮演著重要的語義信息，要么是模型出于某種原因，導(dǎo)致了對于這些位置的歸納偏置，即bias。之后論文通過實(shí)驗(yàn)進(jìn)一步驗(yàn)證了后者更可能是主要原因。
基于上面的現(xiàn)象，StreamingLLM中保留了一定的初始Token，通過這種Attention sink的方案穩(wěn)定Attention計(jì)算。之后通過Rolling KV的方式保留一定數(shù)量的最新Token。
如圖4所示。
圖4 Streaming LLM中的KV cache方案，
引用自[1]
到此為止，我們從介紹了StreamingLLM的來龍去脈，下面讓我們深入細(xì)節(jié)，進(jìn)一步思考。
KV cache
上一節(jié)我們多次提到了KV cache這樣的字眼，簡單說：KV cache是一種在自回歸大模型下的優(yōu)化方法，該方法通過緩存一定的過去信息（K和V），用于加速當(dāng)前step模型的推理過程，避免重復(fù)計(jì)算。
詳細(xì)來看，如圖5所示：
Step 1：正常的Attention計(jì)算，之后將K矩陣和V矩陣的結(jié)果緩存下來KcacheVcache。
Step 2：根據(jù)第一步生成的Token，生成當(dāng)前step的Qstep/Kstep/Vstep向量（圖5中Step N藍(lán)色部分），并將Kstep/Vstep和第一步生成的KcacheVcache拼接到一起生成真正的K/V向量，從而進(jìn)行后面的計(jì)算。
圖5 KV cache機(jī)制，引用自[3]
根據(jù)這樣的原理，不難得到KV cache在顯存中的占用：2（K+V）*L（模型層數(shù)）*B（batch大?。?S（序列長度）*D（隱層緯度）*2（fp16字節(jié)），以LLaMA 2 7B模型為例帶入得到2*32*1*4096*4096*2=2GB。
讓我們回想一下剛剛介紹的StreamingLLM算法，當(dāng)超過最大長度之后，需要Rolling KV從而生成下一個Token，這樣勢必會造成大量的KV cache讀寫操作，給本就訪存密集的推理過程進(jìn)一步增加大量額外訪存量，大大降低推理性能。
帶著這樣的疑問，我們深入StreamingLLM作者給出的開源實(shí)現(xiàn)方案。實(shí)際上，從代碼中可以看到是如下流程：
Step 1：通過設(shè)置一個max_gen_len參數(shù)控制了本次推理最大的生成序列長度，并和當(dāng)前prompt長度一起重新開辟了KV cache最大的顯存占用空間。
Step 2：對KV cache空間進(jìn)行一次整體的Rolling。
Step 3：后續(xù)正常推理過程。
可以看到這樣的實(shí)現(xiàn)方法解決了頻繁的KV cache Rolling帶來的訪存問題，但是在每次問答中并不能進(jìn)行無限長度輸出。
Inference-infinitely
基于上述的問題，我們提出了Inference-infinitely，并在Teco-Inference-Engine上進(jìn)行了實(shí)現(xiàn)，提供了每次問答都可以無限輸出的能力。
圖6 Inference-infinitely示意圖
圖6中描述了Inference-infinitely的算法：
Step 1：保留初始Token（黃色）。
Step 2：當(dāng)達(dá)到本次輸出的最大長度之后，將KV cache通過參數(shù)W（圖中為128）進(jìn)行一次整體滑動（藍(lán)色），并生成新的Token（紅色）。
Step 3：每滑動一次提供W個Token生成的空間（灰色），當(dāng)buffer填滿之后，再進(jìn)行Step 2。
通過這樣的算法，很好的解決了上述提到的KV cache Rolling帶來的頻繁訪存，以及每次問答的生成都可以有無限的輸出長度，同時也保證了模型的推理效果。
至此，本文介紹了Tecorigin 在大模型推理——無限長序列生成上的一些探索和努力。未來，會有更多的大模型技術(shù)跟大家一起分享、交流、討論。
免費(fèi)試用申請
如果您對我們的產(chǎn)品感興趣，
可點(diǎn)“http://www.sc-wtoodl.cn/cn/developer.html ”，進(jìn)行試用申請。
參考文獻(xiàn)
[1] [2309.17453] Efficient Streaming Language Models with Attention Sinks (arxiv.org)
[2] [2307.03172] Lost in the Middle: How Language Models Use Long Contexts (arxiv.org)
[3] 使用 FasterTransformer 和 Triton 推理服務(wù)器加速大型 Transformer 模型的推理 - NVIDIA 技術(shù)博客