從零復(fù)現(xiàn)Llama3代碼庫(kù)爆火，大神Kapathy一鍵三連，GitHub狂攬2k+

西風(fēng) 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

讓大神Andrej Karpathy一鍵三連??（點(diǎn)贊+轉(zhuǎn)發(fā)+評(píng)論），一個(gè)教你從頭開始實(shí)現(xiàn)Llama3的代碼庫(kù)爆火。

X上轉(zhuǎn)贊收藏量超6.8k，GitHub攬星2k+。

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

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

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

完全展開后，比起模塊相互嵌套和調(diào)用時(shí)，更容易理解每一步具體在做什么。

網(wǎng)友們對(duì)其也是贊不絕口，紛紛致敬：

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

（量子位在不改變?cè)獾幕A(chǔ)上，進(jìn)行了編譯整理）

從頭實(shí)現(xiàn)llama3

在運(yùn)行納哥提供的文件前，大伙兒需要預(yù)先下載Meta官方提供的Llama3模型權(quán)重。

納哥表示自己沒搞分詞器，推薦用Karpathy的現(xiàn)成簡(jiǎn)潔版BPE代碼。

PS：

“字節(jié)級(jí)（byte-level）”BPE算法，在UTF-8編碼的字符串上運(yùn)行，廣泛應(yīng)用于大模型分詞。Karpathy提供的這個(gè)代碼庫(kù)包含兩個(gè)分詞器，都能在給定文本上訓(xùn)練分詞器的詞匯表和合并規(guī)則、將文本編碼為token、將token解碼為文本。

讀取模型文件的方式通常取決于model classes的編寫方式以及class中變量的命名。但由于納哥是從頭開始實(shí)現(xiàn)Llama3，所以將逐個(gè)張量地讀取文件內(nèi)容。

通過此配置可以推斷出模型的結(jié)構(gòu)和參數(shù)信息，例如模型包含的Transformer層數(shù)、多頭注意力塊中的頭數(shù)，以及詞匯表的大小等細(xì)節(jié)。

將文本轉(zhuǎn)換為token時(shí)，納哥使用tiktoken作為分詞器。

接下來，納哥展示了在代碼中將token轉(zhuǎn)換為高維的嵌入表示。這是代碼庫(kù)中唯一使用內(nèi)置神經(jīng)網(wǎng)絡(luò)模塊的部分。

[17×1]的token矩陣變成了[17×4096]的嵌入矩陣。也就是說，每個(gè)token被轉(zhuǎn)換為一個(gè)長(zhǎng)度為4096的嵌入向量，總共有17個(gè)這樣的嵌入向量。

然后，納哥對(duì)嵌入進(jìn)行RMS歸一化。經(jīng)過這一步后，嵌入的形狀不會(huì)改變，只有數(shù)值被歸一化了。納哥強(qiáng)調(diào)需要一個(gè)norm_eps，避免意外將RMS值設(shè)為0導(dǎo)致除以0的錯(cuò)誤。

以下是公式：

構(gòu)建Transformer的第一層，進(jìn)行歸一化處理，從模型字典中訪問layer.0（即第一層）。歸一化之后，張量的形狀仍然是[17×4096]，與嵌入時(shí)相同，但數(shù)值已被歸一化。

跟著納哥從頭實(shí)現(xiàn)注意力機(jī)制，加載Transformer第一層的注意力頭。

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

納哥表示乍一看有點(diǎn)奇怪，因?yàn)槔硐肭闆r是每個(gè)注意力頭的q、k、v和o向量是獨(dú)立的。而代碼作者將它們捆綁在一起，是為了方便并行計(jì)算注意力頭的矩陣乘法。

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

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

在這里，納哥訪問了第一層第一個(gè)注意力頭的query權(quán)重矩陣，query權(quán)重矩陣的大小是[128×4096]。

將query權(quán)重矩陣與token嵌入相乘，獲得每個(gè)token的query向量。結(jié)果的形狀為[17×128]，有17個(gè)token，每個(gè)token對(duì)應(yīng)一個(gè)長(zhǎng)度為128的query向量。

接下來需要位置編碼。

現(xiàn)在已經(jīng)為prompt中的每個(gè)token生成了query向量，但每個(gè)單獨(dú)的query向量并不知道它在prompt中的具體位置。

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

在這個(gè)prompt中，使用了三次”the”，需要根據(jù)它們?cè)趐rompt中的位置，使這三個(gè)”the”token的query向量有所不同（每個(gè)向量的大小為[1×128]）。

通過使用RoPE（旋轉(zhuǎn)位置嵌入）來進(jìn)行這些旋轉(zhuǎn)操作。

上一步中，納哥將query向量分成對(duì)，并對(duì)每一對(duì)應(yīng)用一個(gè)旋轉(zhuǎn)角度偏移。

由此，得到的向量大小為 [17x64x2]，這是將長(zhǎng)度為128的query向量對(duì)每個(gè)prompt中的token分成64對(duì)。這64對(duì)中的每一對(duì)都會(huì)根據(jù)m*(theta) 進(jìn)行旋轉(zhuǎn)，其中m是要旋轉(zhuǎn)query的token的位置。

使用復(fù)數(shù)的點(diǎn)積來旋轉(zhuǎn)一個(gè)向量：

現(xiàn)在每個(gè)token的query元素都有一個(gè)復(fù)數(shù)（角度變化向量），可以將query向量（之前分成的對(duì)）轉(zhuǎn)換為復(fù)數(shù)，然后通過點(diǎn)積根據(jù)位置旋轉(zhuǎn)query向量。

獲得旋轉(zhuǎn)后的向量后，可以通過將復(fù)數(shù)重新視為實(shí)數(shù)來得到成對(duì)的query向量。

旋轉(zhuǎn)后的對(duì)現(xiàn)在已經(jīng)合并，有一個(gè)新的query向量（旋轉(zhuǎn)后的query向量），其形狀為[17×128]，其中17是token的數(shù)量，128是query向量的維度。

key與query幾乎相同。

納哥表示自己不會(huì)詳細(xì)講解key的數(shù)學(xué)原理，只需要記住以下幾點(diǎn)：

key生成的key向量維度也是128；key的權(quán)重只有query的四分之一，這是因?yàn)閗ey的權(quán)重在同一時(shí)間內(nèi)被4個(gè)頭共享，來減少計(jì)算量；key也會(huì)旋轉(zhuǎn)添加位置信息，原因與query相同。

此時(shí)，納哥已經(jīng)為每個(gè)token獲得了旋轉(zhuǎn)后的query和key。每個(gè)query和key現(xiàn)在的形狀都是[17×128]。

下一步，納哥將對(duì)query矩陣和key矩陣進(jìn)行相乘操作。這樣做會(huì)生成一個(gè)評(píng)分矩陣，將每個(gè)token關(guān)聯(lián)起來。這些評(píng)分描述了每個(gè)token的query與每個(gè)token的key之間的相關(guān)性，這就是自注意力機(jī)制。

注意力評(píng)分矩陣（qk_per_token）的形狀為[17×17]，其中17是prompt中的token數(shù)量。

接下來需要對(duì)query key評(píng)分進(jìn)行掩碼處理。在Llama3的訓(xùn)練過程中，未來token的qk評(píng)分是被掩碼的，只通過過去的token來預(yù)測(cè)token。

因此，在推理時(shí)，要將未來的token評(píng)分設(shè)置為0。

接下來是value，接近注意力機(jī)制的最后一步。

這些評(píng)分（0-1）用于確定每個(gè)token使用多少value矩陣。

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

第一層，第一個(gè)注意力頭的value權(quán)重矩陣如下所示：

然后是value向量。

使用value權(quán)重來獲取每個(gè)token的注意力值，矩陣的大小是[17×128]，其中17是prompt中的token數(shù)量，128是每個(gè)token的value向量的維度。

注意力：與每個(gè)token的value相乘后得到的注意力向量的形狀為[17×128]。

現(xiàn)在有了第一層第一個(gè)頭的注意力value。然后納哥運(yùn)行一個(gè)循環(huán)，對(duì)第一層的每個(gè)頭執(zhí)行與上面的計(jì)算完全相同的數(shù)學(xué)運(yùn)算。

然后得到了第一層所有32個(gè)頭的qkv_attention矩陣，接下來將所有注意力得分合并成一個(gè)大小為[17×4096]的大矩陣。

對(duì)于第0層注意力機(jī)制的最后步驟，其一是將注意力得分矩陣與權(quán)重矩陣相乘。

這是一個(gè)簡(jiǎn)單的線性層，所以只需進(jìn)行矩陣乘法。

現(xiàn)在得到了注意力機(jī)制后的嵌入value變化，應(yīng)該被添加到原始的token嵌入中。

對(duì)嵌入增量進(jìn)行歸一化處理，然后通過嵌入增量運(yùn)行一個(gè)前饋神經(jīng)網(wǎng)絡(luò)。

在Llama3中，加載前饋權(quán)重并實(shí)現(xiàn)前饋網(wǎng)絡(luò)。使用了一種名為SwiGLU的前饋網(wǎng)絡(luò)，這種網(wǎng)絡(luò)結(jié)構(gòu)在模型需要的時(shí)候，能夠有效地增加非線性。

現(xiàn)在完成了第一層之后每個(gè)token的新嵌入?，F(xiàn)在只剩下31層了，只需通過一個(gè)循環(huán)來完成。

納哥表示可以將這個(gè)編輯后的嵌入想象成包含了第一層中所有查詢信息的嵌入。隨著層數(shù)的增加，每一層都會(huì)對(duì)輸入的信息進(jìn)行越來越復(fù)雜的處理，直到最終得到一個(gè)能夠全面了解下一個(gè)需要預(yù)測(cè)的token的嵌入。

之前做的所有事情，對(duì)每一層都重復(fù)一次。

然后得到了最終的嵌入，這是模型對(duì)下一個(gè)token的最優(yōu)預(yù)測(cè)。這個(gè)嵌入的形狀與常規(guī)的token嵌入相同，為[17×4096]，其中17是token的數(shù)量，4096是嵌入的維度。

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

使用輸出解碼器將最終的嵌入轉(zhuǎn)換成一個(gè)token。

接下來看納哥使用最后一個(gè)token的嵌入來預(yù)測(cè)下一個(gè)value，希望預(yù)測(cè)的結(jié)果是42。

因?yàn)楦鶕?jù)《銀河系漫游指南》一書中的說法，42是“生命、宇宙及一切的終極問題的答案”。大多數(shù)LLM在這里都會(huì)回答42，這將驗(yàn)證整個(gè)代碼的正確性。

模型預(yù)測(cè)下一個(gè)token的編號(hào)為2983。這個(gè)編號(hào)對(duì)應(yīng)數(shù)字42嗎？

OK，結(jié)束。

“讓研究變得更加觸手可及”

簡(jiǎn)單介紹一下Nishant Aklecha。

Nishant Aklecha是構(gòu)建和改進(jìn)定制語言模型平臺(tái)Glaive AI的研究員，曾任職于摩根士丹利，負(fù)責(zé)訓(xùn)練和微調(diào)大語言模型。

此外，他還和朋友一同創(chuàng)立了一個(gè)研究實(shí)驗(yàn)室，名為A10（AAAAAAAAAA）。

他們的目標(biāo)可以總結(jié)成一句話：讓研究變得更加觸手可及。

除了放出這個(gè)代碼庫(kù)，Nishant Aklecha可謂好人做到底。

網(wǎng)友想更好地理解這個(gè)代碼庫(kù)的內(nèi)容，Nishant直接一個(gè)YouTube視頻甩了過來：

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

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

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