第4课 手写 Model.py 大模型代码逻辑

内容总结与提炼

• 核心主题与目标：本视频的核心目标是带领观众从零开始，手写一个功能完整、结构清晰的大语言模型 Model.py 文件。与系列课第一阶段的步骤化模拟不同，本次旨在创建一个可封装、可复用的类结构，将训练、推理、微调等流程分离开，实现一个类似 Hugging Face 但更具可控性的 Transformer 模型。最终目标是产出一个既可用于学习，也可用于生产环境的模型代码。

• 关键知识点梳理：

  1. 环境与框架：使用 PyTorch 框架。
  2. 模块化设计：将 Transformer 模型拆解为多个独立的 nn.Module 子类，包括：
     • FeedForward: 前馈网络层。
     • Attention: 单头注意力机制。
     • MultiHeadAttention: 多头注意力机制。
     • TransformerBlock: 组合了多头注意力和前馈网络的完整 Transformer 模块。
     • Model: 整合所有模块，并处理输入、输出和损失计算的顶层模型。
  3. 核心概念实现：
     • 超参数（Hyperparameters）：定义如 d_model, context_len, n_head 等关键参数。
     • 注意力机制（Attention）：实现 QKV 的生成、Scaled Dot-Product Attention、Causal Masking（因果掩码）和 Softmax。
     • 位置编码（Positional Encoding）：使用 sin 和 cos 函数动态生成位置信息矩阵。
     • 残差连接（Residual Connection） 与 层归一化（Layer Normalization）：在 TransformerBlock 中正确应用。
     • 损失计算（Loss Calculation）：使用 F.cross_entropy 计算模型预测与真实标签之间的损失。
  4. PyTorch 技巧：
     • nn.Sequential：用于串联多个网络层。
     • nn.ModuleList：用于存储和管理多个子模块（如多个注意力头）。
     • register_buffer：用于存储非模型参数但需要随模型移动（如到 GPU）的状态，如此处的 mask。
     • 张量变形（view, reshape）：为满足特定函数（如 cross_entropy）的输入要求而调整张量形状。

• 逻辑结构与关系：视频采用**自底向上（Bottom-Up）**的构建逻辑。

  1. 基础设置：首先定义全局超参数。
  2. 构建基础组件：依次实现最小的功能单元，如 FeedForward 和单头的 Attention。
  3. 组合组件：利用已实现的基础组件构建更复杂的模块，如用 Attention 构建 MultiHeadAttention。
  4. 构建核心模块：将 MultiHeadAttention 和 FeedForward 组合成一个完整的 TransformerBlock。
  5. 整合为最终模型：将多个 TransformerBlock 堆叠，并添加词嵌入、位置编码和最终的输出层，形成最终的 Model 类。 这种结构清晰地展示了 Transformer 模型各部分的功能及其相互关系，便于理解和修改。

目录

• 引言与目标
• 超参数定义
• 构建前馈网络 (FeedForward)
• 构建单头注意力机制 (Attention)
• 构建多头注意力机制 (MultiHeadAttention)
• 构建 Transformer 块 (TransformerBlock)
• 构建最终模型 (Model)
• 代码修正与补充
• AI 总结

---

引言与目标

本节课的目标是手写一个可封装、可复用的 Model.py 文件，它将包含 Transformer 的完整架构实现。这与之前系列课中为了直观理解而分步打印矩阵变化的方式不同，这次的代码将是模块化的、生产级别的。

• 优势：
  • 可控性强：可以方便地修改模型架构，例如增减 FeedForward 层的数量、更换激活函数、调整位置编码方式（如从 sin/cos 更换为 RoPE）。
  • 结构清晰：将模型、训练、推理、微调等功能分离，便于维护和扩展。
• 代码结构：
  • 本次课程主要完成 Model.py，包含 Transformer 的所有核心类。
  • 后续课程将编写用于推理和训练的脚本来调用这个模型。
• PyTorch 讲解：课程中会即时讲解用到的 PyTorch 知识，即使没有经验也能跟上。

超参数定义

在模型实现开始前，首先定义一些关键的超参数（Hyperparameters）。初期将它们直接写在 model.py 文件中，后期可以移至专门的配置文件。

• d_model = 512: 模型的维度，即词嵌入向量的长度。
• context_len = 16: 上下文长度，即模型一次能处理的 token 数量。
• n_head = 8: 注意力头的数量。
• head_size = d_model // n_head: 每个注意力头的维度 (512 / 8 = 64)。
• dropout = 0.1: Dropout 的比率，以 10% 的概率随机丢弃神经元，防止过拟合。
• n_blocks = 12: Transformer Block 的层数（即 Nx 的值）。
• batch_size = 4: 训练时每个批次的样本数。
• device: 动态设置计算设备。如果 CUDA (GPU) 可用，则使用 cuda，否则使用 cpu。

  device = 'cuda' if torch.cuda.is_available() else 'cpu'

构建前馈网络 (FeedForward)

前馈网络（Feed-Forward Network）是 Transformer Block 中的一个重要组成部分。它包含两个线性层和一个非线性激活函数。

• 定义 FeedForward 类：

  • 继承自 torch.nn.Module。
  • 在构造函数 __init__ 中，使用 nn.Sequential 将多个操作串联起来，形成一个网络流。

• 网络结构：

  1. nn.Linear(d_model, 4 * d_model): 第一个线性层，将输入维度从 d_model 扩展到 4 * d_model。这是 Transformer 论文中的标准做法。
  2. nn.ReLU(): 使用 ReLU作为激活函数，引入非线性。
  3. nn.Linear(4 * d_model, d_model): 第二个线性层，将维度从 4 * d_model 压缩回 d_model。
  4. nn.Dropout(dropout): 在最后应用 Dropout。

• forward 方法：

  • 定义了数据通过该模块时的前向传播路径。
  • 直接调用在 __init__ 中定义好的 nn.Sequential 对象即可。

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

构建单头注意力机制 (Attention)

这是实现自注意力机制的核心部分，首先从单个头的逻辑开始。

• 定义 Attention 类：

  • 同样继承自 nn.Module。
  • 与理论不同点：在实际代码实现中，每个注意力头都拥有自己独立的 W_q, W_k, W_v 权重矩阵，而不是共享。

• __init__ 构造函数：

  1. 定义权重矩阵：
     • self.W_q = nn.Linear(d_model, head_size, bias=False): Query 的线性变换层。
     • self.W_k = nn.Linear(d_model, head_size, bias=False): Key 的线性变换层。
     • self.W_v = nn.Linear(d_model, head_size, bias=False): Value 的线性变换层。
     • 注意：此处将 bias 设置为 False，因为在多头注意力的输出层通常不需要偏置项。
  2. 定义因果掩码 (Causal Mask)：
     • 使用 self.register_buffer('mask', ...) 来定义掩码。
     • register_buffer 的作用：它将一个张量注册为模型的缓冲区。这个张量不是模型的参数（即在反向传播中不会被更新），但它会成为模型状态的一部分，意味着它会随着模型一起被移动到 GPU (.to(device))，并且在保存和加载模型状态时会被一并处理。
     • 为何使用 buffer：因为在训练过程中，输入的序列长度是动态变化的（从 1 到 context_len），所以掩码的尺寸也需要相应地裁剪。buffer 适合存储这种需要根据输入动态调整但本身不参与训练的变量。
     • 生成掩码：通过 torch.tril(torch.ones(context_len, context_len)) 创建一个下三角矩阵，上三角部分为 0，下三角部分（包括对角线）为 1。
  3. 定义 Dropout：self.dropout = nn.Dropout(dropout)，用于对注意力权重进行随机丢弃。

• forward 方法：

  1. 获取 Q, K, V：
     • q = self.W_q(x)
     • k = self.W_k(x)
     • v = self.W_v(x)
     • 输入 x 的维度是 (B, T, C) 即 (batch, time_step, d_model)。
     • 输出的 Q, K, V 维度是 (B, T, H) 即 (batch, time_step, head_size)。
  2. 计算注意力分数 (Attention Scores)：
     • scores = q @ k.transpose(-2, -1): Q 和 K 的转置进行矩阵乘法。
     • 维度变化：(B, T, H) @ (B, H, T) -> (B, T, T)。
  3. 缩放 (Scaling)：
     • scores = scores / (head_size**0.5): 除以 $\sqrt{d_k}$ (即 head_size 的平方根)，以防止梯度消失或爆炸。
  4. 应用掩码 (Masking)：
     • scores = scores.masked_fill(self.mask[:T, :T] == 0, float('-inf')): 将掩码中值为 0 的位置（上三角）替换为负无穷大。
     • [:T, :T] 是为了处理动态的序列长度 T。
  5. Softmax：
     • scores = F.softmax(scores, dim=-1): 在最后一个维度上进行 Softmax，得到归一化的注意力权重。
  6. 应用 Dropout：scores = self.dropout(scores)。
  7. 与 V 相乘：
     • output = scores @ v: 将注意力权重应用到 V 上。
     • 维度变化：(B, T, T) @ (B, T, H) -> (B, T, H)。
  8. return output: 返回单头注意力的输出。

构建多头注意力机制 (MultiHeadAttention)

将多个单头注意力机制并行运行，并将它们的结果拼接起来。

• 定义 MultiHeadAttention 类：
  1. __init__ 构造函数：
     • self.heads = nn.ModuleList([Attention() for _ in range(n_head)]): 使用列表推导式和 nn.ModuleList 创建 n_head 个独立的 Attention 实例。nn.ModuleList 确保这些子模块能被 PyTorch 正确识别和管理。
     • self.W_o = nn.Linear(d_model, d_model): 定义输出线性层，用于将拼接后的多头结果投影回 d_model 维度。
     • self.dropout = nn.Dropout(dropout): 输出层的 Dropout。
• forward 方法：
  1. 并行计算：out = [h(x) for h in self.heads]，通过列表推导式让每个头独立处理输入 x。
  2. 拼接 (Concatenate)：out = torch.cat(out, dim=-1)，在最后一个维度（特征维度）上将所有头的输出拼接起来。
     • 维度变化：8 个 (B, T, 64) 的张量拼接成一个 (B, T, 512) 的张量。
  3. 最终投影：out = self.dropout(self.W_o(out))，将拼接后的结果通过输出线性层和 Dropout。
  4. return out。

构建 Transformer 块 (TransformerBlock)

一个完整的 Transformer Block 包含一个多头注意力模块和一个前馈网络模块，每个模块前后都有残差连接和层归一化。

• 定义 TransformerBlock 类：
  • __init__ 构造函数：
    • self.multi_head_attention = MultiHeadAttention(): 实例化多头注意力模块。
    • self.feed_forward = FeedForward(): 实例化前馈网络模块。
    • self.ln1 = nn.LayerNorm(d_model): 第一个层归一化。
    • self.ln2 = nn.LayerNorm(d_model): 第二个层归一化。
• forward 方法：
  1. 第一个子层（多头注意力）：
     • x = x + self.multi_head_attention(self.ln1(x)):
       • 首先对输入 x 进行层归一化 (self.ln1(x))。
       • 将结果送入多头注意力模块。
       • 将注意力模块的输出与原始输入 x 相加（残差连接）。
  2. 第二个子层（前馈网络）：
     • x = x + self.feed_forward(self.ln2(x)):
       • 对上一步的输出 x 进行第二次层归一化 (self.ln2(x))。
       • 将结果送入前馈网络。
       • 将前馈网络的输出与该子层的输入相加（第二次残差连接）。
  3. return x。

构建最终模型 (Model)

这是顶层类，它将所有组件整合在一起，形成一个完整的语言模型。

• 定义 Model 类：
  • __init__ 构造函数：
    • self.token_embedding_table = nn.Embedding(vocab_size, d_model): 词嵌入表，将 token 索引映射为 d_model 维的向量。vocab_size 是词汇表大小。
    • self.blocks = nn.Sequential(*[TransformerBlock() for _ in range(n_blocks)]): 使用 nn.Sequential 和列表解包 * 将 n_blocks 个 TransformerBlock 串联起来。
    • self.ln_f = nn.LayerNorm(d_model): 在所有 TransformerBlock 之后应用的最终层归一化。
    • self.lm_head = nn.Linear(d_model, vocab_size): 语言模型头（输出层），一个线性层，将 d_model 维的向量投影到词汇表大小的维度，得到每个 token 的 logits。
• 位置编码（Positional Encoding）：
  • 这部分在 forward 方法中动态生成，以适应不同的序列长度。
  • 公式原理：
    • $P{E}_{(pos,2i)}=\sin (pos/10000^{2i/d_{model}})$
    • $P{E}_{(pos,2i+1)}=\cos (pos/10000^{2i/d_{model}})$
  • 代码实现：
    1. 创建一个形状为 (context_len, d_model) 的全零张量 pe。
    2. 生成 position 张量（0 到 context_len-1）。
    3. 计算分母项 div_term，即 $10000^{2i/d_{model}}$。
    4. 分别计算 sin 和 cos 值，并交错填充到 pe 的偶数和奇数列。
• forward 方法：
  • forward(self, x_batch, y_batch=None): 接收训练输入 x_batch 和目标 y_batch（在推理时为 None）。
  1. 获取输入维度：B, T = x_batch.shape。
  2. 词嵌入：tok_emb = self.token_embedding_table(x_batch)。
  3. 位置编码：pos_emb = self.position_encoding_table[:T, :]，获取对应长度的位置编码。
  4. 相加：x = tok_emb + pos_emb，将词嵌入和位置编码相加。
  5. 通过 Transformer Blocks：x = self.blocks(x)。
  6. 最终层归一化：x = self.ln_f(x)。
  7. 输出投影：logits = self.lm_head(x)，得到预测的 logits。
  8. 计算损失 (Loss)：
     • 仅在 y_batch 不为 None 时（即训练模式）执行。
     • 维度调整：F.cross_entropy 要求 logits 是二维 (N, C)，target 是一维 (N)。因此需要将 logits 和 y_batch 从 (B, T, C) 和 (B, T) 分别 view 或 reshape 为 (B*T, C) 和 (B*T)。
     • loss = F.cross_entropy(logits.view(B*T, vocab_size), y_batch.view(B*T))。
  9. 返回结果：返回 logits 和 loss。
• generate 方法 (推理)：
  • 这是用于生成新 token 的方法，在视频末尾被提及但未详细实现。它会循环调用 forward 方法来逐个生成 token。

代码修正与补充

视频末尾对代码进行了一些细节修正和补充。

1. bias=False：在 Attention 类的 W_q, W_k, W_v 线性层中，明确设置 bias=False。这是一种常见的实践，因为偏置项在后续的层归一化中可能会被抵消，去掉它可以减少不必要的参数。
2. torch.cat 的 dim 参数：在 MultiHeadAttention 中，torch.cat(out, dim=-1) 明确了在最后一个维度上进行拼接。
3. F.softmax 的 dim 参数：在 Attention 类中，F.softmax(scores, dim=-1) 明确了在最后一个维度（即对每个 query 的 key 权重）上进行 softmax。
4. view 方法的使用：修正了 cross_entropy 前的维度变换，应在张量对象上调用 .view() 方法，例如 logits.view(...) 和 y_batch.view(...)。