這是一份關於我新藝術專案 microgpt 的簡要指南，它是一個包含 200 行純 Python 程式碼、無任何依賴的單一檔案，能夠訓練和推理一個 GPT。這個檔案包含了所有必需的完整演算法內容：文件資料集、分詞器、自動微分引擎、類似 GPT-2 的神經網路架構、Adam 優化器、訓練迴圈和推理迴圈。其他一切都只是效率問題。我無法再進一步簡化了。這個腳本是我多個專案（micrograd、makemore、nanogpt 等）的結晶，以及我十多年來將 LLM 簡化至最基本要素的執著追求，我認為它很美妙🥹。它甚至完美地分為三欄：
以下是我引導有興趣的讀者逐步了解程式碼的指南。
大型語言模型的燃料是文字資料流，可選擇性地分成一組文件。在生產級應用中，每個文件都會是網際網路網頁，但對於 microgpt，我們使用一個更簡單的範例：32,000 個名字，每行一個：
資料集看起來像這樣。每個名字都是一個文件：
模型的目標是學習資料中的模式，然後生成分享統計模式的新文件。作為預覽，在腳本結束時，我們的模型將生成（「幻覺」！）聽起來合理的新名字。跳到後面，我們將得到：
這看起來沒什麼，但從像 ChatGPT 這樣的模型的角度來看，你與它的對話只是一個看起來很有趣的「文件」。當你用提示初始化文件時，模型從其角度的回應只是一個統計文件補全。
在底層，神經網路處理的是數字，而不是字元，所以我們需要一種方法將文字轉換為一系列整數 token ID，然後再轉換回來。生產級分詞器如 tiktoken（GPT-4 使用）為了效率會處理字元塊，但最簡單的分詞器只是為資料集中每個唯一的字元分配一個整數：
在上面的程式碼中，我們收集資料集中所有唯一的字元（也就是所有小寫字母 a-z），對它們進行排序，然後每個字母根據其索引獲得一個 ID。請注意，整數值本身沒有任何意義；每個 token 只是單獨的離散符號。它們可以是不同的表情符號，而不是 0、1、2。此外，我們還創建了一個額外的特殊 token，稱為 BOS（Beginning of Sequence，序列開頭），它充當分隔符：它告訴模型「一個新文件在此開始/結束」。稍後在訓練期間，每個文件都會用 BOS 包裹起來：[BOS, e, m, m, a, BOS]。模型會學習到 BOS 啟動一個新名字，而另一個 BOS 則結束它。因此，我們最終的詞彙量是 27（26 個小寫字母 a-z，加上 1 個 BOS token）。
訓練神經網路需要梯度：對於模型中的每個參數，我們需要知道「如果我稍微增加這個數字，損失會上升還是下降，以及上升或下降多少？」。計算圖有許多輸入（模型參數和輸入 token），但最終匯聚到一個單一的純量輸出：損失（我們將在下面精確定義損失是什麼）。反向傳播從那個單一輸出開始，並向後穿過圖，計算損失相對於每個輸入的梯度。它依賴於微積分的鏈式法則。在生產環境中，PyTorch 等函式庫會自動處理這個問題。在這裡，我們在一個名為 Value 的類別中從頭開始實現它：
我意識到這是數學和演算法上最密集的部分，我為此準備了一個 2.5 小時的影片：micrograd 影片。簡而言之，Value 包裝了一個單一的純量數字（.data）並追蹤它是如何計算的。將每個操作想像成一個小樂高積木：它接收一些輸入，產生一個輸出（前向傳播），並且知道它的輸出如何相對於其每個輸入而變化（局部梯度）。這就是 autograd 從每個積木所需的所有資訊。其他一切都只是鏈式法則，將積木串聯起來。
每次你用 Value 物件進行數學運算（加、乘等）時，結果都是一個新的 Value，它會記住它的輸入（_children）以及該操作的局部導數（_local_grads）。例如，__mul__ 記錄了 \( rac{\partial(a \cdot b)}{\partial a} = b\) 和 \( rac{\partial(a \cdot b)}{\partial b} = a\)。完整的積木集合：
backward() 方法以逆拓撲順序（從損失開始，到參數結束）遍歷這個圖，在每一步應用鏈式法則。如果損失是 \(L\)，而節點 \(v\) 有一個子節點 \(c\)，其局部梯度為 \( rac{\partial v}{\partial c}\)，則：
如果你不熟悉微積分，這看起來有點嚇人，但這實際上只是以直觀的方式將兩個數字相乘。一種看待它的方式如下：「如果一輛汽車的速度是自行車的兩倍，而自行車的速度是步行男子的四倍，那麼汽車的速度就是男子的 2 x 4 = 8 倍。」鏈式法則也是同樣的道理：你將沿路徑的變化率相乘。
我們通過在損失節點上設置 self.grad = 1 來啟動，因為 \( rac{\partial L}{\partial L} = 1\): 損失相對於自身的變化率顯然是 1。從那裡開始，鏈式法則只是將局部梯度沿著每條路徑乘回到參數上。
請注意 +=（累加，而不是賦值）。當一個值在圖中被使用多次時（即圖分支），梯度會沿著每個分支獨立流回，並且必須加總。這是多元鏈式法則的結果：如果 \(c\) 通過多個路徑影響 \(L\)，則總導數是來自每條路徑的貢獻總和。
在 backward() 完成後，圖中的每個 Value 都有一個 .grad，其中包含 \( rac{\partial L}{\partial v}\)，它告訴我們如果我們稍微改變該值，最終損失會如何變化。
這是一個具體的例子。請注意 a 被使用了兩次（圖分支），所以它的梯度是兩條路徑的總和：
這正是 PyTorch 的 .backward() 所提供的：
這與 PyTorch 的 loss.backward() 執行的演算法相同，只是處理的是純量而不是張量（純量陣列）——演算法上相同，尺寸小得多且更簡單，但當然效率也低得多。
讓我們詳細說明上面的 .backward() 提供的結果。Autograd 計算出如果 L = a*b + a，且 a=2 且 b=3，那麼 a.grad = 4.0，這告訴我們 a 對 L 的局部影響。如果你稍微改變輸入 a，L 會朝哪個方向變化？在這裡，L 對 a 的導數是 4.0，這意味著如果我們稍微增加 a（例如 0.001），L 大約會增加 4 倍（0.004）。同樣，b.grad = 2.0 表示對 b 的相同微調會使 L 大約增加 2 倍（0.002）。換句話說，這些梯度告訴我們每個單獨輸入對最終輸出（損失）的影響的方向（取決於符號的正負）和陡度（幅度）。這使我們能夠迭代地微調神經網路的參數以降低損失，從而提高其預測。
參數是模型的知識。它們是大量浮點數的集合（為了 autograd 而用 Value 包裝），它們最初是隨機的，並在訓練期間進行迭代優化。每個參數的確切作用一旦我們定義下面的模型架構就會更清楚，但現在我們只需要初始化它們：
每個參數都初始化為從高斯分佈中抽取的小隨機數。state_dict 將它們組織成命名矩陣（借用 PyTorch 的術語）：嵌入表、注意力權重、MLP 權重和最終輸出投影。我們還將所有參數展平成一個單一列表 params，以便優化器稍後可以遍歷它們。在我們這個微小模型中，這總共有 4,192 個參數。GPT-2 有 16 億個參數，而現代 LLM 有數千億個參數。
模型架構是一個無狀態函數：它接收一個 token、一個位置、參數以及先前位置的快取鍵/值，並返回模型認為序列中下一個 token 的 logits（分數）。我們遵循 GPT-2 的架構，並進行了少量簡化：使用 RMSNorm 而非 LayerNorm，沒有偏差，以及 ReLU 而非 GeLU。首先，三個小型輔助函數：
linear 是一個矩陣向量乘法。它接收一個向量 x 和一個權重矩陣 w，並計算 w 的每行的一個點積。這是神經網路的基本構件：一個學習到的線性變換。
softmax 將一組原始分數（logits）向量（範圍可以從 \(-\infty\) 到 \(+\infty\)）轉換為一個概率分佈：所有值最終都在 \([0, 1]\) 範圍內，並且總和為 1。我們首先減去最大值以獲得數值穩定性（這在數學上不改變結果，但可以防止 exp 中的溢出）。
rmsnorm（均方根歸一化）對向量進行重新縮放，使其值的均方根為單位。這可以防止激活在通過網路時增長或縮小，從而穩定訓練。它是原始 GPT-2 中使用的 LayerNorm 的一個更簡單的變體。
該函數處理時間序列中的特定位置（pos_id）的一個 token（ID 為 token_id），以及由鍵和值中的激活總結的先前迭代的某些上下文，這些上下文稱為 KV 快取。以下是逐步進行的過程：
嵌入。神經網路無法直接處理像 5 這樣的原始 token ID。它只能處理向量（數字列表）。因此，我們為每個可能的 token 關聯一個學習到的向量，並將其作為其神經簽名輸入。token ID 和位置 ID 分別從各自的嵌入表中（wte 和 wpe）查找一行。這兩個向量相加，為模型提供一個同時編碼 token 是什麼以及它在序列中位置的表示。現代 LLM 通常會省略位置嵌入，並引入其他基於相對位置的方案，例如 RoPE。
注意力區塊。當前 token 被投影到三個向量：查詢（Q）、鍵（K）和值（V）。直觀地說，查詢表示「我在尋找什麼？」，鍵表示「我包含什麼？」，值表示「如果被選中，我提供什麼？」。例如，在名字「emma」中，當模型處理第二個「m」並試圖預測下一個是什麼時，它可能會學到一個類似「最近出現了哪些母音？」的查詢。早期的「e」會有一個與此查詢很好匹配的鍵，因此它會獲得很高的注意力權重，其值（關於是母音的資訊）會流入當前位置。鍵和值會被附加到 KV 快取中，以便先前的位置可用。每個注意力頭計算其查詢與所有快取鍵之間的點積（按 \(\sqrt{d_{head}}\) 縮放），應用 softmax 得到注意力權重，並對快取值進行加權求和。所有頭的輸出被串聯起來，並通過 attn_wo 進行投影。值得強調的是，注意力區塊是 token 在位置 t 可以「查看」過去 0..t-1 的 token 的確切且唯一的地方。注意力是一種 token 通訊機制。
MLP 區塊。MLP 是「多層感知機」的縮寫，它是一個兩層的前饋網路：投影到嵌入維度的 4 倍，應用 ReLU，再投影回來。這是模型在每個位置進行大部分「思考」的地方。與注意力不同，這個計算完全局限於時間 t。Transformer 將通訊（注意力）與計算（MLP）交織在一起。
殘差連接。注意力區塊和 MLP 區塊都將其輸出加回到其輸入（x = [a + b for ...]）。這使得梯度可以直接流經網路，並使更深層的模型可訓練。
輸出。最終的隱藏狀態通過 lm_head 投影到詞彙量大小，為詞彙表中的每個 token 產生一個 logit。在我們的例子中，這只是 27 個數字。更高的 logit = 模型認為相應的 token 在下一個位置出現的可能性更大。
你可能會注意到我們在訓練期間使用了 KV 快取，這很不尋常。人們通常只將 KV 快取與推理聯繫起來。但 KV 快取在概念上始終存在，即使在訓練期間也是如此。在生產級實現中，它只是隱藏在處理序列中所有位置的向量化注意力計算中。由於 microgpt 一次處理一個 token（沒有批次維度，沒有並行時間步長），我們顯式地構建了 KV 快取。並且與典型的推理設置不同，在推理設置中 KV 快取持有分離的張量，在這裡，快取的鍵和值是計算圖中活躍的 Value 節點，所以我們實際上會對它們進行反向傳播。
現在我們將所有內容連接起來。訓練迴圈重複進行：(1) 選擇一個文件，(2) 對其 token 運行模型前向傳播，(3) 計算損失，(4) 反向傳播以獲得梯度，以及 (5) 更新參數。
分詞。每個訓練步驟選擇一個文件並用 BOS 包裹起來：名字「emma」變成 [BOS, e, m, m, a, BOS]。模型的任務是根據之前的 token 預測每個下一個 token。
前向傳播和損失。我們一次將 token 輸入模型，同時構建 KV 快取。在每個位置，模型輸出 27 個 logits，我們通過 softmax 將其轉換為概率。每個位置的損失是正確下一個 token 的負對數概率：\(-\log p( ext{target})\)。這被稱為交叉熵損失。直觀地說，損失衡量了預測錯誤的程度：模型對實際發生的事情有多驚訝。如果模型為正確的 token 分配概率 1.0，它根本不驚訝，損失為 0。如果它分配接近 0 的概率，模型非常驚訝，損失趨於 \(+\infty\)。我們對文件中的每個位置損失進行平均，得到一個單一的純量損失。
反向傳播。調用一次 loss.backward() 就會通過整個計算圖運行反向傳播，從損失一直回溯到 softmax、模型，並進入每個參數。之後，每個參數的 .grad 都告訴我們如何更改它以降低損失。
Adam 優化器。我們可以直接執行 p.data -= lr * p.grad（梯度下降），但 Adam 更聰明。它為每個參數維護兩個運行平均值：m 追蹤近期梯度的平均值（動量，像滾動的球），v 追蹤近期梯度平方的平均值（為每個參數自適應學習率）。m_hat 和 v_hat 是偏差校正，用於考慮 m 和 v 初始化為零且需要預熱的事實。學習率在訓練過程中線性衰減。更新後，我們將 .grad = 0 重置以進行下一步。
在 1,000 步之後，損失從約 3.3（在 27 個 token 中隨機猜測：\(-\log(1/27) \approx 3.3\)）下降到約 2.37。越低越好，最低可能值是 0（完美預測），所以仍有改進空間，但模型顯然正在學習名字的統計模式。
訓練完成後，我們可以從模型中採樣新名字。參數被凍結，我們只需在迴圈中運行前向傳播，將每個生成的 token 作為下一個輸入：
我們以 BOS token 開始每個採樣，它告訴模型「開始一個新名字」。模型產生 27 個 logits，我們將它們轉換為概率，然後根據這些概率隨機採樣一個 token。該 token 被作為下一個輸入反饋回來，我們重複此過程，直到模型再次產生 BOS（表示「我完成了」）或達到最大序列長度。
溫度參數控制隨機性。在 softmax 之前，我們將 logits 除以溫度。溫度為 1.0 時，直接從模型學到的分佈中採樣。較低的溫度（如這裡的 0.5）會使分佈變尖銳，使模型更保守，更有可能選擇其首選選項。接近 0 的溫度將始終選擇最可能的單一 token（貪婪解碼）。較高的溫度會使分佈變平，產生更多樣化但可能不太連貫的輸出。
你只需要 Python（無需 pip install，無依賴）：
該腳本在我的 macbook 上運行大約需要 1 分鐘。你將看到每一步打印出的損失：
觀察它從約 3.3（隨機）下降到約 2.37。這個數字越低，網路對序列中下一個 token 的預測就越好。在訓練結束時，訓練序列統計模式的知識被濃縮在模型參數中。固定這些參數後，我們現在可以生成新的、幻覺出來的名字。你將再次看到：
作為在你的電腦上運行該腳本的替代方案，你可以嘗試直接在此 Google Colab 筆記本上運行它，並向 Gemini 詢問有關它的問題。試著玩玩這個腳本！你可以嘗試不同的資料集。或者你可以訓練更長時間（增加 num_steps）或增加模型大小以獲得越來越好的結果。
要查看程式碼如何像洋蔥一樣一層一層地構建起來，建議的進程如下所示：
我創建了一個名為 build_microgpt.py 的 Gist，在 Revisions 中你可以看到所有這些版本以及每個步驟之間的差異。我認為這可能是逐步了解程式碼庫的一種有幫助的方式，你一次添加一個組件。
microgpt 包含了訓練和運行 GPT 的完整演算法精髓。但在這個專案和像 ChatGPT 這樣的生產級 LLM 之間，有許多事情會發生變化。沒有一件事情會改變核心演算法和整體佈局，但它們是使其能夠大規模運行的關鍵。按照相同的順序走過各個部分：
資料。生產級模型不使用 32K 個短名字，而是訓練數萬億個網際網路文本 token：網頁、書籍、程式碼等。資料會進行去重、品質過濾，並仔細混合不同領域的內容。
分詞器。生產級模型不使用單個字元，而是使用子詞分詞器，如 BPE（字元對編碼），它學習將經常共同出現的字元序列合併為單個 token。常見的單詞如「the」變成一個單獨的 token，罕見的單詞則被分解成幾個部分。這提供了約 100K 個 token 的詞彙量，並且效率更高，因為模型在每個位置看到更多內容。
自動微分。microgpt 在純 Python 中處理純量 Value 物件。生產級系統使用張量（大型多維數字陣列）並在 GPU/TPU 上運行，這些設備每秒執行數十億次浮點運算。PyTorch 等函式庫處理張量上的自動微分，而像 FlashAttention 這樣的 CUDA 核心會融合多個操作以提高速度。數學原理相同，只是對大量純量進行並行處理。
架構。microgpt 擁有 4,192 個參數。GPT-4 級別的模型擁有數千億個參數。總體而言，它是一個非常相似的 Transformer 神經網路，只是更寬（嵌入維度為 10,000+）且更深（100+ 層）。現代 LLM 還包含一些額外的積木類型，並改變它們的順序：例如 RoPE（旋轉位置嵌入）取代學習到的位置嵌入，GQA（分組查詢注意力）以減少 KV 快取大小，門控線性激活取代 ReLU，混合專家（MoE）層等。但注意力（通訊）和 MLP（計算）在殘差流上交織的核心結構得到了很好的保留。
訓練。生產級訓練不使用每個步驟一個文件，而是使用大型批次（每個步驟數百萬個 token）、梯度累積、混合精度（float16/bfloat16）和仔細的超參數調整。訓練前沿模型需要數千個 GPU 運行數月。
優化。microgpt 使用 Adam，並帶有簡單的線性學習率衰減，僅此而已。大規模優化本身就是一個獨立的學科。模型以降低的精度（bfloat16 甚至 fp8）在大型 GPU 集群上進行訓練以提高效率，這帶來了其自身的數值挑戰。優化器設置（學習率、權重衰減、beta 參數、預熱計劃、衰減計劃）必須精確調整，並且正確的值取決於模型大小、批次大小和資料集組成。縮放定律（例如 Chinchilla）指導如何在固定的計算預算之間分配模型大小和訓練 token 數量。在大規模情況下，任何這些細節的錯誤都可能浪費數百萬美元的計算成本，因此團隊會在投入全面訓練運行之前進行廣泛的小規模實驗來預測正確的設置。
訓練後。訓練結束後的基本模型（稱為「預訓練」模型）是一個文件補全器，而不是聊天機器人。將其變成 ChatGPT 分為兩個階段。首先，SFT（監督式微調）：你只需將文件替換為經過策劃的對話並繼續訓練。演算法上沒有變化。其次，RL（強化學習）：模型生成回應，這些回應會被評分（由人類、另一個「評審」模型或演算法），然後模型從該反饋中學習。從根本上說，模型仍在訓練文件，但這些文件現在由模型本身產生的 token 組成。
推理。為數百萬用戶提供模型服務需要自己的工程堆疊：將請求批次處理、KV 快取管理和分頁（vLLM 等）、用於加速的投機解碼、量化（以 int8/int4 而非 float16 運行）以減少記憶體，以及將模型分佈到多個 GPU 上。從根本上說，我們仍然在預測序列中的下一個 token，但投入了大量的工程工作來使其更快。
所有這些都是重要的工程和研究貢獻，但如果你理解了 microgpt，你就理解了演算法的精髓。
模型是否「理解」任何東西？這是一個哲學問題，但從機械上講：沒有什麼魔法發生。模型是一個將輸入 token 映射到下一個 token 的概率分佈的大型數學函數。在訓練期間，參數會被調整，以使正確的下一個 token 更具概率。這是否構成「理解」取決於你，但其機制完全包含在上面的 200 行程式碼中。
為什麼它有效？模型有數千個可調參數，優化器在每一步都會微調它們一點點以降低損失。經過許多步驟，參數會收斂到捕捉資料統計規律的值。對於名字來說，這意味著：名字通常以輔音開頭，「qu」傾向於一起出現，名字很少連續出現三個輔音等。模型不會學習明確的規則，它學習的是一個恰好反映了這些規則的概率分佈。
這與 ChatGPT 有何關聯？ChatGPT 就是這個核心迴圈（預測下一個 token，採樣，重複）的巨大規模化版本，並經過訓練後使其具有對話能力。當你與它聊天時，系統提示、你的訊息和它的回覆都只是序列中的 token。模型一次一個 token 地補全文件，就像 microgpt 補全名字一樣。
「幻覺」是怎麼回事？模型通過從概率分佈中採樣來生成 token。它沒有真理的概念，它只知道哪些序列在給定的訓練資料下是統計上合理的。microgpt 「幻覺」出像「karia」這樣一個名字，與 ChatGPT 自信地陳述一個錯誤事實是相同的現象。兩者都是聽起來合理但碰巧不真實的補全。
為什麼它這麼慢？microgpt 在純 Python 中一次處理一個純量。一個訓練步驟需要幾秒鐘。在 GPU 上執行相同的數學運算可以並行處理數百萬個純量，速度要快幾個數量級。
我能讓它生成更好的名字嗎？是的。訓練更長時間（增加 num_steps），使模型更大（n_embd、n_layer、n_head），或使用更大的資料集。這些是與大規模情況下同樣重要的調整參數。
如果我更改資料集會怎樣？模型將學習資料中的任何模式。替換為城市名稱、寶可夢名稱、英文單字或短詩的檔案，模型將學會生成這些內容。程式碼的其他部分無需更改。