特斯拉前AI總監(jiān)Andrej Karpathy：大模型有內(nèi)存限制，這個(gè)妙招挺好用！

Andrej Karpathy 用通俗易懂的語(yǔ)言介紹了 speculative execution。

「如今，LLM（大語(yǔ)言模型）并不是單點(diǎn)突破的 —— 而是需要多個(gè)重要組件有效協(xié)同工作的系統(tǒng)。Speculative decoding 是幫助我們從系統(tǒng)角度思考的一個(gè)很好的例子?！箰?ài)丁堡大學(xué)博士生符堯表示道。

符堯上述觀點(diǎn)評(píng)論的是特斯拉前 AI 總監(jiān)、年初重回 OpenAI 的 Andrej Karpathy 剛剛發(fā)布的一條推特。

人形機(jī)器人公司 1X Technologies 的 AI 副總裁 Eric Jang 評(píng)價(jià)道：「Karpathy 很好的解釋了 LLM 的 speculative execution。其他自回歸模型可能會(huì)以類(lèi)似的方式加速。連續(xù)（擴(kuò)散）模型可能從 K 步中獲益較少（可能在第 1 步后偏離猜測(cè)），但可以將其應(yīng)用于 VQ-latents 的離散代碼?！?/p>

看完上述評(píng)價(jià)，我們大概也了解了，Karpathy 說(shuō)的「Speculative execution」，這是優(yōu)化技術(shù)的一類(lèi)，采用這個(gè)技術(shù)的計(jì)算機(jī)系統(tǒng)會(huì)根據(jù)現(xiàn)有信息，利用空轉(zhuǎn)時(shí)間提前執(zhí)行一些將來(lái)可能用得上，也可能用不上的指令。如果指令執(zhí)行完成后發(fā)現(xiàn)用不上，系統(tǒng)會(huì)拋棄計(jì)算結(jié)果，并回退執(zhí)行期間造成的副作用（如緩存）。

為了讓大家更好的理解 Karpathy 的內(nèi)容。我們先介紹一下「Speculative decoding」方法，對(duì)后續(xù)理解更加有益，其主要用于加速大模型的推理。據(jù)了解，GPT-4 泄密報(bào)告也提到了 OpenAI 線上模型推理使用了它（不確定是否 100%）。

關(guān)于「Speculative decoding」，已有幾篇重要文獻(xiàn)可供參考，這也是 Karpathy 為了寫(xiě)這則推特所參考的論文，包括谷歌今年 1 月發(fā)表的論文《Fast Inference from Transformers via Speculative Decoding》、DeepMind 今年 2 月發(fā)表的論文《Accelerating Large Language Model Decoding with Speculative Sampling》，以及谷歌等機(jī)構(gòu) 2018 年的論文《Blockwise Parallel Decoding for Deep Autoregressive Models 》 。

簡(jiǎn)單來(lái)說(shuō)，「Speculative decoding」使用兩個(gè)模型：一個(gè)是原始目標(biāo)模型稱(chēng)為大模型，另一個(gè)是比原始模型小得多的近似模型稱(chēng)為小模型。主要思想是先讓小模型提前解碼多個(gè) token 進(jìn)行猜測(cè)，并將它們作為單個(gè) batch 輸入到一個(gè)大模型中進(jìn)行審核修正，其效果和直接用大模型解碼等價(jià)。如果小模型猜測(cè)的不準(zhǔn)確，那么大型模型會(huì)放棄小模型預(yù)測(cè)的 token，繼續(xù)使用大型模型進(jìn)行解碼。

由于小模型計(jì)算量小，從而大大減少了內(nèi)存訪問(wèn)需求。

介紹完「Speculative decoding」，我們?cè)倩氐?Karpathy 的推特。Karpathy 是針對(duì)下面內(nèi)容回復(fù)的。

Karpathy 表示：對(duì)于 LLM 來(lái)說(shuō)，「Speculative execution」 是一種極好的推理 — 時(shí)間優(yōu)化方法。

它取決于以下方面：在單個(gè)輸入 token 上分發(fā) LLM 所花費(fèi)的時(shí)間與在批處理中分發(fā) K 個(gè)輸入 token 所花費(fèi)的時(shí)間一樣多。產(chǎn)生這樣的原因是因?yàn)椴蓸訃?yán)重受內(nèi)存限制：模型運(yùn)行時(shí)的大部分工作不是在做計(jì)算，而是從 VRAM 讀取 transformer 的權(quán)重到片上緩存進(jìn)行處理。如果你要做的工作是來(lái)讀取這些權(quán)值，你可以把它們應(yīng)用到一整批輸入向量上。

但是我們不能一次性采樣一批 K 個(gè) token，因?yàn)槊?N 個(gè) token 都取決于我們?cè)诘?N-1 步采樣的 token。由于存在串行依賴(lài)性，因此基線實(shí)現(xiàn)只是從左到右逐一進(jìn)行。

現(xiàn)在最聰明的想法是使用一個(gè)小而便宜的草稿模型（draft model），先生成 K 個(gè) token 候選序列，即一個(gè)「草稿」。然后用大模型批量的將輸入組合在一起。速度幾乎與僅輸入一個(gè) token 一樣快。接著從左到右遍歷模型和樣本 token 預(yù)測(cè)的 logits。任何與「草稿」一致的樣本都允許立即跳到下一個(gè) token。如果存在分歧，那么就丟棄「草稿」并承擔(dān)一些一次性工作的成本（對(duì)「草稿」進(jìn)行采樣并為所有后續(xù) token 進(jìn)行前向傳遞）。

這種方法起作用的原因在于，很多「草稿」token 都會(huì)被接受，因?yàn)樗鼈兒苋菀祝约词故歉〉牟莞迥Ｐ鸵材艿玫剿鼈?。?dāng)這些簡(jiǎn)單的 token 被接受時(shí)，我們會(huì)跳過(guò)這些部分。大模型不同意的 hard token 會(huì)回落到原始速度，但由于一些額外的工作，實(shí)際上速度會(huì)慢一些。

Karpathy 表示，這個(gè)奇怪的技巧之所以有效，是因?yàn)?LLM 在推理時(shí)受到內(nèi)存限制，在對(duì)單個(gè)序列進(jìn)行采樣的 batch size=1 設(shè)置中，很大一部分本地 LLM 用例都屬于這種情況。因?yàn)榇蠖鄶?shù) token 都很「簡(jiǎn)單」。