從零復現Llama3代碼庫爆火，大神Kapathy一鍵三連，GitHub狂攬2k+

讓大神Andrej Karpathy一鍵三連❤️（點贊+轉發+評論），一個教你從頭開始實現Llama3的代碼庫爆火。

X上轉贊收藏量超6.8k，GitHub攬星2k+。

火就火在，它教你從頭用Meta開源的權重進行推理，詳細解釋和展開了注意力機制中多個頭的矩陣乘法、位置編碼以及所有中間層。

換句話說，他解釋了每行代碼都在幹啥。

Karpathy看後直呼打造者Nishant Aklecha（後文暫稱“納哥”）是個有品的人：

網友們對其也是讚不絕口，紛紛致敬：

話不多說，一起來看納哥是如何手把手教的。

（量子位在不改變原意的基礎上，進行了編譯整理）

從頭實現llama3

在運行納哥提供的文件前，大夥兒需要預先下載Meta官方提供的Llama3模型權重。

納哥表示自己沒搞分詞器，推薦用Karpathy的現成簡潔版BPE代碼。

PS：

“字節級（byte-level）”BPE算法，在UTF-8編碼的字符串上運行，廣泛應用於大模型分詞。Karpathy提供的這個代碼庫包含兩個分詞器，都能在給定文本上訓練分詞器的詞彙表和合並規則、將文本編碼爲token、將token解碼爲文本。

讀取模型文件的方式通常取決於model classes的編寫方式以及class中變量的命名。但由於納哥是從頭開始實現Llama3，所以將逐個張量地讀取文件內容。

通過此配置可以推斷出模型的結構和參數信息，例如模型包含的Transformer層數、多頭注意力塊中的頭數，以及詞彙表的大小等細節。

將文本轉換爲token時，納哥使用tiktoken作爲分詞器。

接下來，納哥展示了在代碼中將token轉換爲高維的嵌入表示。這是代碼庫中唯一使用內置神經網絡模塊的部分。

[17x1]的token矩陣變成了[17x4096]的嵌入矩陣。也就是說，每個token被轉換爲一個長度爲4096的嵌入向量，總共有17個這樣的嵌入向量。

然後，納哥對嵌入進行RMS歸一化。經過這一步後，嵌入的形狀不會改變，只有數值被歸一化了。納哥強調需要一個norm_eps，避免意外將RMS值設爲0導致除以0的錯誤。

以下是公式：

構建Transformer的第一層，進行歸一化處理，從模型字典中訪問layer.0（即第一層）。歸一化之後，張量的形狀仍然是[17x4096]，與嵌入時相同，但數值已被歸一化。

跟着納哥從頭實現注意力機制，加載Transformer第一層的注意力頭。

從模型中加載query、key、value和output向量時，它們的形狀分別是 [4096x4096]、[1024x4096]、[1024x4096] 和 [4096x4096]。

納哥表示乍一看有點奇怪，因爲理想情況是每個注意力頭的q、k、v和o向量是獨立的。而代碼作者將它們捆綁在一起，是爲了方便並行計算注意力頭的矩陣乘法。

把所有這些向量解包開來：

下一步，納哥將從多個注意力頭中解包query，解包後的形狀是[32x128x4096]，32是Llama3中的注意力頭數量，128是query向量的大小，4096是token嵌入的大小。

在這裡，納哥訪問了第一層第一個注意力頭的query權重矩陣，query權重矩陣的大小是[128x4096]。

將query權重矩陣與token嵌入相乘，獲得每個token的query向量。結果的形狀爲[17x128]，有17個token，每個token對應一個長度爲128的query向量。

接下來需要位置編碼。

現在已經爲prompt中的每個token生成了query向量，但每個單獨的query向量並不知道它在prompt中的具體位置。

例如，query：“the answer to the ultimate question of life, the universe, and everything is ”（生命、宇宙和一切的終極問題的答案是）。

在這個prompt中，使用了三次”the”，需要根據它們在prompt中的位置，使這三個”the”token的query向量有所不同（每個向量的大小爲[1x128]）。

通過使用RoPE（旋轉位置嵌入）來進行這些旋轉操作。

上一步中，納哥將query向量分成對，並對每一對應用一個旋轉角度偏移。

由此，得到的向量大小爲 [17x64x2]，這是將長度爲128的query向量對每個prompt中的token分成64對。這64對中的每一對都會根據m*(theta) 進行旋轉，其中m是要旋轉query的token的位置。

使用複數的點積來旋轉一個向量：

現在每個token的query元素都有一個複數（角度變化向量），可以將query向量（之前分成的對）轉換爲複數，然後通過點積根據位置旋轉query向量。

獲得旋轉後的向量後，可以通過將複數重新視爲實數來得到成對的query向量。

旋轉後的對現在已經合併，有一個新的query向量（旋轉後的query向量），其形狀爲[17x128]，其中17是token的數量，128是query向量的維度。

key與query幾乎相同。

納哥表示自己不會詳細講解key的數學原理，只需要記住以下幾點：

key生成的key向量維度也是128；key的權重只有query的四分之一，這是因爲key的權重在同一時間內被4個頭共享，來減少計算量；key也會旋轉添加位置信息，原因與query相同。

此時，納哥已經爲每個token獲得了旋轉後的query和key。每個query和key現在的形狀都是[17x128]。

下一步，納哥將對query矩陣和key矩陣進行相乘操作。這樣做會生成一個評分矩陣，將每個token關聯起來。這些評分描述了每個token的query與每個token的key之間的相關性，這就是自注意力機制。

注意力評分矩陣（qk_per_token）的形狀爲[17x17]，其中17是prompt中的token數量。

接下來需要對query key評分進行掩碼處理。在Llama3的訓練過程中，未來token的qk評分是被掩碼的，只通過過去的token來預測token。

因此，在推理時，要將未來的token評分設置爲0。

接下來是value，接近注意力機制的最後一步。

這些評分（0-1）用於確定每個token使用多少value矩陣。

和key一樣，value的權重也在每4個注意力頭之間共享，所以下面value權重矩陣的形狀是[8x128x4096]。

第一層，第一個注意力頭的value權重矩陣如下所示：

然後是value向量。

使用value權重來獲取每個token的注意力值，矩陣的大小是[17x128]，其中17是prompt中的token數量，128是每個token的value向量的維度。

注意力：與每個token的value相乘後得到的注意力向量的形狀爲[17x128]。

現在有了第一層第一個頭的注意力value。然後納哥運行一個循環，對第一層的每個頭執行與上面的計算完全相同的數學運算。

然後得到了第一層所有32個頭的qkv_attention矩陣，接下來將所有注意力得分合併成一個大小爲[17x4096]的大矩陣。

對於第0層注意力機制的最後步驟，其一是將注意力得分矩陣與權重矩陣相乘。

這是一個簡單的線性層，所以只需進行矩陣乘法。

現在得到了注意力機制後的嵌入value變化，應該被添加到原始的token嵌入中。

對嵌入增量進行歸一化處理，然後通過嵌入增量運行一個前饋神經網絡。

在Llama3中，加載前饋權重並實現前饋網絡。使用了一種名爲SwiGLU的前饋網絡，這種網絡結構在模型需要的時候，能夠有效地增加非線性。

現在完成了第一層之後每個token的新嵌入。現在只剩下31層了，只需通過一個循環來完成。

納哥表示可以將這個編輯後的嵌入想象成包含了第一層中所有查詢信息的嵌入。隨着層數的增加，每一層都會對輸入的信息進行越來越複雜的處理，直到最終得到一個能夠全面瞭解下一個需要預測的token的嵌入。

之前做的所有事情，對每一層都重複一次。

然後得到了最終的嵌入，這是模型對下一個token的最優預測。這個嵌入的形狀與常規的token嵌入相同，爲[17x4096]，其中17是token的數量，4096是嵌入的維度。

最後，將嵌入解碼成token值。

使用輸出解碼器將最終的嵌入轉換成一個token。

接下來看納哥使用最後一個token的嵌入來預測下一個value，希望預測的結果是42。

因爲根據《銀河系漫遊指南》一書中的說法，42是“生命、宇宙及一切的終極問題的答案”。大多數LLM在這裡都會回答42，這將驗證整個代碼的正確性。

模型預測下一個token的編號爲2983。這個編號對應數字42嗎？

OK，結束。

簡單介紹一下Nishant Aklecha。

Nishant Aklecha是構建和改進定製語言模型平臺Glaive AI的研究員，曾任職於摩根士丹利，負責訓練和微調大語言模型。

此外，他還和朋友一同創立了一個研究實驗室，名爲A10（AAAAAAAAAA）。

他們的目標可以總結成一句話：讓研究變得更加觸手可及。

除了放出這個代碼庫，Nishant Aklecha可謂好人做到底。

網友想更好地理解這個代碼庫的內容，Nishant直接一個YouTube視頻甩了過來：

之前Nishant Aklecha還曾寫過一篇Blog，詳解了潛在一致性模型（LCM），同樣收穫了不少好評。

啥也不說了，感興趣的家人們趕緊碼住吧。

GitHub鏈接：https://github.com/naklecha/llama3-from-scratch

參考鏈接：[1]https://x.com/naklecha/status/1792244347225641338[2]https://naklecha.notion.site/explained-latent-consistency-models-13a9290c0fd3427d8d1a1e0bed97bde2[3]https://www.youtube.com/watch?v=o29P0Kpobz0&t=530s[4]https://www.youtube.com/watch?v=eMlx5fFNoYc