Offloading Inference 主要在探討如何讓 GPU 與其他裝置一起協同推論，例如有些運算放在 CPU，有些記憶體暫存在硬碟裡面。這類的方法通常用在參數量大到不行的模型上，例如我的 GPU 記憶體可能就 10GB 而已，想要使用參數量 30B 的模型進行推論，光是把模型讀進記憶體裡面都是個問題，更不用說是進行推論了。
正當 GPU 努力工作的時候，看了一眼在旁邊喝咖啡聊是非的 CPU & Disk 們，心中鐵定很不是滋味。這些裝置也是有能力負擔一些 Tensor 的存放，甚至能做些運算，閒置在一旁總是不太好。今天我們透過 FlexGen 這個套件來一同看看 Offloading Inference 的概念。
(Powered By Microsoft Designer)
Decoder LM 的 Inference 可以分成兩個階段：
Prefill 階段，模型會對最一開始的輸入進行第一次推論，這時除了預測下一個 Token 的機率以外，還會產生一個很大的初始 KV 快取。
接下來進入 Decoding 階段，這個階段模型會 Token By Token 的做文本生成，每次推論都會產生一組新的 KV 快取，並與原本的 KV 快取串接在一起，所以 KV 快取會越來越大。
在這個過程中，主要的記憶體消耗為模型權重本身以及 KV 快取，其中 KV 快取佔用的記憶體會越來越大，最終超越模型權重本身佔用的記憶體。
FlexGen 將這個 Offloading Generative Inference 視為一種圖形拜訪問題，在這邊一個方格就是一個 Batch，一樣顏色的方塊代表模型的同一個 Layer。
在這個拜訪問題裡面有以下限制：
多數的 Offloading 系統都是以 Row-By-Row 的方式完成計算，因為這樣可以最快完成一筆推論，然後把那個超肥大的 KV 快取丟掉，再繼續進行下一筆推論。
但因為模型沒辦法被完整讀取進 GPU 記憶體裡面，所以這些 Layers 是輪流被放進 GPU 運算的，因此 Row-By-Row 的做法會不斷重複讀取模型權重，造成記憶體讀寫很大的負擔。
但如果我們以 Col-By-Col 的方式做計算，就可以減少模型權重的 IO 負擔。這時 IO 問題就會轉移到隱藏層狀態與 KV 上面。上圖的每個方格在輪流運算時，會發生六件事情：
這些操作基本上都是互相獨立的，因此可以同時進行來縮短需要的時間，這個技巧被稱為 Overlapping。
Offloading Inference 對 Tensor 存放的位置比較講究，切分 Tensor 的方式有很多種粒度 (Granularity)，例如：
較粗的粒度 (Coarser Granularity) 會有比較小的運行開銷 (Runtime Overhead)，例如模型都不切，能夠整份放進 GPU 記憶體，那就不用一直讀讀寫寫，直接算完就好，但是在 Offloading Inference 裡面沒辦法這麼理想。
較細的粒度 (Finer Granularity) 會有更好的靈活度，FlexGen 更偏好細粒度的切法，因此在模型權重上採用 Layer Granularity，而在隱藏層狀態和 KV 快取上採用 Tensor Granularity。
雖然 CPU 比 GPU 慢很多，但是在計算 Attention 時還是有幫助的，考慮以下情況：
KV 快起的大小為 (B, S, H, 4) 位元組，而隱藏層狀態僅為 (B, H, 4) 位元組。
註：因為一次只 Decode 一個 Token，所以隱藏層狀態實際上的長度 S = 1。
理論上 KV 快取會比隱藏層狀態大很多，所以算上記憶體搬運的時間，當 Sequence Length 越長，搬進 CPU 運算與搬進 GPU 運算的時間就會差越多。因此 FlexGen 在規劃運算流程時，也會考量到裝置的 IO 速度。
成本模型 (Cost Model) 是在估算某組設定下所需要的花費時間，基本上就是：
這裡成本模型的概念很簡單，就是以上兩者加在一起，這個時間成本被定義為 T。
因為每個 Layer 的運算裡面，記憶體的 Load/Unload 與實際運算都是平行的，所以 Latency 就是這幾個操作裡面最慢的那個。
張量只會在 (GPU <=> CPU) 與 (CPU <=> Disk) 之間做傳輸，因此在 GPU 裡的記憶體不會直接傳進硬碟裡面。
將問題定義與成本模型建立之後，就可以開始針對一系列的參數搜尋最佳的 Offloading 運算策略。首先定義以下參數：
bls 代表 Block Sizegbs 代表 GPU Batch Sizeg, c, d 分別代表 GPU, CPU, Diskwg, wc, wd 代表模型權重在三個裝置上的放置比例hg, hc, hd 代表隱藏層狀態在三個裝置上的放置比例cg, cc, cd 代表 KV 快取在三個裝置上的放置比例目標是找到一組參數，可以使 T / bls 最小，且符合以下限制：
wg + wc + wd = 1
hg + hc + hd = 1
cg + cc + cd = 1
至此 FlexGen 將 Offloading Inference 變成一個線性規劃 (Linear Programming) 的問題，而這個問題並不難解。首先列舉 bls, gbs 的組合，因為 gbs 為 4 的倍數，且 bls 小於 20，所以能列舉的組合並不多。
接著拿一些 Sample Data Point 去實際估算硬體的速度參數，就能求得完整的一組策略參數。有了這組策略參數後，就可以開始進行推論啦！
量化能減少模型權重佔用的記憶體，在 Offloading Inference 裡面也是如此。但 FlexGen 不僅對模型權重兩化，甚至連 KV 快取都被量化了！系統直接使用 Min/Max Quantization 轉換成 INT4，FlexGen 實驗證實這樣做依然可以維持很好的 PPL 或準確率。
將權重與 KV 快取量化為 INT4 後，可以減少四倍 (FP16 => INT4) 的 IO 量。
這其實令筆者相當驚訝，因為直接做 INT4 量化通常會帶來很大的誤差，這裡不僅將模型量化成 INT，連 KV 快取都變成 INT4，這樣的結果卻不會跟原本的模型差很多，也許這就是 LLM 相當 Robust 的展現吧。
註：若是有做 CPU Delegation 就不做 Quantization。
稀疏注意力 (Sparse Attention) 的概念相對單純，注意力機制會針對一個序列上的每個 Token 算出一個注意力分數 (Attention Score)，用來代表當前位置要給該 Token 多少「注目」。而 Sparse Attention 是在算完 Attention Score 之後，只取 Top-K 分數最高的分數繼續做運算，其他 Token 的分數就通通丟掉掉。
論文作者使用 Google Cloud 的 NVIDIA T4，以 OPT 6.7B, 30B, 175B 三種規模的 LLM 做測試，並使用合成的資料集，每個 Prompt 都被 Padding 到相同長度，然後計算整個模型的吞吐量做為評估公式。
對標的 Baseline 系統為 DeepSpeed ZeRO, Hugging Face Accelerate 與 Petals。
以下是主要的實驗結果：
其他系統多是輸在無法將模型放入 GPU，所以必須使用原本就很慢的 CPU Offloading。或者是很勉強的放進一張 GPU 裡面，但是 Batch Size 超級小，導致吞吐量極低。使用 Pipeline 平行化可以大幅提昇吞吐量，尤其是生成的 Tokens 越多時，Pipeline 平行化的優勢就會展現出來。
更令筆者感到興趣的是 Quantization 的實驗：
其中 FP16 是原本的權重，4-Bit 代表有量化，而 4-Bit-S 代表有量化且加上 Sparse Attention。可以看到無論有沒有加上 Sparse Attention，4-Bit Quantization 都跟 FP16 差距沒有很大。
看完論文並瞭解 FlexGen 的工作原理後，我們就來實際操作看看吧！首先使用 pip 安裝 FlexGen 套件：
pip install flexgen
建議開個新的 Conda 環境來安裝，這東西大概裝了十幾分鐘才裝完。
接著筆者借了一台 GPU 只有 12GB 記憶體的電腦來跑以下指令：
export TRANSFORMERS_CACHE="./Models/Cache"
python3 -m flexgen.flex_opt \
--model facebook/opt-30b --percent 0 100 100 0 100 0
因為那台電腦的 Home 目錄所在硬碟快被塞爆了，所以把 TRANSFORMERS_CACHE 指定到另外一顆硬碟上。開始執行後，那台電腦就用他孱弱的網速開始努力下載 OPT-30B 的模型權重並進行格式轉換。
(Few Days Later ...)
經過兩個小時的奮鬥之後，程式終於跑完了！得到以下結果：
Outputs:
----------------------------------------------------------------------
0: Paris is the capital city of France and the most populous city in the country. It is the second largest city in the European Union after London. Paris is also the seat of the French government
----------------------------------------------------------------------
3: Paris is the capital city of France and the most populous city in the country. It is the second largest city in the European Union after London. Paris is also the seat of the French government
----------------------------------------------------------------------
TorchDevice: cuda:0
cur_mem: 0.0079 GB, peak_mem: 4.4778 GB
TorchDevice: cpu
cur_mem: 56.5031 GB, peak_mem: 0.0000 GB
model size: 55.803 GB cache size: 2.789 GB hidden size (p): 0.029 GB
peak gpu mem: 4.478 GB projected: False
prefill latency: 41.840 s prefill throughput: 48.948 token/s
decode latency: 300.872 s decode throughput: 0.412 token/s
total latency: 342.713 s total throughput: 0.373 token/s
對一個只有 12GB 記憶體的 Pascal GPU 來說，平均 2.7 秒處理一個 Token 應該算不錯吧？至少他跑起來了 XD
我想這就是理論與現實之間的差距，當我們把過多的 IO 負擔丟向別的地方時，就會在其他地方產生更多成本出來，其中就包含時間成本。
看完實驗結果後，真正引起筆者興趣的反而是這個 Petals 套件，這個套件由 BigScience 所開發，使用 BitTorrent-Style 的 Offloading 框架。首先安裝此套件：
pip install petals
這個套件裝起來就滿快了，然後使用以下程式碼做測試：
from petals import AutoDistributedModelForCausalLM as ModelCls
from transformers import AutoTokenizer as TkCls
from transformers import TextStreamer
# 微調過的 Llama 2 70B
model_name = "petals-team/StableBeluga2"
model: ModelCls = ModelCls.from_pretrained(model_name)
tokenizer: TkCls = TkCls.from_pretrained(model_name)
streamer = TextStreamer(tokenizer)
prompt = "This is a long story, "
tokens = tokenizer(prompt, return_tensors="pt")
input_ids = tokens["input_ids"]
# 開始進行文本生成
outputs = model.generate(
input_ids,
max_new_tokens=32,
streamer=streamer,
)
可以到此網頁確認哪些模型可用，看到這網頁就不難理解為何官方說他們是 "BitTorrent-Style" 的了。
雖然是跑 70B 的模型，但過程只需要下載其中一個 Shard 就能跑（約 1GB 左右），因此很快就能開始進入生成環節。接著就會看到程式努力的吐出生成 Tokens，約 1 秒多可以生成 1 個 Token。
其運作原理是把大家的 GPU 運算能力集合起來，把 Shard 切細一點，我們一人跑一點點，輪流接力做運算，過程中透過網路來傳輸中間層的運算結果。這真的是很有意思的想法，雖然很慢但至少能跑。其中也有 Falcon 180B 這種模型能跑，各位可以去玩看看。
今天不知為何開始很熱血的在看論文，並跟大家分享了一下 FlexGen 的技術細節，看完這篇 Paper 後，筆者對整個 Offloading Framework 也有比較初步的瞭解與概念。可惜的是現在大公司都有幾萬張 A100 工作站，且小型研究單位也能充分利用量化技術的情況下，真的需要使用 Offloading 的情境實在太少見了。從 FlexGen 最後的 Commit 時間來看，這一塊在 LLM 領域裡面相對不活躍一些。
對於底層推論框架的介紹大致到這邊告一段落，接下來將會談論一些開發 LLM 相關應用時，會用到的一些技巧，例如 In-Context Learning, Retrieval Model 等等，明天見啦！
iThome鐵人賽
看影片追技術