LLM演進史-重建GPT2 -(1)

本文基於 Andrej Karpathy 的 4 小時復現 GPT-2，個人看完後覺得是非常好的視頻，這是 LLM 演進史的終結篇，本文基於此進行文字版本的補充。前面的內容請參考 https://blog.nagi.fun 主頁的內容，該博主已經寫得非常詳盡了。

本系列準備一共分為三個部分，分別是主體部分實現，加速實現以及分佈式訓練。

實現 GPT-2 nn.Module#

Config 配置#

@dataclass
class GPTConfig():
    block_size: int=1024     # 序列長度限制（上下文窗口長度）
    vocab_size: int=50527    # 詞表大小
    n_layer: int=12          # Transformer層數
    n_head: int=12           # 注意力頭數量
    n_embd: int=768          # 嵌入維度（每個token的向量長度）

@dataclass裝飾器定義了一個名為GPTConfig的配置類

(如果不懂裝飾器可以自行 CSDN 或者知乎)

為什麼要用 dataclass：

•普通類需要手動編寫__init__方法，裝飾後非常簡單
•支持顯式聲明，且直接pirnt(GPTConfig(n_head=16))可以直接打印出參數

BackBone#


class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
		        # word token embedding
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            # word position embedding
            wpe = nn.Embedding(config.block_size, config.n_embd),
            # 主體block
            h = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
            # word token embedding,
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.attn = CasualSelfAttention(config)
        self.mlp = mlp(config)
    
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

transformer

self.transformer：transformer 架構本體

nn.ModuleDict：nn.Module 裡面的字典，nn.ModuleDict(ln_f = nn.LayerNorm(config.n_embd) ,)可以理解成{ln_f: nn.LayerNorm(config.n_embd)}

wte ：word 到 embedding 的線性層 [詞表的長度，詞嵌入維度]，把詞變為特徵向量

wpe ：word position 到 embedding 的線性層 [序列長度，詞嵌入維度]，把位置信息變為特徵向量

h ：transformer 的 encoder 主體，每一個 Block 由一個 attention 和一個 mlp 組成

ln_f ：LayerNorm，對於 Pre-Norm 後得到的大方差做一個歸一化，後續有解釋

lm_head ：最後一層輸出，把詞的特徵向量轉為具體的詞

Block ：Transformer 由多個相同的 Block 組成

Tips⚠️：此處會發現 GPT2 的 LN 層在 attention 和 mlp 之前，而與上圖原文中的 LN 層和殘差 (先殘差再歸一化) 連接在一起不同。

Karpathy的解釋是：原本的模型是先殘差連接再進行LN歸一化處理，這表明連上來的殘差也被歸一化了，但這是不好的。一個純淨的殘差是更好的，因為在反向傳播的時候，當梯度回流時，加法是把其梯度均勻分給它的兩個分支，這意味著梯度通過殘差這條路徑直接流向輸入，從優化的角度來看，乾淨的殘差是可取的。

說實話我沒看懂他的解釋，因此去網上搜索了一下相關的內容，然後發現GPT2的做法叫pre-norm，而Attention is all you need的做法叫post-norm。

pre-norm

蘇神對於這兩種差異的解釋非常到位，殘差連接是 $x+F(x)$ , 如果 $x$ 的方差為 $\sigma^2_1$ ， $F(x)$ 的方差為 $\sigma^2_2$ ，那麼殘差連接後的方差就是 $\sigma^2_1+\sigma^2_2$ ，也就是說殘差會放大方差，我們要想辦法縮小這個方差，樸素的方法是添加歸一化，也就是 $x_{t+1}=Norm(x_t+F(x))$ ，然而，這種做法雖然穩定了前向傳播的方差，但事實上已經嚴重削弱了殘差的恆等分支，所以反而失去了殘差 “易於訓練” 的優點，通常要 warmup 並設置足夠小的學習率才能使它收斂。而本身 transformer 的結構有兩個特點： warm-up 階段超參數敏感和優化過程收斂速度慢。(這塊筆者也不知道為什麼)，相當於在 post-norm 的情況下就更難收斂了，並且訓練成本也會有一定程度的上升。

然後來解釋一下削弱了殘差的恆等分支 (其實也就是 Karpathy 說的乾淨的殘差)，假設初始狀態下 $x，F(x)$ 方差都為 1， $x+F(x)$ 方差為 2，而 Normalization 操作負責將方差重新降為 1，這就說明初始階段 Post Norm 相當

x_{t+1}=\frac{x_t+F(x_t)}{\sqrt{2}}

遞歸下去

x_l=\frac{x_{l-1}}{\sqrt{2}}+\frac{F_{l-1}(x_{l-1})}{\sqrt{2}}=\frac{x_{l-2}}{{2}}+\frac{F_{l-2}(x_{l-2})}{{2}}+\frac{F_{l-1}(x_{l-1})}{\sqrt{2}}

x_l=\frac{x_{0}}{{2^{l/2}}}+\frac{F_{0}(x_{0})}{{2^{l/2}}}+\frac{F_{1}(x_{1})}{{2}^{(l-1)/2}}+\frac{F_{2}(x_{2})}{{2}^{(l-2)/2}}+...+\frac{F_{l-1}(x_{l-1})}{{2}^{1/2}}

本來殘差的意思是給前面輸入的層搞一條“綠色通道”，讓梯度可以更直接地回傳，但是在Post Norm中，這條“綠色通道”被嚴重削弱了，越靠近前面的通道反而權重越小，使得在多次殘差連接後，前面的殘差無法感知到末尾的梯度變化，殘差“名存實亡”，因此還是不容易訓練。

論文請見《ON LAYER NORMALIZATION IN THE TRANSFORMER ARCHITECTURE》

而修正後的 Pre-Norm 其形式為

x_{t+1}=x_t+F_t(Norm(x_t))

展開迭代後：

x_{t+1}=x_t+F_t(Norm(x_t)) = x_{t-1}+F_{t-1}(Norm(x_{t-1}))+F_t(Norm(x_t))

x_{t}=x_{0}+F_{0}(Norm(x_{0}))+F_1(Norm(x_1))+...+F_{l-1}(Norm(x_{l-1}))

每一條殘差通道都是平權的，殘差的作用會比 Post Norm 更加明顯，所以它也更好優化。當然，這樣最後的 $x_l$ 方差將會很大，所以在接預測層之前 $x_l$ 也還要加個 Normalization，這正是ln_f

Karpathy 講到 Attention 是 tokens 進行通信的地方，是一個池化函數，是一個加權和函數，是一個 ruduce operation 。
MLP 發生在每個單獨的 token 上，在 tokens 之間沒有信息被收集或者交換，是一個 map operation 。

MLP#

class mlp(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, config.n_embd*4)
        self.Gelu = nn.GELU(approximate='tanh')
        self.c_proj = nn.Linear(config.n_embd*4, config.n_embd)
    
    def forward(self, x):
        x = self.c_fc(x)
        x = self.Gelu(x)
        x = self.c_proj(x)
        return x

很簡單的 MLP 線性映射，從 [n_embd, 4 * n_embd] 到 [4 * n_embd, n_embd] ，中間再加一個非線性層 GELU 激活函數。GELU 的函數圖像和 Relu 很像，但是在尾部處可導，也就是解決了 Relu 在 x 小於 0 時候導數恆為 0 的問題，而這種平滑會產生更好的效果

Karpathy 這裡講到為什麼使用 tanh 的近似值，談到這是一個歷史遺留問題，在 tensorflow 時期用精確的 GELU 特別慢，所以開發了利用 tanh 近似 GELU 的函數。

GELU

Attention#

class CasualSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.c_attn = nn.Linear(config.n_embd, config.n_embd*3)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.n_embd = config.n_embd
        self.n_head = config.n_head
        # 後續在權重模塊會講到它的作用
        self.c_proj.NANOGPT_SCALE_INIT = 1 
        self.register_buffer("bias", torch.tril((torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size)))
    
    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        # 得到qkv
        q, k, v = qkv.split(self.n_embd, dim=2)
        # query, key, value 全部被分成[B, n_head, T, n_embd//n_head]
        query = q.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
        key = k.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
        value = v.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)
        # QK^T/d
        att = query @ key.transpose(-1, -2) * (1.0/math.sqrt(key.size(-1)))
        mask_att = att.masked_fill(self.bias[:,:,:T,:T]==0, float('-inf'))
        wei = F.softmax(mask_att, dim=-1)
        out = wei @ value
        out = out.transpose(1,2).contiguous().view(B, T, C)
        out = self.c_proj(out)
        return out

self.c_attn： $W_q,W_k,W_v$ 的組合，將輸入 $x$ 變為輸入的 $Q,K,V$

self.c_proj：計算完 $\frac{QK^T}{\sqrt{d_k}}V$ 後的一層線性層

self.n_embd：每個 token 的特徵向量空間

self.n_head：多頭注意力機制的頭的個數

self.bias：這裡的 bias 是掩碼的意思，也就是上三角矩陣，防止前面的 token 學習到後面 token 的方法。具體原理如下：對於輸入 $x$ :

x=\begin{bmatrix} x_{11} & x_{12} & x_{13} \\ x_{21} & x_{22} & x_{23} \\ x_{31} & x_{32} & x_{33} \end{bmatrix},bias=\begin{bmatrix} 1 & 0 & 0 \\ 1 & 1 & 0 \\ 1 & 1 & 1 \end{bmatrix} ,mask_{att}=\begin{bmatrix} x_{11} & -inf & -inf \\ x_{21} & x_{22} & -inf \\ x_{31} & x_{32} & x_{33} \end{bmatrix}

-inf 在後續的softmax 過程中會變成接近於 0 的一個值，從而對分類不起影響

contiguous()：transpose 並不改變物理排序，只會改變形式排序，而使用該函數才可以修正物理排序，舉個例子說明，對於數組 $[[[1,2][7,8]],[[3,4][5,6]]]\space \space \space \space \space shape=[1,2,2]$ ， transpose(1,2)後等於 $[[[1,2][2,3]],[[3,5][4,6]]]\space \space \space \space \space shape=[1,2,2]$ ，但是這兩個數組的物理存儲都為 $[1,2,7,8,3,4,5,6]$ ，因此在對 transpose 後的數組進行 view 操作時會報錯。

DownLoad from HuugingFace#

    # 在 class GPT 下
    @classmethod
    def from_pretrained(cls, model_type):
        """Loads pretrained GPT-2 model weights from huggingface"""
        # 四種類型的model
        assert model_type in {'gpt2','gpt2-medium','gpt2-large','gpt2-xl'}
        # 打印出你加載的是哪一種
        print("Loading weights from pretrained gpt:%s"%model_type)
        # 每種GPT所對應的超參數不一樣大
        config_args ={
        'gpt2' : dict(n_layer=12,n_head=12,n_embd=768), # 124M params
        'gpt2-medium' : dict(n_layer=24,n_head=16,n_embd=1024), #350M params
        'gpt2-large' : dict(n_layer=36,n_head=20,n_embd=1280), #774M params
        'gpt2-xl' : dict(n_layer=48,n_head=25,n_embd=1600), #1558M params
        }[model_type]
        # 詞表大小永遠是50527
        config_args['vocab_size'] = 50257 
        # 單個block的大小始終是1024
        config_args['block_size'] = 1024
        # 給模型導入超參數
        config = GPTConfig(**config_args)
        model = GPT(config)
        # sd是模型參數名字典
        sd = model.state_dict()
        sd_keys = sd.keys()
        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask

        # 從HF上下載權重，sd_hf是下載下來的模型參數名字典
        model_hf = GPT2LMHeadModel.from_pretrained(model_type, cache_dir="/home/shong_Tan/project/gpt_2/model_weight", local_files_only=True)
        sd_hf = model_hf.state_dict()

        sd_keys_hf = sd_hf.keys()
        # 丟掉HF權重中掩碼bias
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')]
        # 丟掉HF權重中bias
        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')]
        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
        # 確保sd和hf_sd的參數名一樣多
        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
        # 確保sd和hf_sd的transformer塊的權重名一樣
        for k in sd_keys_hf:
            if any(k.endswith(w) for w in transposed):
                assert sd_hf[k].shape[::-1] == sd[k].shape
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k].t())
            else:
                assert sd_hf[k].shape == sd[k].shape, f"mismatched keys: {sd_hf[k].shape} != {sd[k].shape}"
                with torch.no_grad():
                    sd[k].copy_(sd_hf[k])
        return model

看代碼註釋就好了

Tips⚠️：從 HF 下載下來的lm_head.weight 和 transformer.wte.weight 的形狀大小是相同的，都是 $[50527, 768]$ ，它們一個是輸入嵌入，一個是輸出 logits，這兩個應該是一致的，這才能反應出當 token 被嵌入特徵向量時的語義在經過交互後，當輸出的時候還是這個特徵向量，它又可以被轉回到原來的 token。同時， $50527*768 \approx 40M$ ，可以節省大量的顯存。

Forward#

    # 在 class GPT 下
    def forward(self, idx, target):
		    # 進入時維度為[batch, token長度]
        B, T = idx.size()
        # token長度不能超過上下文
        assert T <= self.config.block_size, f"超出輸入上下文長度限制 {T-self.config.block_size} token"
        # pos [0,1,2,..,T-1]，並且記住要放到device上去
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        # 位置嵌入
        pos = self.transformer.wpe(pos) #(T, n_embd)
        # token嵌入
        tok = self.transformer.wte(idx) #(B, T, n_embd)
        # 相加是在 (T, n_embd)這個維度上的數值相加
        x = tok + pos
        # 經過transformer塊
        for block in self.transformer.h:
            x = block(x)
        # 最後一層歸一化
        x = self.transformer.ln_f(x)
        # 線性層輸出
        logits = self.lm_head(x) #(B, T, vocab_size)
        loss = None
        # 有target的話也就是有標籤就進行訓練，計算損失函數，反之只用推理就可以
        if target is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), target.view(-1))
        return logits, loss

# 小試牛刀
num_return_sequences = 5
max_length = 30
model = GPT.from_pretrained('gpt2')
# eval會在評估的時候禁用dropout層，對於batchnorm也會有不一樣的反應，以及凍結參數
model.eval()
# 模型移到gpu
model.to('cuda')

# 以下是分詞，調openai的tiktoken庫就行了，如果想要知道原理的，建議去看文章開頭的博主blog
import tiktoken
enc = tiktoken.get_encoding('gpt2')
tokens = enc.encode("Hello, I'm a language model, ")
tokens = torch.tensor(tokens, dtype=torch.long) # [8, ]
tokens = tokens.unsqueeze(0).repeat(num_return_sequences, 1) # [5,8]
x = tokens.to('cuda')

while x.size(1) < max_length:
		with torch.no_grad():
				# 輸入模型得到結果
				logits, loss  = model(x) # x: [B, T]    logits:[B,T,C]
				# 取出最後一個token的預測
				logits = logits[:, -1, :] # [B, 1, C]
				# 對最後一個維度C取softmax
				probs = F.softmax(logits, dim=-1) # [B, 1, C]
				# 在最後一個維度C選出topk大的概率和對應的index
				topk_probs, topk_indices = torch.topk(probs, 50, dim=-1) 
				# 從topk的概率中隨機選取一個概率
				ix = torch.multinomial(topk_probs, 1)
				# 找到被選中的概率對應的index
				xcol = torch.gather(topk_indices, -1, ix)
				# 將得到的輸出token加在x上又作為輸入 [B, T+1]
				x = torch.cat((x, xcol), dim=1)

分詞形式：將 "Hello, I'm a language model," 變為 [15496, 11, 314, 1101, 257, 3303, 2746, 11, 220]

參考下面網站可以自己嘗試

https://tiktokenizer.vercel.app/

初始化#

Dataset#

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'
# 這是mac的M芯片系列
elif hasattr(torch.backends, "eps") and torch.backends.mps.is_available():
    device = 'mps'
print("Using device: ", device)
import tiktoken
enc = tiktoken.get_encoding('gpt2')
with open('input.txt', 'r') as f:
    data = f.read()
text = data[:1000]
tokens = enc.encode(text)
B, T = 4, 32
buf = torch.tensor(tokens[:B*T+1])
buf.to(device)
# 本質上是通過前n個詞推測n+1個詞
x = buf[:-1].view(B, T)
y = buf[1:].view(B,T)

model.GPT(GPTConfig())
model.to(device)
logits, loss = model(x)
print(loss.item())

這裡loss的值大概是 11 左右，因為 $-log(\frac{1}{50257})\approx 11$

訓練單個 batch 代碼

# 使用AdamW優化器，Adam與SGD的區別自己去了解吧 
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(50):
		# 優化器梯度清零
		optimizer.zero_grad()
		# 得到logits和損失
		logits, loss = model(x, y)
		# 反向傳播求導
		loss.backward()
		# 利用導數更新原有參數
		optimizer.step()

Adam 優化器能比 SGD 更快地收斂

Dataloader 函數

class DataLoaderLite():

    def __init__(self, B, T):
        self.B = B
        self.T = T
        # 讀入整個input.txt
        with open('input.txt', 'r') as f:
            data = f.read()
        enc = tiktoken.get_encoding('gpt2')
        tokens = enc.encode(data)
        self.tokens = torch.tensor(tokens, dtype=torch.long)
        print(f"load {len(self.tokens)} tokens")
        print(f"1 epoch = {len(self.tokens)//(B*T)} batched")
        # 定義在當前batch的位置
        self.current_position = 0

    def next_batch(self):
        B, T = self.B, self.T
        buf = self.tokens[self.current_position: self.current_position+B*T+1]
        x = buf[:-1].view(B, T)
        y = buf[1:].view(B, T)
        # 每個batch有B*T個元素對
        self.current_position += B*T
        # 如果batch把tokens用完了，又回到toknes[0]
        if self.current_position+B*T+1 > len(self.tokens):
            self.current_position = 0
        return x, y

修正訓練代碼

train_loader = DataLoaderLite(4, 32 )
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for i in range(50):
		optimizer.zero_grad()
    x, y = train_loader.next_batch()
    x, y = x.to(device), y.to(device)
    logits, loss = model(x, y)
    loss.backward()
		optimizer.step()
		print(f"step: {i}, loss: {loss.item()}")

權重#

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, ' '):
                std = std * (2*self.config.n_layer**-0.5)
            torch.nn.init.normal_(module.weight, mean=0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)

        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0, std=0.02)

std = std * (2*self.config.n_layer)**-0.5 ：此處的方差來源是考慮了殘差流的貢獻，每次殘差連接都說明輸入 input 有了一次等額貢獻，需要一個因子來處理， $\frac{1}{\sqrt{2*n_{layer}}}$ ，這裡是控制了 Pre-Norm 的殘差連接導致的方差過大，這裡的 2 是因為每層裡面 Attention 和 MLP 都用到了一次殘差。

std ：std 的值的來源也是有根據的，按照 GPT2 裡文檔的說法，最好在 $\frac{1}{\sqrt{n_{embd}}}$ 左右