使用PyTorch加速生成式 AI模型

　　PyTorch 團(tuán)隊(duì)發(fā)布了一篇名為《Accelerating Generative AI with PyTorch II: GPT， Fast》的博文，重點(diǎn)介紹如何使用純?cè)��?PyTorch 加速生成式 AI 模型。

　　正如最近在PyTorch 開(kāi)發(fā)者大會(huì)上宣布的那樣，PyTorch 團(tuán)隊(duì)從頭開(kāi)始編寫(xiě)了一個(gè) LLM，其速度幾乎比基線快10倍，并且沒(méi)有損失準(zhǔn)確性，所有這些都使用本機(jī) PyTorch 優(yōu)化。團(tuán)隊(duì)利用了廣泛的優(yōu)化，包括:

　　Torch.compile:PyTorch 模型的編譯器

　　GPU 量化:通過(guò)降低精度運(yùn)算來(lái)加速模型

　　推測(cè)性解碼:使用小型“草案”模型來(lái)預(yù)測(cè)大型“目標(biāo)”模型的輸出，加速大模型

　　張量并行:通過(guò)在多個(gè)設(shè)備上運(yùn)行模型來(lái)加速模型。

　　而且，更令人驚訝的是，他們可以用不到1000行的本機(jī) PyTorch 代碼來(lái)完成此任務(wù)。

　　具體步驟如下:

　　第1步:通過(guò) TORCH.COMPILE 和靜態(tài) KV 緩存(107.0TOK/S)減少 CPU 開(kāi)銷

　　Torch.compile 允許我們將更大的區(qū)域捕獲到單個(gè)編譯區(qū)域中，特別是在使用 mode=”reduce-overhead” 運(yùn)行時(shí)，對(duì)于減少 CPU 開(kāi)銷非常有效。在這里，我們還指定 fullgraph=True，它驗(yàn)證模型中沒(méi)有“圖形中斷”(即 torch.compile 無(wú)法編譯的部分)。換句話說(shuō)，它確保 torch.compile 充分發(fā)揮其潛力。

　　要應(yīng)用它，我們只需用它包裝一個(gè)函數(shù)(或模塊)即可。

　　torch.compile(decode_one_token，mode="reduce-overhead"，fullgraph=True)

　　然而，這里有一些細(xì)微差別，使得人們通過(guò)將 torch.compile 應(yīng)用于文本生成來(lái)獲得顯著的性能提升有些不簡(jiǎn)單。

　　第一個(gè)障礙是 kv 緩存。kv-cache 是一種推理時(shí)間優(yōu)化，可緩存為先前標(biāo)記計(jì)算的激活(請(qǐng)參閱此處以獲取更深入的解釋)。然而，當(dāng)我們生成更多令牌時(shí)，kv-cache 的“邏輯長(zhǎng)度”就會(huì)增長(zhǎng)。由于兩個(gè)原因，這是有問(wèn)題的。一是每次緩存增長(zhǎng)時(shí)重新分配(和復(fù)制!)kv-cache 的成本非常高。另一個(gè)問(wèn)題是，這種動(dòng)態(tài)使得減少開(kāi)銷變得更加困難，因?yàn)槲覀儾辉倌軌蚶��?cudagraphs 等方法。

　　為了解決這個(gè)問(wèn)題，我們使用“靜態(tài)”kv-cache，這意味著我們靜態(tài)分配 kv-cache 的最大大小，然后屏蔽掉計(jì)算的注意力部分中未使用的值。

　　第二個(gè)障礙是預(yù)填充階段。Transformer 文本生成最好被視為一個(gè)兩階段過(guò)程:1. 處理整個(gè)提示的預(yù)填充，以及2. 自回歸生成每個(gè)標(biāo)記的解碼。

　　盡管一旦 kv-cache 靜態(tài)化，解碼就可以完全靜態(tài)化，但由于提示長(zhǎng)度可變，預(yù)填充階段仍然需要更多的動(dòng)態(tài)性。因此，我們實(shí)際上需要使用單獨(dú)的編譯策略來(lái)編譯這兩個(gè)階段。

　　雖然這些細(xì)節(jié)有點(diǎn)棘手，但實(shí)際實(shí)現(xiàn)起來(lái)一點(diǎn)也不困難(參見(jiàn) gpt-fast)!而且性能的提升是巨大的。

　　突然之間，我們的性能提高了4倍以上!當(dāng)工作負(fù)載受到開(kāi)銷限制時(shí)，這種性能提升通常很常見(jiàn)。

　　步驟2:通過(guò) INT8僅權(quán)重量化緩解內(nèi)存帶寬瓶頸(157.4TOK/S)

　　那么，鑒于我們已經(jīng)看到應(yīng)用 torch.compile 帶來(lái)了巨大的加速，是否有可能做得更好?思考這個(gè)問(wèn)題的一種方法是計(jì)算我們與理論峰值的接近程度。在這種情況下，最大的瓶頸是將權(quán)重從 GPU 全局內(nèi)存加載到寄存器的成本。換句話說(shuō)，每次前向傳遞都要求我們“接觸”GPU 上的每個(gè)參數(shù)。那么，理論上我們能夠以多快的速度“觸及”模型中的每個(gè)參數(shù)?

　　為了衡量這一點(diǎn)，我們可以使用模型帶寬利用率(MBU)。這衡量了我們?cè)谕评磉^(guò)程中能夠使用的內(nèi)存帶寬的百分比。

　　計(jì)算它非常簡(jiǎn)單。我們只需獲取模型的總大小(# params * 每個(gè)參數(shù)的字節(jié)數(shù))并將其乘以每秒可以進(jìn)行的推理數(shù)量。然后，我們將其除以 GPU 的峰值帶寬即可得到 MBU。

　　例如，對(duì)于我們上面的情況，我們有一個(gè)7B 參數(shù)模型。每個(gè)參數(shù)都存儲(chǔ)在 fp16中(每個(gè)參數(shù)2個(gè)字節(jié))，我們實(shí)現(xiàn)了107個(gè)令牌/秒。最后，我們的 A100-80GB 理論內(nèi)存帶寬為2TB/s。

　　將所有這些放在一起，我們得到 **72% MBU!**這相當(dāng)不錯(cuò)，考慮到即使只是復(fù)制內(nèi)存也很難突破85%。

　　但是……這確實(shí)意味著我們非常接近理論極限，并且我們顯然在從內(nèi)存加載權(quán)重方面遇到了瓶頸。我們做什么并不重要——如果不以某種方式改變問(wèn)題陳述，我們可能只能再爭(zhēng)取10% 的性能。

　　讓我們?cè)倏匆幌律厦娴牡仁�。我們無(wú)法真正改變模型中參數(shù)的數(shù)量。我們無(wú)法真正改變 GPU 的內(nèi)存帶寬(好吧，無(wú)需支付更多的錢(qián))。但是，我們可以更改每個(gè)參數(shù)存儲(chǔ)的字節(jié)數(shù)!

　　因此，我們得出了下一個(gè)技術(shù)——int8量化。這里的想法很簡(jiǎn)單。如果從內(nèi)存加載權(quán)重是我們的主要瓶頸，為什么我們不把權(quán)重做得更小呢?

　　請(qǐng)注意，這僅量化權(quán)重 - 計(jì)算本身仍然在 bf16中完成。這使得這種形式的量化易于應(yīng)用，并且精度幾乎沒(méi)有降低。

　　此外，torch.compile還可以輕松生成int8量化的高效代碼。讓我們?cè)俅慰纯瓷厦娴幕鶞?zhǔn)測(cè)試，這次包含了僅 int8權(quán)重量化。

　　從深藍(lán)色線(torch.compile + int8)可以看出，使用torch.compile + int8僅權(quán)重量化時(shí)，性能有顯著提升!而且，淺藍(lán)色線(沒(méi)有torch.compile + int8)實(shí)際上甚至比f(wàn)p16性能還差很多!這是因?yàn)闉榱死��?int8量化的性能優(yōu)勢(shì)，我們需要融合內(nèi)核。這顯示了 torch.compile 的好處之一 - 可以為用戶自動(dòng)生成這些內(nèi)核!

　　將 int8量化應(yīng)用于我們的模型，我們看到性能提高了50%，達(dá)到157.4個(gè)令牌/秒!

　　第3步:使用推測(cè)解碼重新構(gòu)建問(wèn)題

　　即使使用了量化等技術(shù)，我們?nèi)匀幻媾R另一個(gè)問(wèn)題。為了生成100個(gè)代幣，我們必須加載權(quán)重100次。

　　即使權(quán)重被量化，我們?nèi)匀槐仨氁槐橛忠槐榈丶虞d我們的權(quán)重，對(duì)于我們生成的每個(gè)令牌一次!有沒(méi)有辦法解決?

　　乍一看，答案似乎是否定的——我們的自回歸一代存在嚴(yán)格的序列依賴性。然而，事實(shí)證明，通過(guò)利用推測(cè)解碼，我們能夠打破這種嚴(yán)格的串行依賴性并獲得加速!

　　想象一下，您有一位高級(jí)工程師(稱為 Verity)，他做出了正確的技術(shù)決策，但編寫(xiě)代碼的速度相當(dāng)慢。然而，您還有一名初級(jí)工程師(稱為 Drake)，他并不總是做出正確的技術(shù)決策，但可以比 Verity 更快(而且更便宜!) 編寫(xiě)代碼。我們?nèi)绾卫��?Drake(初級(jí)工程師)更快地編寫(xiě)代碼，同時(shí)確保我們?nèi)匀蛔龀稣�確的技術(shù)決策?

　　首先，Drake 經(jīng)歷了編寫(xiě)代碼的勞動(dòng)密集型過(guò)程，并在此過(guò)程中做出技術(shù)決策。接下來(lái)，我們將代碼交給 Verity 進(jìn)行審查。

　　在審查代碼后，Verity 可能會(huì)認(rèn)為 Drake 做出的前3個(gè)技術(shù)決策是正確的，但后2個(gè)需要重做。因此，Drake 回去，放棄了他最后的2個(gè)決定，并從那里重新開(kāi)始編碼。

　　值得注意的是，雖然 Verity(高級(jí)工程師)只看過(guò)一次代碼，但我們能夠生成3段與她編寫(xiě)的代碼相同的經(jīng)過(guò)驗(yàn)證的代碼!因此，假設(shè) Verity 能夠比她自己編寫(xiě)這3段代碼更快地審查代碼，那么這種方法就會(huì)領(lǐng)先。

　　在變壓器推理的背景下，Verity 將由更大的模型發(fā)揮作用，我們希望其輸出用于我們的任務(wù)，稱為驗(yàn)證器模型。同樣，Drake 將由一個(gè)較小的模型來(lái)扮演，該模型能夠比較大的模型(稱為草稿模型)更快地生成文本。因此，我們將使用草稿模型生成8個(gè)令牌，然后使用驗(yàn)證者模型并行處理所有8個(gè)令牌，并丟棄不匹配的令牌。

　　如上所述，推測(cè)解碼的一個(gè)關(guān)鍵特性是它不會(huì)改變輸出的質(zhì)量。只要使用草稿模型生成令牌 + 驗(yàn)證令牌所需的時(shí)間少于生成這些令牌所需的時(shí)間，我們就會(huì)領(lǐng)先。

　　在原生 PyTorch 中完成這一切的一大好處是，這項(xiàng)技術(shù)實(shí)際上非常容易實(shí)現(xiàn)!這是完整的實(shí)現(xiàn)，大約50行原生 PyTorch。

　　盡管推測(cè)性解碼保證我們?cè)跀?shù)學(xué)上與常規(guī)生成相比具有相同的結(jié)果，但它確實(shí)具有運(yùn)行時(shí)性能根據(jù)生成的文本以及草稿和驗(yàn)證器模型的對(duì)齊程度而變化的屬性。例如，當(dāng)運(yùn)行 CodeLlama-34B + CodeLlama-7B 時(shí)，我們能夠在生成代碼時(shí)獲得2倍的令牌/秒提升。另一方面，當(dāng)使用 Llama-7B + TinyLlama-1B 時(shí)，我們只能獲得大約1.3倍的令牌/秒提升。

　　步驟4:使用 INT4量化和 GPTQ (202.1TOK/S) 進(jìn)一步減小權(quán)重的大小

　　當(dāng)然，如果將權(quán)重從16位減少到8位可以通過(guò)減少我們需要加載的字節(jié)數(shù)來(lái)實(shí)現(xiàn)加速，那么將權(quán)重減少到4位將導(dǎo)致更大的加速!

　　不幸的是，當(dāng)權(quán)重減少到4位時(shí)，模型的準(zhǔn)確性開(kāi)始成為一個(gè)更大的問(wèn)題。從我們的初步評(píng)估中，我們看到雖然使用僅 int8權(quán)重量化沒(méi)有明顯的精度下降，但使用僅 int4權(quán)重量化卻有。

　　我們可以使用兩個(gè)主要技巧來(lái)限制 int4量化的精度下降。

　　第一個(gè)是擁有更細(xì)粒度的縮放因子�？紤]縮放因子的一種方法是，當(dāng)我們有量化張量表示時(shí)，它處于浮點(diǎn)張量(每個(gè)值都有縮放因子)和整數(shù)張量(沒(méi)有值有縮放因子)之間的滑動(dòng)比例。例如，對(duì)于 int8量化，我們每行都有一個(gè)縮放因子。然而，如果我們想要更高的精度，我們可以將其更改為“每32個(gè)元素一個(gè)縮放因子”。我們選擇組大小為32來(lái)最小化準(zhǔn)確性下降，這也是社區(qū)中的常見(jiàn)選擇。

　　另一種是使用比簡(jiǎn)單地對(duì)權(quán)重進(jìn)行舍入更先進(jìn)的量化策略。例如，GPTQ等方法利用示例數(shù)據(jù)來(lái)更準(zhǔn)確地校準(zhǔn)權(quán)重。在本例中，我們基于 PyTorch 最近發(fā)布的torch.export在存儲(chǔ)庫(kù)中原型化了 GPTQ 的實(shí)現(xiàn)。

　　此外，我們需要將 int4反量化與矩陣向量乘法融合的內(nèi)核。在這種情況下，torch.compile 不幸地?zé)o法從頭開(kāi)始生成這些內(nèi)核，因此我們?cè)?PyTorch 中利用一些手寫(xiě)的 CUDA 內(nèi)核。

　　這些技術(shù)需要一些額外的工作，但將它們組合在一起會(huì)產(chǎn)生更好的性能!

　　第5步:將所有內(nèi)容組合在一起(244.7TOK/S)

　　最后，我們可以將所有技術(shù)組合在一起以獲得更好的性能!

　　第6步:使用張量并行性

　　到目前為止，我們一直限制自己在單個(gè) GPU 上最大限度地減少延遲。然而，在許多設(shè)置中，我們可以使用多個(gè) GPU。這使我們能夠進(jìn)一步改善延遲!

　　為了直觀地了解為什么這可以讓我們改善延遲，讓我們看一下 MBU 的先驗(yàn)方程，特別是分母。在多個(gè) GPU 上運(yùn)行使我們能夠獲得更多的內(nèi)存帶寬，從而獲得更高的潛在性能。

　　至于選擇哪種并行策略，請(qǐng)注意，為了減少一個(gè)示例的延遲，我們需要能夠同時(shí)在更多設(shè)備上利用內(nèi)存帶寬。這意味著我們需要將一個(gè)令牌的處理拆分到多個(gè)設(shè)備上。換句話說(shuō)，我們需要使用張量并行性。

　　幸運(yùn)的是，PyTorch 還提供了與 torch.compile 組合的張量并行的低級(jí)工具。我們還在開(kāi)發(fā)用于表達(dá)張量并行性的更高級(jí)別的 API，請(qǐng)繼續(xù)關(guān)注!

　　然而，即使沒(méi)有更高級(jí)別的 API，添加張量并行性實(shí)際上仍然很容易。我們的實(shí)現(xiàn)只有150行代碼，并且不需要任何模型更改。

　　我們?nèi)匀荒軌蚶�用前面提到的所有�?yōu)化，所有這些優(yōu)化都可以繼續(xù)與張量并行性組合。將這些組合在一起，我們能夠以55個(gè)令牌/秒的速度為 Llama-70B 提供 int8量化服務(wù)!

　　結(jié)論

　　讓我們看看我們能夠?qū)崿F(xiàn)什么目標(biāo)。

　　簡(jiǎn)單性:忽略量化，model.py(244LOC)+generate.py(371LOC)+tp.py(151LOC)得出766LOC，實(shí)現(xiàn)快速推理+推測(cè)解碼+張量并行。

　　性能:使用 Llama-7B，我們能夠使用編譯 + int4Quant + 推測(cè)解碼來(lái)達(dá)到241tok/s。借助 llama-70B，我們還能夠引入張量并行性以達(dá)到80tok/s。這些都接近或超過(guò) SOTA 性能數(shù)據(jù)!

　　PyTorch 始終保持簡(jiǎn)單性、易用性和靈活性。然而，使用 torch.compile，我們也可以提高性能。