輔助生成：低延遲文本生成的新方向
大型語言模型如今風靡一時，許多公司投入大量資源來擴展它們規模并解鎖新功能。然而，作為注意力持續時間不斷縮短的人類，我們并不喜歡大模型緩慢的響應時間。由于延遲對于良好的用戶體驗至關重要，人們通常使用較小的模型來完成任務，盡管它們的質量較低 (例如 代碼補全任務)。
為什么文本生成這么慢？是什么阻止你在不破產的情況下部署低延遲大型語言模型？在這篇博文中，我們將重新審視自回歸文本生成的瓶頸，并介紹一種新的解碼方法來解決延遲問題。你會發現，通過使用我們的新的輔助生成方法，你可以將硬件中的延遲降低多達 10 倍！
理解文本生成延遲
文本生成的核心很容易理解。讓我們看看核心部分 (即 ML 模型)，它的輸入包含一個文本序列，其中包括到目前為止生成的文本，以及其他特定于模型的組件 (例如 Whisper 還有一個音頻輸入)。該模型接受輸入并進行前向傳遞: 輸入被喂入模型并一層一層順序傳遞，直到預測出下一個 token 的非標準化對數概率 (也稱為 logits)。一個 token 可能包含整個詞、子詞，或者是單個字符，這取決于具體模型。如果你想深入了解文本生成的原理，GPT-2 插圖 是一個很好的參考。
模型的前向傳遞提供了下一個 token 的概率，你可以自由操作 (例如，將不需要的單詞或序列的概率設置為 0)。文本生成的步驟就是從這些概率中選擇下一個 token。常見的策略包括選擇最有可能的 token (貪心解碼)，或從它們的分布中抽樣 (多項式抽樣)。在選擇了下一個 token 之后，我們將模型前向傳遞與下一個 token 迭代地連接起來，繼續生成文本。這個解釋只是解碼方法的冰山一角; 請參閱我們 關于文本生成的博客 以進行深入探索。
從上面的描述中可以看出，文本生成的延遲瓶頸很明顯: 運行大型模型的前向傳遞很慢，你可能需要依次執行數百次迭代。但讓我們深入探討一下: 為什么前向傳遞速度慢？前向傳遞通常以矩陣乘法為主，通過查閱相應的 維基百科，你可以看出內存帶寬是此操作的限制 (例如，從 GPU RAM 到 GPU 計算核心)。換句話說， 前向傳遞的瓶頸來自將模型權重加載到設備的計算核心中，而不是來自執行計算本身。
目前，你可以探索三個主要途徑來充分理解文本生成，所有這些途徑都用于解決模型前向傳遞的性能問題。首先，對于特定硬件的模型優化。例如，如果你的設備可能與 Flash Attention 兼容，你可以使用它通可以過重新排序操作或 INT8 量化 來加速注意力層，其減少了模型權重的大小。
其次，如果你有并發文本生成需求，你可以對輸入進行批處理，從而實現較小的延遲損失并大幅增加吞吐量。你可以將模型對于多個輸入并行計算，這意味著你將在大致相同的內存帶寬負擔情況下獲得了更多 token。批處理的問題在于你需要額外的設備內存 (或在某處卸載內存)。你可以看到像 FlexGen 這樣的項目以延遲為代價來優化吞吐量。
# Example showcasing the impact of batched generation. Measurement device: RTX3090
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2").to("cuda")
inputs = tokenizer(["Hello world"], return_tensors="pt").to("cuda")
def print_tokens_per_second(batch_size):
new_tokens = 100
cumulative_time = 0
# warmup
model.generate(
**inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
)
for _ in range(10):
start = time.time()
model.generate(
**inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size
)
cumulative_time += time.time() - start
print(f"Tokens per second: {new_tokens * batch_size * 10 / cumulative_time:.1f}")
print_tokens_per_second(1) # Tokens per second: 418.3
print_tokens_per_second(64) # Tokens per second: 16266.2 (~39x more tokens per second)
最后，如果你有多個可用設備，你可以使用 Tensor 并行 分配工作負載并獲得更低的延遲。使用 Tensor 并行，你可以將內存帶寬負擔分攤到多個設備上，但除了在多個設備運行計算的成本之外，你還需要考慮設備間的通信瓶頸。該方法的收益在很大程度上取決于模型大小: 對于可以輕松在單個消費級設備上運行的模型，通常效果并不顯著。根據這篇 DeepSpeed 博客，你會發現你可以將大小為 17B 的模型分布在 4 個 GPU 上，從而將延遲減少 1.5 倍 (圖 7)。
這三種類型的改進可以串聯使用，從而產生 高通量解決方案。然而，在應用特定于硬件的優化后，降低延遲的方法有限——并且現有的方法很昂貴。讓我們接下來解決這個問題！
重新回顧語言模型解碼器的正向傳播
上文我們講到，每個模型前向傳遞都會產生下一個 token 的概率，但這實際上是一個不完整的描述。在文本生成期間，典型的迭代包括模型接收最新生成的 token 作為輸入，加上所有其他先前輸入的緩存內部計算，再返回下一個 token 得概率。緩存用于避免冗余計算，從而實現更快的前向傳遞，但它不是強制性的 (并且可以設置部分使用)。禁用緩存時，輸入包含到目前為止生成的整個 token 序列，輸出包含 _所有位置_的下一個 token 對應的概率分布！如果輸入由前 N 個 token 組成，則第 N 個位置的輸出對應于其下一個 token 的概率分布，并且該概率分布忽略了序列中的所有后續 token。在貪心解碼的特殊情況下，如果你將生成的序列作為輸入傳遞并將 argmax 運算符應用于生成的概率，你將獲得生成的序列。
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")
inputs = tok(["The"], return_tensors="pt")
generated = model.generate(**inputs, do_sample=False, max_new_tokens=10)
forward_confirmation = model(generated).logits.argmax(-1)
# We exclude the opposing tips from each sequence: the forward pass returns
# the logits for the next token, so it is shifted by one position.
print(generated[:-1].tolist() == forward_confirmation[1:].tolist()) # True
這意味著你可以將模型前向傳遞用于不同的目的: 除了提供一些 token 來預測下一個標記外，你還可以將序列傳遞給模型并檢查模型是否會生成相同的序列 (或部分相同序列)。
(請訪問閱讀原文查看動態演示)
讓我們想象，你可以訪問一個神奇的無延遲的預測輔助模型，該模型針對任何給定輸入生成與你的模型相同的序列。順便說一句，這個模型不能直接用，只能輔助你的生成程序。使用上述屬性，你可以使用此輔助模型獲取候選輸出 token，然后使用你的模型進行前向傳遞以確認它們的正確性。在這個烏托邦式的場景中，文本生成的延遲將從 O(n) 減少到 O(1)，其中生成的 token 數量為 n。對于需要多次迭代生成的過程，我們談論的是其數量級。
向現實邁出一步，我們假設輔助模型失去了它的預測屬性。根據你的模型，現在它是一個無延遲模型，但它會弄錯一些候選 token。由于任務的自回歸性質，一旦輔助模型得到一個錯誤的 token，所有后續候選 token 都必須無效。但是，你可以使用模型更正錯誤 token 并反復重復此過程后再次查詢輔助模型。即使輔助模型失敗了幾個 token，文本生成的延遲也會比原始形式小得多。
顯然，世界上沒有無延遲的輔助模型。然而，找到一個近似于模型的文本生成輸出的其它模型相對容易，例如經過類似訓練的相同架構的較小版本模型通常符合此需求。當模型大小的差異變得顯著時，使用較小的模型作為輔助模型的成本在跳過幾個前向傳遞后就顯得無關緊要了！現在，你了解了 _ 輔助生成 _ 的核心。
使用輔助模型的貪心解碼
輔助生成是一種平衡行為。你希望輔助模型快速生成候選序列，同時盡可能準確。如果輔助模型的質量很差，你將承擔使用輔助模型的成本，而收益卻很少甚至沒有。另一方面，優化候選序列的質量可能意味著使用更慢的輔助模型，從而導致網絡減速。雖然我們無法為你自動選擇輔助模型，但我們包含了一個額外的要求和一個啟發式方法，以確保模型與輔助模型一起花費的時間保持在可控范圍內。
首先，我們要求輔助模型必須具有與你的模型完全相同的分詞器。如果沒有此要求，則必須添加昂貴的 token 解碼和重新編碼步驟。此外，這些額外的步驟必須在 CPU 上進行，這反過來可能增加了設備間數據傳輸。能夠快速地使用輔助模型對于輔助生成的好處是至關重要的。
最后，啟發式。至此，你可能已經注意到電影盜夢空間和輔助生成之間的相似之處——畢竟你是在文本生成中運行文本生成。每個候選 token 有一個輔助模型前向傳播，我們知道前向傳播是昂貴的。雖然你無法提前知道輔助模型將獲得的 token 數量，但你可以跟蹤此信息并使用它來限制向輔助模型請求的候選 token 數量——輸出的某些部分比其它一些部分更容易被預計。
總結一下，這是我們最初實現的輔助生成的循環 (代碼):
- 使用貪心解碼與輔助模型生成一定數量的
候選 token。當第一次調用輔助生成時，生成的候選 token的數量被初始化為5。 - 使用我們的模型，對
候選 token進行前向計算，獲得每個 token 對應的概率。 - 使用 token 選擇方法 (使用
.argmax()進行貪心搜索或使用.multinomial()用于采樣方法) 來從概率中選取next_tokens。 - 比較步驟 3 中選擇的
next_tokens和候選 token中相同的 token 數量。請注意，我們需要從左到右進行比較， 在第一次不匹配后，后續所有候選 token都無效。5. 使用步驟 4 得到的匹配數量將候選 token分割。也就是，將輸入 tokens 加上剛剛驗證得到的正確的 tokens。 - 調整下一次迭代中生成的
候選 token的數量 —— 使用啟發式方法，如果步驟 3 中所有 token 都匹配，則候選 token的長度增加2，否則減少1。
(請訪問閱讀原文查看動態演示)
我們在 ?? Transformers 中設計了 API，因此使用該方法對你來說是無痛的。你需要做的就是將輔助模型作為 assistant_model 參數傳入從而獲得延遲收益！我們暫時限制了輔助生成的批量大小為 1。
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
print(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']
額外的內部復雜性是否值得？讓我們看一下貪心解碼情況下的延遲數 (采樣結果在下一節)?？紤]批量大小為 1，這些結果是直接從 ?? Transformers 中提取的，沒有任何額外的優化，因此你應該能夠在你的設置中復現它們。
Space 體驗地址: https://hf.co/spaces/joaogante/assisted_generation_benchmarks
通過觀察收集到的數據，我們發現輔助生成可以在不同的設置中顯著減少延遲，但這不是靈丹妙藥——你應該在應用之前對其進行系統的評估以清晰使用該方法的代價。對于輔助生成方法，我們可以得出結論:
- ?? 需要訪問至少比你的模型小一個數量級的輔助模型 (差異越大越好) ;
- ?? 在存在 INT8 的情況下獲得高達 3 倍的加速，否則能夠達到 2 倍的加速;
- ?? 如果你正在使用不適合你的模型的 GPU 并且依賴于內存卸載的模型，你可以看到高達 10 倍的加速;
- ?? 在輸入驅動任務中大放異彩，例如自動語音識別或摘要。
輔助生成的采樣方法
貪心解碼適用于以輸入為基礎的任務 (自動語音識別、翻譯、摘要……) 或事實知識尋求。對于需要大量創造力的開放式任務，例如使用語言模型作為聊天機器人的大多數任務，應該改用采樣方法。雖然輔助生成方法是為貪心解碼而設計的，但這并不意味著你不能使用多項式采樣進行輔助生成！
從 next token 的概率分布中抽取樣本將導致我們的基于貪心的輔助生產更頻繁地失敗，從而降低其延遲優勢。但是，我們可以使用采樣中的溫度系數來控制下一個標記的概率分布有多尖銳。在一種極端情況下，當溫度接近 0 時，采樣將近似于貪心解碼，有利于最有可能的 token。在另一個極端，當溫度設置為遠大于 1 的值時，采樣將是混亂的，從均勻分布中抽取。因此，低溫對你的輔助模型更有利，能夠保留輔助生成的大部分延遲優勢，如下所示。
不妨親眼看一看，感受一下輔助生成的魅力？
Space 體驗地址: https://hf.co/spaces/joaogante/assisted_generation_demo
未來發展方向
輔助生成表明當前文本生成策略已經到了可優化的階段。我們意識到它目前的難點不在于計算量的問題，因此可以應用簡單的啟發式方法來充分利用可用的內存帶寬，緩解瓶頸。我們相信，進一步優化輔助模型將使我們獲得更大的延遲降低——例如，如果我們請求輔助模型生成多個連續候選 token，我們可能能夠跳過更多的前向傳遞。自然地，使用高質量的小模型作為輔助模型對于實現和擴大收益至關重要。
該方法最初在我們的 ?? Transformers 庫下發布，用于 .generate() 函數，我們預期將其納入整個 Hugging Face 宇宙。它的實現也是完全開源的。因此，如果你正在進行文本生成而沒有使用我們的工具，你可以隨時將其作為參考。
最后，輔助生成重新提出了文本生成中的一個關鍵問題: 模型中所有新 token 都是給定模型以自回歸方式計算的結果，同質地前向傳遞每一個 token。這篇博文提出了這樣的想法: 生成的大部分序列也可以由小尺寸的模型同樣生成。為此，我們需要新的模型架構和解碼方法——我們很高興看到未來會帶來什么！
相關工作
在這篇博文最初發布后，我注意到其他作品也探索了相同的核心原則 (使用前向傳遞來驗證更長的連續性)。特別地，請看以下作品:
Citation
@misc {gante2023assisted,
author = { {Joao Gante} },
title = { Assisted Generation: a new direction toward low-latency text generation },
year = 2023,
url = { https://huggingface.co/blog/assisted-generation },
doi = { 10.57967/hf/0638 },
publisher = { Hugging Face Blog }
}
致謝
我要感謝 Sylvain Gugger、Nicolas Patry 和 Lewis Tunstall 分享了許多寶貴的建議來改進這篇博文。最后，感謝 Chunte Lee 設計了精美的封面，你可以在我們的網頁上看到。
原文鏈接: https://huggingface.co/blog/assisted-generation
作者: Joao Gante
譯者: gxy-gxy
排版/審校: zhongdongy (阿東)
