新火種
2024-05-22 
從零復(fù)現(xiàn)Llama3代碼庫爆火，大神Kapathy一鍵三連，GitHub狂攬2k+
讓大神Andrej Karpathy一鍵三連??（點贊+轉(zhuǎn)發(fā)+評論），一個教你從頭開始實現(xiàn)Llama3的代碼庫爆火。
X上轉(zhuǎn)贊收藏量超6.8k，GitHub攬星2k+。
火就火在，它教你從頭用Meta開源的權(quán)重進行推理，詳細解釋和展開了注意力機制中多個頭的矩陣乘法、位置編碼以及所有中間層。
換句話說，他解釋了每行代碼都在干啥。
Karpathy看后直呼打造者Nishant Aklecha（后文暫稱“納哥”）是個有品的人：
網(wǎng)友們對其也是贊不絕口，紛紛致敬：
話不多說，一起來看納哥是如何手把手教的。
（新火種在不改變原意的基礎(chǔ)上，進行了編譯整理）
從頭實現(xiàn)llama3
在運行納哥提供的文件前，大伙兒需要預(yù)先下載Meta官方提供的Llama3模型權(quán)重。
納哥表示自己沒搞分詞器，推薦用Karpathy的現(xiàn)成簡潔版BPE代碼。
PS：
“字節(jié)級（byte-level）”BPE算法，在UTF-8編碼的字符串上運行，廣泛應(yīng)用于大模型分詞。Karpathy提供的這個代碼庫包含兩個分詞器，都能在給定文本上訓(xùn)練分詞器的詞匯表和合并規(guī)則、將文本編碼為token、將token解碼為文本。
讀取模型文件的方式通常取決于model classes的編寫方式以及class中變量的命名。但由于納哥是從頭開始實現(xiàn)Llama3，所以將逐個張量地讀取文件內(nèi)容。
通過此配置可以推斷出模型的結(jié)構(gòu)和參數(shù)信息，例如模型包含的Transformer層數(shù)、多頭注意力塊中的頭數(shù)，以及詞匯表的大小等細節(jié)。
將文本轉(zhuǎn)換為token時，納哥使用tiktoken作為分詞器。
接下來，納哥展示了在代碼中將token轉(zhuǎn)換為高維的嵌入表示。這是代碼庫中唯一使用內(nèi)置神經(jīng)網(wǎng)絡(luò)模塊的部分。
[17×1]的token矩陣變成了[17×4096]的嵌入矩陣。也就是說，每個token被轉(zhuǎn)換為一個長度為4096的嵌入向量，總共有17個這樣的嵌入向量。
然后，納哥對嵌入進行RMS歸一化。經(jīng)過這一步后，嵌入的形狀不會改變，只有數(shù)值被歸一化了。納哥強調(diào)需要一個norm_eps，避免意外將RMS值設(shè)為0導(dǎo)致除以0的錯誤。
以下是公式：
構(gòu)建Transformer的第一層，進行歸一化處理，從模型字典中訪問layer.0（即第一層）。歸一化之后，張量的形狀仍然是[17×4096]，與嵌入時相同，但數(shù)值已被歸一化。
跟著納哥從頭實現(xiàn)注意力機制，加載Transformer第一層的注意力頭。
從模型中加載query、key、value和output向量時，它們的形狀分別是 [4096×4096]、[1024×4096]、[1024×4096] 和 [4096×4096]。
納哥表示乍一看有點奇怪，因為理想情況是每個注意力頭的q、k、v和o向量是獨立的。而代碼作者將它們捆綁在一起，是為了方便并行計算注意力頭的矩陣乘法。
把所有這些向量解包開來：
下一步，納哥將從多個注意力頭中解包query，解包后的形狀是[32x128x4096]，32是Llama3中的注意力頭數(shù)量，128是query向量的大小，4096是token嵌入的大小。
在這里，納哥訪問了第一層第一個注意力頭的query權(quán)重矩陣，query權(quán)重矩陣的大小是[128×4096]。
將query權(quán)重矩陣與token嵌入相乘，獲得每個token的query向量。結(jié)果的形狀為[17×128]，有17個token，每個token對應(yīng)一個長度為128的query向量。
接下來需要位置編碼。
現(xiàn)在已經(jīng)為prompt中的每個token生成了query向量，但每個單獨的query向量并不知道它在prompt中的具體位置。
例如，query：“the answer to the ultimate question of life, the universe, and everything is ”（生命、宇宙和一切的終極問題的答案是）。
在這個prompt中，使用了三次”the”，需要根據(jù)它們在prompt中的位置，使這三個”the”token的query向量有所不同（每個向量的大小為[1×128]）。
通過使用RoPE（旋轉(zhuǎn)位置嵌入）來進行這些旋轉(zhuǎn)操作。
上一步中，納哥將query向量分成對，并對每一對應(yīng)用一個旋轉(zhuǎn)角度偏移。
由此，得到的向量大小為 [17x64x2]，這是將長度為128的query向量對每個prompt中的token分成64對。這64對中的每一對都會根據(jù)m*(theta) 進行旋轉(zhuǎn)，其中m是要旋轉(zhuǎn)query的token的位置。
使用復(fù)數(shù)的點積來旋轉(zhuǎn)一個向量：
現(xiàn)在每個token的query元素都有一個復(fù)數(shù)（角度變化向量），可以將query向量（之前分成的對）轉(zhuǎn)換為復(fù)數(shù)，然后通過點積根據(jù)位置旋轉(zhuǎn)query向量。
獲得旋轉(zhuǎn)后的向量后，可以通過將復(fù)數(shù)重新視為實數(shù)來得到成對的query向量。
旋轉(zhuǎn)后的對現(xiàn)在已經(jīng)合并，有一個新的query向量（旋轉(zhuǎn)后的query向量），其形狀為[17×128]，其中17是token的數(shù)量，128是query向量的維度。
key與query幾乎相同。
納哥表示自己不會詳細講解key的數(shù)學(xué)原理，只需要記住以下幾點：
key生成的key向量維度也是128；key的權(quán)重只有query的四分之一，這是因為key的權(quán)重在同一時間內(nèi)被4個頭共享，來減少計算量；key也會旋轉(zhuǎn)添加位置信息，原因與query相同。
此時，納哥已經(jīng)為每個token獲得了旋轉(zhuǎn)后的query和key。每個query和key現(xiàn)在的形狀都是[17×128]。
下一步，納哥將對query矩陣和key矩陣進行相乘操作。這樣做會生成一個評分矩陣，將每個token關(guān)聯(lián)起來。這些評分描述了每個token的query與每個token的key之間的相關(guān)性，這就是自注意力機制。
注意力評分矩陣（qk_per_token）的形狀為[17×17]，其中17是prompt中的token數(shù)量。
接下來需要對query key評分進行掩碼處理。在Llama3的訓(xùn)練過程中，未來token的qk評分是被掩碼的，只通過過去的token來預(yù)測token。
因此，在推理時，要將未來的token評分設(shè)置為0。
接下來是value，接近注意力機制的最后一步。
這些評分（0-1）用于確定每個token使用多少value矩陣。
和key一樣，value的權(quán)重也在每4個注意力頭之間共享，所以下面value權(quán)重矩陣的形狀是[8x128x4096]。
第一層，第一個注意力頭的value權(quán)重矩陣如下所示：
然后是value向量。
使用value權(quán)重來獲取每個token的注意力值，矩陣的大小是[17×128]，其中17是prompt中的token數(shù)量，128是每個token的value向量的維度。
注意力：與每個token的value相乘后得到的注意力向量的形狀為[17×128]。
現(xiàn)在有了第一層第一個頭的注意力value。然后納哥運行一個循環(huán)，對第一層的每個頭執(zhí)行與上面的計算完全相同的數(shù)學(xué)運算。
然后得到了第一層所有32個頭的qkv_attention矩陣，接下來將所有注意力得分合并成一個大小為[17×4096]的大矩陣。
對于第0層注意力機制的最后步驟，其一是將注意力得分矩陣與權(quán)重矩陣相乘。
這是一個簡單的線性層，所以只需進行矩陣乘法。
現(xiàn)在得到了注意力機制后的嵌入value變化，應(yīng)該被添加到原始的token嵌入中。
對嵌入增量進行歸一化處理，然后通過嵌入增量運行一個前饋神經(jīng)網(wǎng)絡(luò)。
在Llama3中，加載前饋權(quán)重并實現(xiàn)前饋網(wǎng)絡(luò)。使用了一種名為SwiGLU的前饋網(wǎng)絡(luò)，這種網(wǎng)絡(luò)結(jié)構(gòu)在模型需要的時候，能夠有效地增加非線性。
現(xiàn)在完成了第一層之后每個token的新嵌入?，F(xiàn)在只剩下31層了，只需通過一個循環(huán)來完成。
納哥表示可以將這個編輯后的嵌入想象成包含了第一層中所有查詢信息的嵌入。隨著層數(shù)的增加，每一層都會對輸入的信息進行越來越復(fù)雜的處理，直到最終得到一個能夠全面了解下一個需要預(yù)測的token的嵌入。
之前做的所有事情，對每一層都重復(fù)一次。
然后得到了最終的嵌入，這是模型對下一個token的最優(yōu)預(yù)測。這個嵌入的形狀與常規(guī)的token嵌入相同，為[17×4096]，其中17是token的數(shù)量，4096是嵌入的維度。
最后，將嵌入解碼成token值。
使用輸出解碼器將最終的嵌入轉(zhuǎn)換成一個token。
接下來看納哥使用最后一個token的嵌入來預(yù)測下一個value，希望預(yù)測的結(jié)果是42。
因為根據(jù)《銀河系漫游指南》一書中的說法，42是“生命、宇宙及一切的終極問題的答案”。大多數(shù)LLM在這里都會回答42，這將驗證整個代碼的正確性。
模型預(yù)測下一個token的編號為2983。這個編號對應(yīng)數(shù)字42嗎？
OK，結(jié)束。
“讓研究變得更加觸手可及”
簡單介紹一下Nishant Aklecha。
Nishant Aklecha是構(gòu)建和改進定制語言模型平臺Glaive AI的研究員，曾任職于摩根士丹利，負責(zé)訓(xùn)練和微調(diào)大語言模型。
此外，他還和朋友一同創(chuàng)立了一個研究實驗室，名為A10（AAAAAAAAAA）。
他們的目標可以總結(jié)成一句話：讓研究變得更加觸手可及。
除了放出這個代碼庫，Nishant Aklecha可謂好人做到底。
網(wǎng)友想更好地理解這個代碼庫的內(nèi)容，Nishant直接一個YouTube視頻甩了過來：
之前Nishant Aklecha還曾寫過一篇Blog，詳解了潛在一致性模型（LCM），同樣收獲了不少好評。
啥也不說了，感興趣的家人們趕緊碼住吧。
Tags:
相關(guān)推薦
- 免責(zé)聲明
- 本文所包含的觀點僅代表作者個人看法，不代表新火種的觀點。在新火種上獲取的所有信息均不應(yīng)被視為投資建議。新火種對本文可能提及或鏈接的任何項目不表示認可。 交易和投資涉及高風(fēng)險，讀者在采取與本文內(nèi)容相關(guān)的任何行動之前，請務(wù)必進行充分的盡職調(diào)查。最終的決策應(yīng)該基于您自己的獨立判斷。新火種不對因依賴本文觀點而產(chǎn)生的任何金錢損失負任何責(zé)任。
