大型語言模型如今風(fēng)靡一時，許多公司投入大量資源來擴(kuò)展它們規(guī)模并解鎖新功能。然而，作為注意力持續(xù)時間不斷縮短的人類，我們并不喜歡大模型緩慢的響應(yīng)時間。由于延遲對于良好的用戶體驗(yàn)至關(guān)重要，人們通常使用較小的模型來完成任務(wù)，盡管它們的質(zhì)量較低 (例如?代碼補(bǔ)全任務(wù))。

為什么文本生成這么慢？是什么阻止你在不破產(chǎn)的情況下部署低延遲大型語言模型？在這篇博文中，我們將重新審視自回歸文本生成的瓶頸，并介紹一種新的解碼方法來解決延遲問題。你會發(fā)現(xiàn)，通過使用我們的新的輔助生成方法，你可以將硬件中的延遲降低多達(dá) 10 倍！

理解文本生成延遲

文本生成的核心很容易理解。讓我們看看核心部分 (即 ML 模型)，它的輸入包含一個文本序列，其中包括到目前為止生成的文本，以及其他特定于模型的組件 (例如 Whisper 還有一個音頻輸入)。該模型接受輸入并進(jìn)行前向傳遞: 輸入被喂入模型并一層一層順序傳遞，直到預(yù)測出下一個 token 的非標(biāo)準(zhǔn)化對數(shù)概率 (也稱為 logits)。一個 token 可能包含整個詞、子詞，或者是單個字符，這取決于具體模型。如果你想深入了解文本生成的原理，GPT-2 插圖 是一個很好的參考。

(請訪問閱讀原文查看動態(tài)演示)

模型的前向傳遞提供了下一個 token 的概率，你可以自由操作 (例如，將不需要的單詞或序列的概率設(shè)置為 0)。文本生成的步驟就是從這些概率中選擇下一個 token。常見的策略包括選擇最有可能的 token (貪心解碼)，或從它們的分布中抽樣 (多項(xiàng)式抽樣)。在選擇了下一個 token 之后，我們將模型前向傳遞與下一個 token 迭代地連接起來，繼續(xù)生成文本。這個解釋只是解碼方法的冰山一角; 請參閱我們 關(guān)于文本生成的博客 以進(jìn)行深入探索。

(請訪問閱讀原文查看動態(tài)演示)

從上面的描述中可以看出，文本生成的延遲瓶頸很明顯: 運(yùn)行大型模型的前向傳遞很慢，你可能需要依次執(zhí)行數(shù)百次迭代。但讓我們深入探討一下: 為什么前向傳遞速度慢？前向傳遞通常以矩陣乘法為主，通過查閱相應(yīng)的 維基百科，你可以看出內(nèi)存帶寬是此操作的限制 (例如，從 GPU RAM 到 GPU 計算核心)。換句話說， _前向傳遞的瓶頸來自將模型權(quán)重加載到設(shè)備的計算核心中，而不是來自執(zhí)行計算本身_。

目前，你可以探索三個主要途徑來充分理解文本生成，所有這些途徑都用于解決模型前向傳遞的性能問題。首先，對于特定硬件的模型優(yōu)化。例如，如果你的設(shè)備可能與 Flash Attention 兼容，你可以使用它通可以過重新排序操作或 INT8 量化 來加速注意力層，其減少了模型權(quán)重的大小。

其次，如果你有并發(fā)文本生成需求，你可以對輸入進(jìn)行批處理，從而實(shí)現(xiàn)較小的延遲損失并大幅增加吞吐量。你可以將模型對于多個輸入并行計算，這意味著你將在大致相同的內(nèi)存帶寬負(fù)擔(dān)情況下獲得了更多 token。批處理的問題在于你需要額外的設(shè)備內(nèi)存 (或在某處卸載內(nèi)存)。你可以看到像 FlexGen 這樣的項(xiàng)目以延遲為代價來優(yōu)化吞吐量。

#?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)


最后，如果你有多個可用設(shè)備，你可以使用 Tensor 并行 分配工作負(fù)載并獲得更低的延遲。使用 Tensor 并行，你可以將內(nèi)存帶寬負(fù)擔(dān)分?jǐn)偟蕉鄠€設(shè)備上，但除了在多個設(shè)備運(yùn)行計算的成本之外，你還需要考慮設(shè)備間的通信瓶頸。該方法的收益在很大程度上取決于模型大小: 對于可以輕松在單個消費(fèi)級設(shè)備上運(yùn)行的模型，通常效果并不顯著。根據(jù)這篇 DeepSpeed 博客，你會發(fā)現(xiàn)你可以將大小為 17B 的模型分布在 4 個 GPU 上，從而將延遲減少 1.5 倍 (圖 7)。

這三種類型的改進(jìn)可以串聯(lián)使用，從而產(chǎn)生 高通量解決方案。然而，在應(yīng)用特定于硬件的優(yōu)化后，降低延遲的方法有限——并且現(xiàn)有的方法很昂貴。讓我們接下來解決這個問題！

重新回顧語言模型解碼器的正向傳播

上文我們講到，每個模型前向傳遞都會產(chǎn)生下一個 token 的概率，但這實(shí)際上是一個不完整的描述。在文本生成期間，典型的迭代包括模型接收最新生成的 token 作為輸入，加上所有其他先前輸入的緩存內(nèi)部計算，再返回下一個 token 得概率。緩存用于避免冗余計算，從而實(shí)現(xiàn)更快的前向傳遞，但它不是強(qiáng)制性的 (并且可以設(shè)置部分使用)。禁用緩存時，輸入包含到目前為止生成的整個 token 序列，輸出包含 _所有位置_的下一個 token 對應(yīng)的概率分布！如果輸入由前 N 個 token 組成，則第 N 個位置的輸出對應(yīng)于其下一個 token 的概率分布，并且該概率分布忽略了序列中的所有后續(xù) token。在貪心解碼的特殊情況下，如果你將生成的序列作為輸入傳遞并將 argmax 運(yùn)算符應(yīng)用于生成的概率，你將獲得生成的序列。

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 來預(yù)測下一個標(biāo)記外，你還可以將序列傳遞給模型并檢查模型是否會生成相同的序列 (或部分相同序列)。

(請訪問閱讀原文查看動態(tài)演示)

讓我們想象，你可以訪問一個神奇的無延遲的預(yù)測輔助模型，該模型針對任何給定輸入生成與你的模型相同的序列。順便說一句，這個模型不能直接用，只能輔助你的生成程序。使用上述屬性，你可以使用此輔助模型獲取候選輸出 token，然后使用你的模型進(jìn)行前向傳遞以確認(rèn)它們的正確性。在這個烏托邦式的場景中，文本生成的延遲將從?O(n)?減少到?O(1)，其中生成的 token 數(shù)量為?n。對于需要多次迭代生成的過程，我們談?wù)摰氖瞧鋽?shù)量級。

向現(xiàn)實(shí)邁出一步，我們假設(shè)輔助模型失去了它的預(yù)測屬性。根據(jù)你的模型，現(xiàn)在它是一個無延遲模型，但它會弄錯一些候選 token。由于任務(wù)的自回歸性質(zhì)，一旦輔助模型得到一個錯誤的 token，所有后續(xù)候選 token 都必須無效。但是，你可以使用模型更正錯誤 token 并反復(fù)重復(fù)此過程后再次查詢輔助模型。即使輔助模型失敗了幾個 token，文本生成的延遲也會比原始形式小得多。

顯然，世界上沒有無延遲的輔助模型。然而，找到一個近似于模型的文本生成輸出的其它模型相對容易，例如經(jīng)過類似訓(xùn)練的相同架構(gòu)的較小版本模型通常符合此需求。當(dāng)模型大小的差異變得顯著時，使用較小的模型作為輔助模型的成本在跳過幾個前向傳遞后就顯得無關(guān)緊要了！現(xiàn)在，你了解了 _ 輔助生成 _ 的核心。

使用輔助模型的貪心解碼

輔助生成是一種平衡行為。你希望輔助模型快速生成候選序列，同時盡可能準(zhǔn)確。如果輔助模型的質(zhì)量很差，你將承擔(dān)使用輔助模型的成本，而收益卻很少甚至沒有。另一方面，優(yōu)化候選序列的質(zhì)量可能意味著使用更慢的輔助模型，從而導(dǎo)致網(wǎng)絡(luò)減速。雖然我們無法為你自動選擇輔助模型，但我們包含了一個額外的要求和一個啟發(fā)式方法，以確保模型與輔助模型一起花費(fèi)的時間保持在可控范圍內(nèi)。

首先，我們要求輔助模型必須具有與你的模型完全相同的分詞器。如果沒有此要求，則必須添加昂貴的 token 解碼和重新編碼步驟。此外，這些額外的步驟必須在 CPU 上進(jìn)行，這反過來可能增加了設(shè)備間數(shù)據(jù)傳輸。能夠快速地使用輔助模型對于輔助生成的好處是至關(guān)重要的。

最后，啟發(fā)式。至此，你可能已經(jīng)注意到電影盜夢空間和輔助生成之間的相似之處——畢竟你是在文本生成中運(yùn)行文本生成。每個候選 token 有一個輔助模型前向傳播，我們知道前向傳播是昂貴的。雖然你無法提前知道輔助模型將獲得的 token 數(shù)量，但你可以跟蹤此信息并使用它來限制向輔助模型請求的候選 token 數(shù)量——輸出的某些部分比其它一些部分更容易被預(yù)計。

總結(jié)一下，這是我們最初實(shí)現(xiàn)的輔助生成的循環(huán) (代碼):

1. 使用貪心解碼與輔助模型生成一定數(shù)量的候選 token。當(dāng)?shù)谝淮握{(diào)用輔助生成時，生成的候選 token?的數(shù)量被初始化為?5。

2. 使用我們的模型，對候選 token?進(jìn)行前向計算，獲得每個 token 對應(yīng)的概率。

3. 使用 token 選擇方法 (使用.argmax()?進(jìn)行貪心搜索或使用?.multinomial()?用于采樣方法) 來從概率中選取?next_tokens。

4. 比較步驟 3 中選擇的?next_tokens?和 ?候選 token?中相同的 token 數(shù)量。請注意，我們需要從左到右進(jìn)行比較， 在第一次不匹配后，后續(xù)所有?候選 token都無效。5. 使用步驟 4 得到的匹配數(shù)量將候選 token?分割。也就是，將輸入 tokens 加上剛剛驗(yàn)證得到的正確的 tokens。

5. 調(diào)整下一次迭代中生成的候選 token?的數(shù)量 —— 使用啟發(fā)式方法，如果步驟 3 中所有 token 都匹配，則候選 token?的長度增加?2，否則減少?1。

(請訪問閱讀原文查看動態(tài)演示)

我們在 ?? Transformers 中設(shè)計了 API，因此使用該方法對你來說是無痛的。你需要做的就是將輔助模型作為?assistant_model?參數(shù)傳入從而獲得延遲收益！我們暫時限制了輔助生成的批量大小為?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']


額外的內(nèi)部復(fù)雜性是否值得？讓我們看一下貪心解碼情況下的延遲數(shù) (采樣結(jié)果在下一節(jié))。考慮批量大小為 1，這些結(jié)果是直接從 ?? Transformers 中提取的，沒有任何額外的優(yōu)化，因此你應(yīng)該能夠在你的設(shè)置中復(fù)現(xiàn)它們。

Space 體驗(yàn)地址:
https://hf.co/spaces/joaogante/assisted_generation_benchmarks

通過觀察收集到的數(shù)據(jù)，我們發(fā)現(xiàn)輔助生成可以在不同的設(shè)置中顯著減少延遲，但這不是靈丹妙藥——你應(yīng)該在應(yīng)用之前對其進(jìn)行系統(tǒng)的評估以清晰使用該方法的代價。對于輔助生成方法，我們可以得出結(jié)論:

1. ?? 需要訪問至少比你的模型小一個數(shù)量級的輔助模型 (差異越大越好) ;

2. ?? 在存在 INT8 的情況下獲得高達(dá) 3 倍的加速，否則能夠達(dá)到 2 倍的加速;

3. ?? 如果你正在使用不適合你的模型的 GPU 并且依賴于內(nèi)存卸載的模型，你可以看到高達(dá) 10 倍的加速;

4. ?? 在輸入驅(qū)動任務(wù)中大放異彩，例如自動語音識別或摘要。

輔助生成的采樣方法

貪心解碼適用于以輸入為基礎(chǔ)的任務(wù) (自動語音識別、翻譯、摘要……) 或事實(shí)知識尋求。對于需要大量創(chuàng)造力的開放式任務(wù)，例如使用語言模型作為聊天機(jī)器人的大多數(shù)任務(wù)，應(yīng)該改用采樣方法。雖然輔助生成方法是為貪心解碼而設(shè)計的，但這并不意味著你不能使用多項(xiàng)式采樣進(jìn)行輔助生成！

從?next token?的概率分布中抽取樣本將導(dǎo)致我們的基于貪心的輔助生產(chǎn)更頻繁地失敗，從而降低其延遲優(yōu)勢。但是，我們可以使用采樣中的溫度系數(shù)來控制下一個標(biāo)記的概率分布有多尖銳。在一種極端情況下，當(dāng)溫度接近 0 時，采樣將近似于貪心解碼，有利于最有可能的 token。在另一個極端，當(dāng)溫度設(shè)置為遠(yuǎn)大于 1 的值時，采樣將是混亂的，從均勻分布中抽取。因此，低溫對你的輔助模型更有利，能夠保留輔助生成的大部分延遲優(yōu)勢，如下所示。

不妨親眼看一看，感受一下輔助生成的魅力？

Space 體驗(yàn)地址:https://hf.co/spaces/joaogante/assisted_generation_demo

未來發(fā)展方向

輔助生成表明當(dāng)前文本生成策略已經(jīng)到了可優(yōu)化的階段。我們意識到它目前的難點(diǎn)不在于計算量的問題，因此可以應(yīng)用簡單的啟發(fā)式方法來充分利用可用的內(nèi)存帶寬，緩解瓶頸。我們相信，進(jìn)一步優(yōu)化輔助模型將使我們獲得更大的延遲降低——例如，如果我們請求輔助模型生成多個連續(xù)候選 token，我們可能能夠跳過更多的前向傳遞。自然地，使用高質(zhì)量的小模型作為輔助模型對于實(shí)現(xiàn)和擴(kuò)大收益至關(guān)重要。

該方法最初在我們的 ?? Transformers 庫下發(fā)布，用于?.generate()?函數(shù)，我們預(yù)期將其納入整個 Hugging Face 宇宙。它的實(shí)現(xiàn)也是完全開源的。因此，如果你正在進(jìn)行文本生成而沒有使用我們的工具，你可以隨時將其作為參考。

最后，輔助生成重新提出了文本生成中的一個關(guān)鍵問題: 模型中所有新 token 都是給定模型以自回歸方式計算的結(jié)果，同質(zhì)地前向傳遞每一個 token。這篇博文提出了這樣的想法: 生成的大部分序列也可以由小尺寸的模型同樣生成。為此，我們需要新的模型架構(gòu)和解碼方法——我們很高興看到未來會帶來什么！

相關(guān)工作

在這篇博文最初發(fā)布后，我注意到其他作品也探索了相同的核心原則 (使用前向傳遞來驗(yàn)證更長的連續(xù)性)。特別地，請看以下作品:

• 分塊并行解碼, 來自 Google Brain

• 推測性采樣, 來自 DeepMind

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 分享了許多寶貴的建議來改進(jìn)這篇博文。最后，感謝 Chunte Lee 設(shè)計了精美的封面，你可以在我們的網(wǎng)頁上看到。

原文鏈接:?https://hf.co/blog/assisted-generation

作者: Joao Gante

譯者: gxy-gxy

排版/審校: zhongdongy (阿東)