【大模型】- 预训练 Mini-GPT

宋 算法
  2026-08-19 09:04 2026-08-19 09:04 创建   2026-08-21 16:20:12 2026-08-21 16:20:12 更新      1.7k 字  7 分钟

预训练 Mini-GPT

从头构建并训练一个小语言模型

类型： 学习　｜　语言： Python　｜　🏷 前置：《数据管道》（本系列第 3 篇）

学习目标

  • 实现一个小型 GPT 架构用于预训练
  • 理解下一个 token 预测的训练目标
  • 设置训练循环并监控损失
  • 生成文本以评估训练进度
  • 保存和加载模型检查点

预训练概述

预训练是语言模型学习的主要阶段。模型在大规模文本上训练，以预测下一个 token。这使得模型学习语法、语义、世界知识和推理模式。

模型架构

我们将构建一个小型 GPT 风格的模型：

1
2
3
4
5
Token Embeddings + Position Embeddings → 
Transformer Block × N → 
Layer Norm → 
Linear Head → 
Logits

配置

1
2
3
4
5
6
7
8
9
10
11
12
config = {
    "vocab_size": 50257,  # GPT-2 词汇表大小
    "block_size": 256,    # 上下文长度
    "n_embd": 256,        # 嵌入维度
    "n_head": 8,          # 注意力头数
    "n_layer": 6,         # Transformer 块数
    "dropout": 0.1,       # Dropout 率
    "lr": 3e-4,           # 学习率
    "batch_size": 32,     # 批次大小
    "max_iters": 5000,    # 最大迭代次数
    "eval_interval": 500, # 评估间隔
}

实现

注意力机制

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import torch
import torch.nn as nn
import torch.nn.functional as F

class Head(nn.Module):
    def __init__(self, config, head_size):
        super().__init__()
        self.key = nn.Linear(config["n_embd"], head_size, bias=False)
        self.query = nn.Linear(config["n_embd"], head_size, bias=False)
        self.value = nn.Linear(config["n_embd"], head_size, bias=False)
        self.register_buffer("tril", torch.tril(
            torch.ones(config["block_size"], config["block_size"])
        ))
        self.dropout = nn.Dropout(config["dropout"])
    
    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)
        q = self.query(x)
        v = self.value(x)
        
        # 注意力权重
        wei = q @ k.transpose(-2, -1) * C**-0.5
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf"))
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)
        
        return wei @ v

class MultiHeadAttention(nn.Module):
    def __init__(self, config, n_heads):
        super().__init__()
        head_size = config["n_embd"] // n_heads
        self.heads = [Head(config, head_size) for _ in range(n_heads)]
        self.proj = nn.Linear(config["n_embd"], config["n_embd"])
        self.dropout = nn.Dropout(config["dropout"])
    
    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        return self.dropout(self.proj(out))

Feed-Forward 网络

1
2
3
4
5
6
7
8
9
10
11
12
class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config["n_embd"], 4 * config["n_embd"]),
            nn.GELU(),
            nn.Linear(4 * config["n_embd"], config["n_embd"]),
            nn.Dropout(config["dropout"]),
        )
    
    def forward(self, x):
        return self.net(x)

Transformer 块

1
2
3
4
5
6
7
8
9
10
11
12
class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln1 = nn.LayerNorm(config["n_embd"])
        self.attn = MultiHeadAttention(config, config["n_head"])
        self.ln2 = nn.LayerNorm(config["n_embd"])
        self.ff = FeedForward(config)
    
    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ff(self.ln2(x))
        return x

完整模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        self.token_embedding = nn.Embedding(config["vocab_size"], config["n_embd"])
        self.position_embedding = nn.Embedding(config["block_size"], config["n_embd"])
        self.blocks = nn.Sequential(*[Block(config) for _ in range(config["n_layer"])])
        self.ln_f = nn.LayerNorm(config["n_embd"])
        self.head = nn.Linear(config["n_embd"], config["vocab_size"], bias=False)
        
        # 权重共享
        self.head.weight = self.token_embedding.weight
        
        # 初始化
        self.apply(self._init_weights)
    
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            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.0, std=0.02)
    
    def forward(self, idx, targets=None):
        B, T = idx.shape
        
        # 嵌入
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(T, device=idx.device))
        x = tok_emb + pos_emb
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.head(x)
        
        # 计算损失
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1)
            )
        
        return logits, loss

训练循环

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
model = GPT(config)
optimizer = torch.optim.AdamW(model.parameters(), lr=config["lr"])

def estimate_loss():
    model.eval()
    losses = []
    for _ in range(200):
        xb, yb = get_batch("val")
        logits, loss = model(xb, yb)
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)

for iter in range(config["max_iters"]):
    # 定期评估
    if iter % config["eval_interval"] == 0:
        losses = estimate_loss()
        print(f"step {iter}: train loss: {losses:.4f}")
    
    # 获取批次
    xb, yb = get_batch("train")
    
    # 前向传播
    logits, loss = model(xb, yb)
    
    # 反向传播
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

文本生成

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0):
    model.eval()
    for _ in range(max_new_tokens):
        # 裁剪到 block_size
        idx_cond = idx[:, -config["block_size"]:]
        
        # 获取预测
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :] / temperature
        
        # 采样
        probs = F.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1)
        
        # 追加
        idx = torch.cat((idx, idx_next), dim=1)
    
    return idx

# 生成文本
context = torch.zeros((1, 1), dtype=torch.long)
generated = generate(model, context, max_new_tokens=500)
print(tokenizer.decode(generated[0].tolist()))

保存和加载

1
2
3
4
5
6
7
8
9
10
11
# 保存
torch.save({
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "config": config,
}, "checkpoint.pt")

# 加载
checkpoint = torch.load("checkpoint.pt")
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])

训练监控

监控以下指标：

  • 训练损失：应该稳步下降
  • 验证损失：检查过拟合（如果开始上升）
  • 生成质量：定性检查
  • 学习率：可以衰减以稳定训练

下一步

下一课将扩展到分布式训练，以处理更大的模型和数据集。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

📝 自我检查（课程配套测验）

Q1（学前） GPT 预训练使用什么训练目标？

A. 掩码语言建模（预测被掩码的 token）
B. 下一 token 预测：给定先前 token，预测下一个
C. 句子分类
D. 图文对齐

答案： B　解析： GPT 是因果（自回归）语言模型，用下一 token 预测训练。给定 [t1, t2, …, tn]，学习预测 tn+1。损失是预测与真实下一 token 的交叉熵。

Q2（学前） GPT-2 Small（124M）有多少 transformer 层、注意力头和嵌入维度？

A. 6 层、6 头、512 维
B. 12 层、12 头、768 维
C. 24 层、16 头、1024 维
D. 48 层、25 头、1600 维

答案： B　解析： GPT-2 Small 有 12 层 transformer、每层 12 个注意力头、768 维嵌入。该架构约 1.24 亿参数，可在单 GPU 上几小时完成训练。

Q3（学后） GPT 中因果注意力掩码的作用是什么？

A. 防止关注 padding token
B. 防止每个 token 关注未来 token，确保预测时只能使用过去上下文
C. 掩蔽低置信注意力分数
D. 减少训练内存

答案： B　解析： 因果掩码是三角矩阵，在 softmax 前将未来位置设为负无穷。位置 5 的 token 可关注 1–5 但不能关注 6+。确保模型从左到右生成 token。

Q4（学后） 文本生成中「温度」控制什么？

A. 生成速度
B. token 选择的随机性：温度越低输出越确定，越高越多样
C. 生成 token 数量
D. 模型置信阈值

答案： B　解析： 温度在 softmax 前除以 logits。温度=0.1 使分布非常尖锐（近乎确定）。温度=1.0 是训练分布。温度>1.0 使其平坦，增加随机性。

Q5（学后） 为什么预训练比微调需要更多算力？

A. 预训练使用更大 batch
B. 预训练从零处理数万亿 token 学习通用语言模式，微调在已有能力模型上用数千样本调整
C. 预训练用不同架构
D. 微调不用梯度

答案： B　解析： 预训练从随机权重在数万亿 token 上构建全部语言知识。微调从这些已学权重出发，在更小数据集（数千到数百万样本）上调整。

  • 标题: 【大模型】- 预训练 Mini-GPT
  • 作者: 宋
  • 创建于 : 2026-08-19 09:04:00
  • 更新于 : 2026-08-21 16:20:12
  • 链接: https://sxl-space.tk/2026/08/19/010_LLM/010_LLM-04-PreTrainingMiniGPT/
  • 版权声明: 版权所有 © 宋，禁止转载。
目录
【大模型】- 预训练 Mini-GPT
  1. 1. 预训练 Mini-GPT
  2. 2. 学习目标
  3. 3. 预训练概述
  4. 4. 模型架构
    1. 4.1. 配置
  5. 5. 实现
    1. 5.1. 注意力机制
    2. 5.2. Feed-Forward 网络
    3. 5.3. Transformer 块
    4. 5.4. 完整模型
  6. 6. 训练循环
  7. 7. 文本生成
  8. 8. 保存和加载
  9. 9. 训练监控
  10. 10. 下一步
  11. 11. 下一课将扩展到分布式训练，以处理更大的模型和数据集。
  12. 12. 📝 自我检查（课程配套测验）