從零復現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
本文轉自 量子位 ，作者：量子位
