跳到主要內容
AI News HubLIVE
站內改寫3 分鐘閱讀

從輸入提示到看到第一個詞出現之間發生了什麼

文章摘要

本文詳細解釋了大語言模型在推理過程中的內部機制,包括自迴歸生成、預填充和解碼階段、KV快取的工作原理及其記憶體佔用,以及解碼策略如何影響輸出。

來源Hacker News AI作者: shbhmrzd
從輸入提示到看到第一個詞出現之間發生了什麼
回報錯誤

更正管道尚未開通,可先複製下方文章資訊留存。

查看更正說明
直接讀正文

本文是大語言模型工作原理系列的第四部分,重點介紹訓練好的模型在實際推理時發生了什麼。當使用者輸入一個提示並按下回車後,模型需要逐步生成回覆,而不是一次性輸出全部內容。

自迴歸生成

與訓練時模型能看到整個序列並並行預測每個位置的下一詞元不同,推理時模型只能一個一個地生成詞元。具體過程如下:使用者輸入提示,分詞器將其轉換為詞元ID;然後執行一次前向傳播,模型輸出整個提示之後下一個詞元的logits;從這些logits中取樣一個詞元,將其追加到序列中;再執行一次前向傳播,這次輸入是原始提示加上新生成的詞元;如此重複,直到達到所需長度或遇到停止詞元。這就是自迴歸生成——每一步都依賴前一步的輸出,因此必須序列進行。

例如,使用者輸入“What is gravity?”,分詞後為[What, is, gravity, ?](4個詞元)。第一步輸入這4個詞元,前向傳播後得到下一個詞元的機率分佈,“Gravity”機率最高,取樣得到“Gravity”。第二步輸入[What, is, gravity, ?, Gravity](5個詞元),得到“is”,依此類推,逐步生成“Gravity is the force…”直到停止。每一步模型都要重新處理整個已有序列,這導致大量冗餘計算。

預填充與解碼階段

實際推理分為兩個階段。

預填充階段:一次性並行處理整個提示。例如50個詞元的提示只需一次前向傳播,GPU同時計算所有詞元透過每一層的注意力模式(因果掩碼確保只關注前面的詞元)。預填充的真正目的是計算注意力,為每個詞元生成鍵(Key)和值(Value)向量,這些向量直接存入KV快取供解碼階段複用。短提示很快,但超長提示(數萬詞元)可能因注意力計算複雜度與序列長度成平方關係而變慢。

解碼階段:逐個生成新詞元,每次前向傳播只處理一個新詞元。GPU每次僅做少量矩陣乘法,瓶頸從計算轉向記憶體頻寬——大部分時間花在從記憶體載入模型權重而非做數學運算。這就是為什麼GPU記憶體頻寬(資料從視訊記憶體讀取的速度)對解碼效能比原始算力更重要。典型的聊天場景中,100詞元提示+500詞元回覆,預填充很快,解碼的500次順序前向傳播才是使用者感知到的延遲。這也是回覆流式輸出一個詞元一個詞元出現的原因——每個詞元生成後立即傳送,無需等待全部完成。

KV快取:跨步驟節省工作量

回憶解碼階段每次都要處理到目前為止的整個序列。在第100步,詞元1到99與第99步時完全一樣,注意力機制卻重複計算它們的鍵向量和值向量,造成浪費。KV快取透過儲存每個已計算詞元的K和V向量來避免重複:在第100步,只計算新詞元的K和V,複用快取中詞元1到99的向量。

預填充時,模型處理完整提示並將所有詞元的K和V向量存入快取。解碼時,生成詞元51只需執行一次前向傳播,計算其K和V,然後與快取的詞元1-50的向量拼接,再執行注意力。如此反覆,快取不斷增長但舊詞元的向量永不重算。

KV快取的記憶體佔用

以Llama 3.1-8B為例:32層、8個KV頭(分組查詢注意力)、每個頭維度128、FP16精度(每個值2位元組),支援最長128,000詞元的上下文。計算單個詞元在單個層單個頭的KV:一個Key向量128個數×2位元組=256位元組,Value同樣,所以一個詞元在單個頭單個層需512位元組。每層有8個頭,因此一個詞元在一層需512×8=4,096位元組。32層共4,096×32=131,072位元組≈128KB。對於4,096詞元的序列:128KB×4,096=512MB。若達到128,000詞元上下文,則約16GB。服務多個使用者時需乘以併發序列數,例如32使用者各4,096詞元,KV快取需16GB,與模型權重本身(約16GB)相當。這就是分組查詢注意力(GQA)的意義——將KV頭數從32減至8,快取大小降至1/4。雖然可能會損失一些細粒度注意力模式,但實際記憶體節省遠大於質量差異。

KV快取節省的工作量

無快取時生成1000個詞元需從頭計算K/V向量:第1步計算1個,第2步2個,……總計約500,000次K/V計算。有快取後每個詞元的K/V只計算一次,總共1000次。但注意力點積的數量仍然是1+2+…+1000≈500,000,這一O(n²)成本不變。KV快取主要節省了計算K/V向量的矩陣乘法,而無快取時模型需要為每個歷史詞元在每一層重複執行昂貴的矩陣運算。

解碼策略

模型輸出每個詞元的機率後,還需要透過解碼策略來決定實際選擇哪個詞元。常見的策略包括貪婪解碼(總是選機率最高的詞元)、溫度取樣(調整機率分佈的尖銳程度)、top-k取樣(只從機率最高的k個詞元中取樣)和top-p(nucleus)取樣(從累積機率超過p的最小集合中取樣)。這些策略平衡了回覆的創造性與連貫性。在實際應用中,開發者會根據任務需求調整這些引數,例如聊天場景可能使用較高的溫度以增加多樣性,而事實性問答則傾向於低溫度或貪婪解碼。

總之,大語言模型的推理是一個複雜但精妙的過程,從輸入提示到最終回覆的每一步都涉及大量工程最佳化。理解預填充、解碼、KV快取和解碼策略,有助於開發者更好地利用和最佳化這些模型。

展開要點與分析

文章情報

工程師進階

要點

  • 大語言模型的推理是自迴歸的:每次只生成一個詞,每一步都依賴前面的輸出。
  • 推理分為預填充(並行處理整個提示)和解碼(逐個生成詞)兩個階段。
  • KV快取儲存了歷史詞元的鍵值向量,大幅減少重複計算,但佔用大量視訊記憶體。
  • 解碼策略(如溫度、top-k取樣)決定了回覆的創造性和連貫性。

要點與分析由自動化流程生成,可能有誤,請結合原始來源核實。