大模型入门 ch04：实现一个GPT模型

本文是github上的大模型教程LLMs-from-scratch的学习笔记，教程地址：教程链接

LLM大模型主要是参数量大，而不是代码量大。

这是本节的具体内容

首先实现一个GPT的骨架
分别实现GPT骨架内的各个部分，包括LayerNorm，GELU,Feed forward network和Shorcut connections
将四个部分合成Transformer块
重复多次Transformer块，实现最终的GPT架构

1. GPT骨架

首先确定模型的配置，包括词汇表数量，上下文长度，embeddding维度等等。

GPT_CONFIG_124M = {
    "vocab_size": 50257,    # Vocabulary size
    "context_length": 1024, # Context length
    "emb_dim": 768,         # Embedding dimension
    "n_heads": 12,          # Number of attention heads
    "n_layers": 12,         # Number of layers
    "drop_rate": 0.1,       # Dropout rate
    "qkv_bias": False       # Query-Key-Value bias
}

"vocab_size"表示词汇量为50257个单词，由第2章讨论的BPE分词器支持
"context_length"表示模型的最大输入标记数，这是由第2章中介绍的位置嵌入启用的
"emb_dim"是token输入的嵌入大小，将每个输入token转换为768维向量
"n_heads"是第3章实现的多头注意力机制中注意力头的数量
"n_layers"是模型中transformer块的数量，我们将在接下来的部分中实现
"drop_rate"是第3章中讨论的dropout机制的强度;0.1意味着在训练过程中丢弃10%的隐藏单元以缓解过拟合
“qkv_bias”决定多头注意力机制(来自第3章)中的“线性”层在计算查询(Q)、键(K)和值(V)张量时是否应该包含一个偏差向量。

首先我们搭好这个骨架，后续进行补充完善：

import torch
import torch.nn as nn


class DummyGPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        
        # Use a placeholder for TransformerBlock
        self.trf_blocks = nn.Sequential(
            *[DummyTransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        
        # Use a placeholder for LayerNorm
        self.final_norm = DummyLayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits


class DummyTransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        # A simple placeholder

    def forward(self, x):
        # This block does nothing and just returns its input.
        return x


class DummyLayerNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-5):
        super().__init__()
        # The parameters here are just to mimic the LayerNorm interface.

    def forward(self, x):
        # This layer does nothing and just returns its input.
        return x

2. Layer Normalization

Layer normalization将神经网络的激活值归一化到均值为0，方差为1，使得训练更加稳定和快速。在多头注意力机制的前后，我们都会加入Layernorm层。

归一化可以对不同的维度进行归一化

例如在我们这个例子中，如果对行归一化，那么就是对每一个样本进行归一化，如果是列归一化，就是对每一个特征进行归一化。通过减去平均值再除以均方根，我们可以将输入归一化成均值为0，方差为1。

class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

除了通过减去均值并除以方差来执行规范化之外，我们还添加了两个可训练参数，scale和shift参数
初始值scale(乘以1)和shift(加0)没有任何影响;然而，scale和shift是可训练参数，如果确定这样做会提高模型在训练任务上的性能，LLM会在训练过程中自动调整
这允许模型学习最适合其处理数据的适当缩放和平移
在计算方差的平方根之前，我们还添加了一个较小的值(eps);这是为了避免方差为0时除零误差

3. 前馈网络

在本节中，我们实现了一个小的神经网络子模块，该子模块用作LLMs中transformer模块的一部分
在深度学习中，ReLU(整流线性单元)激活函数由于其简单性和在各种神经网络架构中的有效性而被广泛使用
在llm中，除了传统的ReLU之外，还使用了各种其他类型的激活函数;两个值得注意的例子是GELU(高斯误差线性单元)和SwiGLU (Swish-Gated线性单元)。
GELU和SwiGLU是更复杂的光滑激活函数，分别包含高斯门控线性单元和sigmoid门控线性单元，为深度学习模型提供了更好的性能，不像ReLU的更简单的分段线性函数层-激活函数 (feed forward network with GELU activations)
GELU 可以通过多种方式实现。准确的版本被定义为GELU(x)=x⋅Φ(x)，其中Φ(x)是标准高斯分布的累积分布函数。
-在实践中，通常实现一个计算成本更低的近似: $\text{GELU}(x) \approx 0.5 \cdot x \cdot \left(1 + \tanh\left[\sqrt{\frac{2}{\pi}} \cdot \left(x + 0.044715 \cdot x^3\right)\right]\right)$
(原始的GPT-2模型也是用这个近似进行训练的)

class GELU(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return 0.5 * x * (1 + torch.tanh(
            torch.sqrt(torch.tensor(2.0 / torch.pi)) * 
            (x + 0.044715 * torch.pow(x, 3))
        ))

在这里插入图片描述
于是我们可以组成一个FeedFroward前馈网络层：

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
            GELU(),
            nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
        )

    def forward(self, x):
        return self.layers(x)

4. 残差连接

最初，在计算机视觉深度网络(残差网络)中提出了残差连接，以缓解梯度消失问题
连接为渐变在网络中的流动创建了一条更短的路径
这是通过将一个层的输出添加到后面一层的输出来实现的，通常在中间跳过一个或多个层

5. 连接注意力和全连接层，构成Transform block

在本节中，我们现在将前面的概念结合到的transformer块中
transformer模块结合了上一章中的因果多头注意力模块和线性层，即我们在上一节中实现的前馈神经网络
此外，transformer block也使用dropout和shortcut连接

from previous_chapters import MultiHeadAttention


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"], 
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        # Shortcut connection for attention block
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        # Shortcut connection for feed forward block
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # Add the original input back

        return x

6. 编写GPT模型

transformer块重复多次;在最小的124M GPT-2模型中，我们重复12次。

我们让cfg["n_layers"] = 12

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        
        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

7. 生成文本

GPT模型一次只生成一个token

下面的generate_text_simple函数实现了贪婪解码，这是一种简单快速的生成文本的方法
在贪婪解码中，在每一步，模型选择具有最高概率的单词(或标记)作为它的下一个输出(最高的logit对应最高的概率，因此我们在技术上甚至不需要显式地计算softmax函数)
在下一章中，我们将实现一个更高级的generate_text函数
下图描述了GPT模型在给定输入上下文的情况下如何生成下一个单词标记

def generate_text_simple(model, idx, max_new_tokens, context_size):
    # idx is (batch, n_tokens) array of indices in the current context
    for _ in range(max_new_tokens):
        
        # Crop current context if it exceeds the supported context size
        # E.g., if LLM supports only 5 tokens, and the context size is 10
        # then only the last 5 tokens are used as context
        idx_cond = idx[:, -context_size:]
        
        # Get the predictions
        with torch.no_grad():
            logits = model(idx_cond)
        
        # Focus only on the last time step
        # (batch, n_tokens, vocab_size) becomes (batch, vocab_size)
        logits = logits[:, -1, :]  

        # Apply softmax to get probabilities
        probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)

        # Get the idx of the vocab entry with the highest probability value
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)

        # Append sampled index to the running sequence
        idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

    return idx