小狍子皮皮奇

外观

实现GPT模型

约 2738 字大约 9 分钟

设计思路与核心概念

1. GPT模型的背景与动机

GPT（Generative Pre-trained Transformer）是OpenAI开发的生成式预训练Transformer模型，代表了现代大语言模型的重要里程碑，主要解决以下问题：

  • 语言生成任务：通过自回归方式生成连贯、有意义的文本
  • 预训练范式：在大规模无标注文本上进行预训练，然后微调到具体任务
  • 统一架构：使用单一的Transformer架构处理多种NLP任务
  • 可扩展性：通过增加参数量和数据量持续提升模型能力

2. 核心设计思想

GPT模型的核心思想是基于Transformer的自回归语言建模：

  • 自回归生成：根据前面的词元预测下一个词元，实现文本生成
  • 因果注意力：只能关注当前位置之前的信息，保证生成的因果性
  • 位置编码：通过位置嵌入让模型理解序列中的位置信息
  • 词汇表映射：将离散的词元ID转换为连续的向量表示

3. 模型架构

GPT模型的整体结构

词元ID → 词元嵌入 → 位置嵌入 → 相加 → Dropout →
Transformer块1 → Transformer块2 → ... → Transformer块N →
最终层归一化 → 输出投影 → 词汇表概率分布

数学表示

对于输入序列 X=[x1,x2,...,xT]X = [x_1, x_2, ..., x_T]，GPT的计算过程为：

嵌入层：

E=TokenEmb(X)+PosEmb([1,2,...,T])E = \text{TokenEmb}(X) + \text{PosEmb}([1, 2, ..., T])

Transformer层：

H(l)=TransformerBlock(l)(H(l−1))H^{(l)} = \text{TransformerBlock}^{(l)}(H^{(l-1)})

输出层：

P(xt+1∣x1,...,xt)=softmax(Linear(LayerNorm(Ht(L))))P(x_{t+1}|x_1, ..., x_t) = \text{softmax}(\text{Linear}(\text{LayerNorm}(H^{(L)}_t)))

执行流程

1. 整体执行流程图

2. 详细计算流程图

计算步骤详解

步骤操作输入形状输出形状说明
1词元嵌入[batch, seq_len][batch, seq_len, emb_dim]将词元ID转换为向量
2位置嵌入[seq_len][seq_len, emb_dim]添加位置信息
3嵌入相加两个嵌入[batch, seq_len, emb_dim]融合词元和位置信息
4嵌入Dropout[batch, seq_len, emb_dim][batch, seq_len, emb_dim]防止过拟合
5Transformer块[batch, seq_len, emb_dim][batch, seq_len, emb_dim]特征提取和变换
6最终归一化[batch, seq_len, emb_dim][batch, seq_len, emb_dim]稳定输出
7输出投影[batch, seq_len, emb_dim][batch, seq_len, vocab_size]映射到词汇表

完整代码实现

GPT模型实现.py
import torch
import torch.nn as nn
import math


class GELU(nn.Module):
    """GELU激活函数实现"""
    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))
        ))


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


class MultiHeadAttention(nn.Module):
    """多头注意力机制实现"""
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads

        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.out_proj = nn.Linear(d_out, d_out)
        self.dropout = nn.Dropout(dropout)

        # 注册因果掩码
        self.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1))

    def forward(self, x):
        b, num_tokens, d_in = x.shape

        # 计算查询、键、值
        keys = self.W_key(x)
        queries = self.W_query(x)
        values = self.W_value(x)

        # 重塑为多头格式
        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
        values = values.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算注意力分数
        attn_scores = queries @ keys.transpose(2, 3)

        # 应用掩码
        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]
        attn_scores.masked_fill_(mask_bool, -torch.inf)

        # 应用softmax和dropout
        attn_weights = torch.softmax(attn_scores / math.sqrt(self.head_dim), dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 计算上下文向量
        context_vec = (attn_weights @ values).transpose(1, 2)
        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)

        # 输出投影
        context_vec = self.out_proj(context_vec)

        return context_vec


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)


class TransformerBlock(nn.Module):
    """Transformer块实现"""
    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
        x = self.norm1(x)
        x = self.att(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # 残差连接

        # 第二个子层：前馈网络
        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut  # 残差连接

        return x


class GPTModel(nn.Module):
    """GPT模型完整实现"""
    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"])

        # 创建多个Transformer块
        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)
        )

        # 嵌入相加和dropout
        x = tok_embeds + pos_embeds
        x = self.drop_emb(x)

        # 通过Transformer块
        x = self.trf_blocks(x)

        # 最终层归一化
        x = self.final_norm(x)

        # 输出投影到词汇表
        logits = self.out_head(x)

        return logits


def main():
    """主函数：测试GPT模型"""
    print("=" * 60)
    print("GPT模型测试")
    print("=" * 60)

    # GPT-124M配置
    GPT_CONFIG_124M = {
        "vocab_size": 50257,
        "context_length": 1024,
        "emb_dim": 768,
        "n_heads": 12,
        "n_layers": 12,
        "drop_rate": 0.1,
        "qkv_bias": False
    }

    print("配置信息:")
    for key, value in GPT_CONFIG_124M.items():
        print(f"  {key}: {value}")
    print()

    # 创建测试数据
    torch.manual_seed(123)
    batch_size, seq_len = 2, 4
    batch = torch.randint(0, GPT_CONFIG_124M["vocab_size"], (batch_size, seq_len))

    print(f"输入批次形状: {batch.shape}")
    print(f"输入词元ID:\n{batch}")
    print()

    # 创建GPT模型
    model = GPTModel(GPT_CONFIG_124M)

    # 前向传播
    with torch.no_grad():
        logits = model(batch)

    print(f"输出logits形状: {logits.shape}")
    print(f"输出logits示例 (第一个样本，第一个位置，前5个词汇):")
    print(f"{logits[0, 0, :5].tolist()}")
    print()

    # 模型参数统计
    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型总参数量: {total_params:,}")

    # 各层参数统计
    print("\n各层参数统计:")
    print(f"  词元嵌入层: {model.tok_emb.weight.shape} = {model.tok_emb.weight.numel():,}")
    print(f"  位置嵌入层: {model.pos_emb.weight.shape} = {model.pos_emb.weight.numel():,}")
    print(f"  输出投影层: {model.out_head.weight.shape} = {model.out_head.weight.numel():,}")

    # 权重绑定后的参数量（GPT-2风格）
    total_params_tied = total_params - model.out_head.weight.numel()
    print(f"\n权重绑定后参数量: {total_params_tied:,}")

    # 模型大小估算
    total_size_bytes = total_params * 4  # 假设float32
    total_size_mb = total_size_bytes / (1024 * 1024)
    print(f"模型大小估算: {total_size_mb:.2f} MB")

    print("=" * 60)
    print("测试完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()

运行结果

执行上述完整代码后，得到以下输出结果：

============================================================
GPT模型测试
============================================================
配置信息:
  vocab_size: 50257
  context_length: 1024
  emb_dim: 768
  n_heads: 12
  n_layers: 12
  drop_rate: 0.1
  qkv_bias: False

输入批次形状: torch.Size([2, 4])
输入词元ID:
tensor([[15742, 11036, 24681,  8041],
        [37378,  3791, 20331, 32351]])

输出logits形状: torch.Size([2, 4, 50257])
输出logits示例 (第一个样本，第一个位置，前5个词汇):
[-0.4764159023761749, -1.0278220176696777, 0.7373533844947815, 0.7530407309532166, -1.088520884513855]

模型总参数量: 163,009,536

各层参数统计:
  词元嵌入层: torch.Size([50257, 768]) = 38,597,376
  位置嵌入层: torch.Size([1024, 768]) = 786,432
  输出投影层: torch.Size([50257, 768]) = 38,597,376

权重绑定后参数量: 124,412,160
模型大小估算: 621.83 MB
============================================================
测试完成！
============================================================

结果分析

通过测试结果，我们可以验证GPT模型的正确性和特性：

关键验证点

验证项期望结果实际结果状态
输入形状[2, 4][2, 4]✅ 正确
输出形状[2, 4, 50257][2, 4, 50257]✅ 正确
参数量~124M124,412,160✅ 符合GPT-124M规格
模型大小~500-600MB621.83 MB✅ 合理

参数分布分析

组件参数量占比说明
词元嵌入层38,597,37631.0%词汇表到向量的映射
位置嵌入层786,4320.6%位置信息编码
Transformer块85,028,35268.4%核心计算组件
输出投影层38,597,376-与词元嵌入共享权重

关键发现

  1. 形状正确性：模型成功将词元ID序列转换为词汇表上的概率分布

  2. 参数规模：

    • 总参数量：1.63亿（包含输出层）
    • 有效参数量：1.24亿（权重绑定后）
    • 符合GPT-124M的标准规格

  3. 内存占用：模型大小约622MB，适合在现代GPU上运行

  4. 输出特性：logits值分布合理，没有出现数值异常

  5. 架构验证：成功实现了完整的GPT架构，包括嵌入层、Transformer块和输出层

代码详细解析

1. 类设计说明

GPTModel类结构

class GPTModel(nn.Module):
    def __init__(self, cfg):
        # 初始化嵌入层、Transformer块、归一化和输出层
    def forward(self, in_idx):
        # 实现完整的GPT前向传播流程

设计要点：

  • 继承自nn.Module，符合PyTorch模块化设计
  • 使用配置字典统一管理所有超参数
  • 模块化设计，每个组件职责明确
  • 支持任意长度的序列输入（受context_length限制）

2. 核心组件分析

嵌入层设计

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

关键特性：

  • 词元嵌入：将离散的词元ID映射为连续向量
  • 位置嵌入：为每个位置学习独特的位置表示
  • 嵌入Dropout：防止嵌入层过拟合

Transformer块堆叠

self.trf_blocks = nn.Sequential(
    *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])]
)

设计特点：

  • 深度堆叠：通过多层Transformer块增强模型表达能力
  • 参数共享：每层使用相同的架构但不同的参数
  • 梯度流动：通过残差连接确保梯度有效传播

输出层设计

self.final_norm = LayerNorm(cfg["emb_dim"])
self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)

关键组件：

  • 最终归一化：稳定输出特征分布
  • 输出投影：将隐藏状态映射到词汇表概率

3. 前向传播流程

嵌入阶段

# 词元嵌入
tok_embeds = self.tok_emb(in_idx)

# 位置嵌入
pos_embeds = self.pos_emb(
    torch.arange(seq_len, device=in_idx.device)
)

# 嵌入相加和dropout
x = tok_embeds + pos_embeds
x = self.drop_emb(x)

特征提取阶段

# 通过Transformer块
x = self.trf_blocks(x)

# 最终层归一化
x = self.final_norm(x)

输出生成阶段

# 输出投影到词汇表
logits = self.out_head(x)

关键组件解析：

组件输入形状输出形状作用
tok_emb[batch, seq_len][batch, seq_len, emb_dim]词元向量化
pos_emb[seq_len][seq_len, emb_dim]位置信息编码
trf_blocks[batch, seq_len, emb_dim][batch, seq_len, emb_dim]特征提取变换
final_norm[batch, seq_len, emb_dim][batch, seq_len, emb_dim]输出标准化
out_head[batch, seq_len, emb_dim][batch, seq_len, vocab_size]词汇表映射

4. 参数量和复杂度分析

各组件参数量计算

基于GPT-124M配置：

组件参数量计算公式参数数量
词元嵌入vocab_size × emb_dim38,597,376
位置嵌入context_length × emb_dim786,432
Transformer块n_layers × 7,085,56885,026,816
最终归一化2 × emb_dim1,536
输出投影emb_dim × vocab_size38,597,376
总计163,009,536

权重绑定优化

在GPT-2中，输出投影层与词元嵌入层共享权重：

  • 优化前参数量：163,009,536
  • 优化后参数量：124,412,160
  • 节省参数量：38,597,376 (23.7%)

计算复杂度分析

  • 时间复杂度：O(batch_size × seq_len² × emb_dim × n_layers)
  • 空间复杂度：O(batch_size × seq_len × emb_dim)
  • 主要瓶颈：注意力机制的二次复杂度

更新日志

2025/8/18 00:31
查看所有更新日志
  • bd1d0-迁移目录于
  • 8d9ff-feat: 新增大模型架构系列文档 - 快捷连接、Transformer块、GPT模型实现和文本生成于

版权归属：NateHHX

许可证：署名 4.0 国际 (CC-BY-4.0)