太長不看版: 許多 LLM (如 gemma-2-9b 、 Mixtral-8x22B-Instruct-v0.1 等) 苦於缺乏對應小模型，而無法適用 輔助生成 方案。本文，我們將介紹由英特爾研究院和 Hugging Face 合作開發的 通用輔助生成 技術。有了這項技術，LLM 可與 任意 SLM 搭配組成輔助生成方案。從而，我們可以用輔助生成技術加速 任意 解碼器模型或 混合專家 模型以獲得 1.5x-2.0x 的加速比。重要的是，開銷幾乎為零 🔥🔥🔥！一起了解一下吧！
引言
如今，風頭最勁的開放權重 LLM 引數量一般都有數十億到數千億 (說你呢 Llama-3.1-405B 👋)，這給在生產環境中部署這些餓獸帶來了一系列工程挑戰。挑戰之一就是: 大模型文字生成速度很慢。為此，社群開發了很多不同的技術來加速解碼過程。輔助生成，也稱為 投機解碼，是其中一種非常常用且實用的方法，可在不損失準確性的情況下加速 LLM 推理。本文，我們將瞭解輔助生成的工作原理，並分享我們的最新研究成果，該成果使得對 Hugging Face Hub 14 萬個語言模型 中的 任意一個 模型進行加速成為可能，🚀！
輔助生成
輔助生成的核心是一對模型，分別稱為 目標模型 和 輔助模型 ，其中輔助模型是目標模型的小版，舉個例子，你可以使用 Llama-3.2-1B 作為較大的 Llama-3.1-70b 目標模型的輔助模型。整個生成過程是一個迭代過程: 每一輪，輔助模型會先一個一個自迴歸地生成多個詞元; 接著，目標模型透過一次前向傳播驗證輔助模型本輪生成的所有詞元。加速的奧秘就在於目標模型每次前向傳播中可以驗證多個詞元，而不像原本每次只能生成一個詞元。更詳細的解釋，請參閱 原博文。結合新近推出的 動態投機 策略，輔助生成可將文字生成速度提高 1.5 至 3 倍，具體倍數取決於任務型別及所使用的模型。
但，輔助生成並非無往而不利，一個最明顯的問題就是: 其要求目標模型和輔助模型必須使用相同的分詞器，這意味著兩者必須來自同一個模型系列。然而，許多廣泛使用的模型缺乏合適的“矮小緊”模型，因此與如此大幅的延遲降低無緣。根據我們的經驗，一般來說，輔助模型需要至少比目標模型小 50-100 倍，才會看到有意義的加速。舉幾個例子，CodeLlama-13b 沒有小模型; gemma-2-9b 只有一個 2b 的小模型，顯然不夠小、不夠快，因此加速註定不會太明顯。
通用輔助生成
為了緩解這個痛點，英特爾研究院與 Hugging Face 合作開發了通用輔助生成 (Universal Assisted Generation，UAG) 技術。UAG 可以無視分詞器的差異，配對任意目標模型和輔助模型。例如，可以使用 gemma-2-9b 作為目標模型，並選取 vicuna-68m 作為輔助模型。
該技術背後的主要思想是雙路分詞器對映: 每一輪，輔助模型生成完詞元后，就將其輸出詞元序列解碼為文字，再使用目標模型的分詞器將文字編碼成詞元序列; 同樣地，在目標模型驗證完後，將目標模型的詞元序列用相同的方法轉換回輔助模型的詞元序列，再將其新增至輔助模型的上下文用於下一輪迭代。
由於輔助模型和目標模型的分詞器的詞彙表不同，因此還需要處理由此帶來的差異。為了準確地對輔助模型新生成的詞元序列進行重編碼，必須再多給它一些上文詞元。然後，將整個序列重新編碼為目標模型的詞元格式，並與之前生成的最新的目標詞元對齊，以錨定新生成詞元的確切位置。下面的影片對此過程進行了圖解。
從目標模型到輔助模型的詞元重編碼也遵循與上述影片類似的過程。此時，如遇不匹配的詞元，需從輔助模型的鍵值 (KV) 快取中將它們丟棄掉，以保證資料的完整性。
基準測試
下表展示了不同目標模型與異分詞器輔助模型形成輔助解碼方案時測得的延遲改進。
| 目標模型 | 輔助模型 | 資料集 | 任務 | 加速比 |
|---|---|---|---|---|
codellama/CodeLlama-13b-Instruct-hf |
bigcode/tiny_starcoder_py |
openai/humaneval |
程式碼生成 | 1.90x |
mistralai/Mixtral-8x22B-Instruct-v0.1 |
double7/vicuna-68m |
cnn_dailymail |
摘要 | 1.52x |
google/gemma-2-9b |
double7/vicuna-68m |
cnn_dailymail |
摘要 | 1.76x |
mistralai/Mixtral-8x22B-Instruct-v0.1 |
Qwen/Qwen2-0.5B-Instruct |
tau/scrolls |
長文摘要 | 1.78x |
meta-llama/Llama-3.1-70B |
Qwen/Qwen2-0.5B-Instruct |
tau/scrolls |
長文摘要 | 1.78x |
microsoft/Phi-3-medium-128k-instruct |
Qwen/Qwen2-0.5B-Instruct |
tau/scrolls |
長文摘要 | 1.91x |
請注意，在標準輔助解碼方案下，上表中所有目標模型都會苦於沒有合適的小模型 (低於 10 億引數)。
上述實驗均在 100 個隨機樣本上完成。 Llama 和 Mixtral 目標模型的實驗分別用了 2 張和 4 張 A100 GPU; 其他所有實驗均使用單張 A6000 GPU。
程式碼
通用輔助生成技術已整合至 🤗 Transformers 4.46.0 版。
要使能該技術，需將 tokenizer 和 assistant_tokenizer 傳遞給 generate() ，示例程式碼如下:
from transformers import AutoModelForCausalLM, AutoTokenizer
prompt = "Alice and Bob"
checkpoint = "google/gemma-2-9b"
assistant_checkpoint = "double7/vicuna-68m"
assistant_tokenizer = AutoTokenizer.from_pretrained(assistant_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt")
model = AutoModelForCausalLM.from_pretrained(checkpoint)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint)
outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
tokenizer.batch_decode(outputs, skip_special_tokens=True)
輸出如下:
['Alice and Bob are sitting in a bar. Alice is drinking a beer and Bob is drinking a']
下一步
標準輔助生成方案在 do_sample=True 時，使用的投機取樣演算法為 該論文的演算法 1，但 UAG
目前僅實現了多項分佈取樣。在多項分佈取樣中，如果目標模型與輔助模型取樣得的詞元不相同時，會自動拒絕該詞元，這與投機取樣對此情況的處理不同。在實踐中，這意味著與共享分詞器的標準方案相比，UAG 方案在 do_sample=True 時吞吐量會較低。將來，我們計劃增加對 UAG 投機取樣的支援。
此外，我們還打算將 UAG 整合到 🤗 Transformers 流水線中，以使使用者能夠更簡單、輕鬆地利用它。
參考資源
- Fast Inference from Transformers via Speculative Decoding
- 輔助生成: 低延遲文字生成的新方向
英文原文: https://hf.co/blog/universal_assisted_generation
原文作者: Daniel Korat，Oren Pereg，Moshe Berchansky，Jonathan Mamou，Joao Gante，Lewis Tunstall，Nadav Timor，Moshe Wasserblat
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。