Coding an LLM architecture 编码 LLM 架构#

使用如下配置

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   Dropout 比例
    "qkv_bias": False        # Query-Key-Value bias
}

一个抽象架构如下

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"])
        #1 TransformerBlock 占位符
        self.trf_blocks = nn.Sequential(               
            *[DummyTransformerBlock(cfg)               
              for _ in range(cfg["n_layers"])]         
        )         
        #2 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__()

    def forward(self, x):     
        return x

class DummyLayerNorm(nn.Module):          
    def __init__(self, normalized_shape, eps=1e-5):   
        super().__init__()

    def forward(self, x):
        return x

现在我们就有了一个初始的框架，调用也可以跑通

tokenizer = tiktoken.get_encoding("gpt2")
batch = []
txt1 = "Every effort moves you"
txt2 = "Every day holds a"

batch.append(torch.tensor(tokenizer.encode(txt1)))
batch.append(torch.tensor(tokenizer.encode(txt2)))
batch = torch.stack(batch, dim=0)

torch.manual_seed(123)
model = DummyGPTModel(GPT_CONFIG_124M)
logits = model(batch)
print("Output shape:", logits.shape)
print(logits)

> Output shape: torch.Size([2, 4, 50257])
> tensor([[[-1.2034,  0.3201, -0.7130,  ..., -1.5548, -0.2390, -0.4667],
         [-0.1192,  0.4539, -0.4432,  ...,  0.2392,  1.3469,  1.2430],
         [ 0.5307,  1.6720, -0.4695,  ...,  1.1966,  0.0111,  0.5835],
         [ 0.0139,  1.6754, -0.3388,  ...,  1.1586, -0.0435, -1.0400]],

        [[-1.0908,  0.1798, -0.9484,  ..., -1.6047,  0.2439, -0.4530],
         [-0.7860,  0.5581, -0.0610,  ...,  0.4835, -0.0077,  1.6621],
         [ 0.3567,  1.2698, -0.6398,  ..., -0.0162, -0.1296,  0.3717],
         [-0.2407, -0.7349, -0.5102,  ...,  2.0057, -0.3694,  0.1814]]],
       grad_fn=<UnsafeViewBackward0>)

嵌入有 50,257 个维度，因为这些维度中的每一个都指向词汇表中的一个唯一标记。当我们实现后处理代码时，我们将这些 50,257 维向量转换回标记 ID，然后我们可以将其解码成单词。

Normalizing activations with layer normalization 使用层归一化规范化激活#

NeatReader-1745407843002

层归一化背后的主要思想是调整神经网络层的激活值（输出），使其均值为 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

Implementing a feed forward network with GELU activations 实现具有高斯误差线性单元（GELU）激活函数的前馈网络#

NeatReader-1745408279264

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))
        ))

我们使用 GELU 函数来实现小型神经网络模块 FeedForward

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)

FeedForward 模块是一个小型神经网络，由两个 Linear 层和一个 GELU 激活函数组成。在拥有 1.24 亿参数的 GPT 模型中，它通过 GPT_CONFIG_124M 字典接收输入批次，每个批次中的词元嵌入维度为 768，其中 GPT_CONFIG_ 124M[“emb_dim”] = 768

NeatReader-1745416486781

做到这里就完成了如图的部分

NeatReader-1745416830263

Adding shortcut connections 添加快捷连接#

捷径连接（也称为跳跃连接或残差连接）背后的概念。最初，捷径连接是为计算机视觉中的深度网络（特别是残差网络）提出的，以缓解梯度消失的问题。梯度消失问题是指在训练过程中引导权重更新的梯度在反向传播通过各层时逐渐变小，使得早期层难以有效训练的问题。

NeatReader-1745458823052

class ExampleDeepNeuralNetwork(nn.Module):
    def __init__(self, layer_sizes, use_shortcut):
        super().__init__()
        self.use_shortcut = use_shortcut
        self.layers = nn.ModuleList([       #1
            nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), 
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), 
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), 
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]), 
                          GELU()),
            nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]), 
                          GELU())
        ])

    def forward(self, x):
        for layer in self.layers:
            layer_output = layer(x)         #2
            if self.use_shortcut and x.shape == layer_output.shape:    #3 这里就是判断如果开启了 shortcut
                x = x + layer_output # 将输入和输出累加作为下一次的输入
            else:
                x = layer_output
        return x

Connecting attention and linear layers in a transformer block 在 Transformer 模块中连接注意力层和线性层#

让我们实现 Transformer 模块，它是 GPT 和其他 LLM 架构的基本构建块。在拥有 1.24 亿参数的 GPT - 2 架构中，这个模块会重复十几次，它结合了我们之前介绍过的几个概念：多头注意力机制、层归一化、随机失活、前馈层和高斯误差线性单元（GELU）激活函数。

NeatReader-1745475979670

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 = x #1 使用 shortcut
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut      #2 将结果 connection 起来

        shortcut = x         #3 循环一次
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut      #4
        return x

在这两个组件的每一个之前应用层归一化（ LayerNorm ），并在它们之后应用随机失活（dropout）来对模型进行正则化并防止过拟合。这也被称为前置层归一化（Pre-LayerNorm）。较旧的架构，如原始的 Transformer 模型，反而在自注意力和前馈网络之后应用层归一化，称为后置层归一化（Post-LayerNorm），这通常会导致较差的训练动态。

使用如下代码进行尝试

torch.manual_seed(123)
x = torch.rand(2, 4, 768)  # 1
block = TransformerBlock(GPT_CONFIG_124M)
output = block(x)

print("Input shape:", x.shape)
print("Output shape:", output.shape)

> Input shape: torch.Size([2, 4, 768])
Output shape: torch.Size([2, 4, 768])

到这里，我们就实现了所有部分 NeatReader-1745478540422

Coding the GPT model 编写 GPT 模型#

NeatReader-1745478628870

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
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

生成文本#

NeatReader-1745478818639

GPT 模型从输出张量转换为生成文本的过程涉及几个步骤。

NeatReader-1745478883309

代码如下

def generate_text_simple(model, idx,                 #1  idx 是当前上下文中一个形状为(batch, n_tokens)的索引数组。
                         max_new_tokens, context_size): 
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]    #2 如果当前上下文超过支持的上下文大小，则裁剪上下文，例如，如果 LLM 仅支持 5 个词元，而上下文大小为 10，则仅使用最后 5 个词元作为上下文
        with torch.no_grad():
            logits = model(idx_cond)

        logits = logits[:, -1, :]                    #3
        probas = torch.softmax(logits, dim=-1)           #4
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)    #5
        idx = torch.cat((idx, idx_next), dim=1)     #6 将采样得到的索引追加到正在生成的序列中，其中`idx`的形状为（批次，n_tokens + 1）

    return idx
Backward Coding attention mechanisms 注意力机制 Pretraining on unlabeled data 在未标签的数据进行预训练 Forward
comments powered by Disqus