从零训练一个 Mini GPT — 实战学习笔记

这份笔记来自一次真实的训练实践：在 RTX 3070 (8GB VRAM) + 16GB RAM 上，从零实现并训练了一个 ~10M 参数的 GPT 语言模型。包含完整代码、踩过的坑、训练数据，适合作为深度学习对话的上下文。

---

一、GPT 是什么？一句话版本

GPT = 只有 Decoder 的 Transformer，任务是"给定前面的文字，预测下一个字"（自回归语言模型）。

训练时：喂一段文本，模型尝试预测每个位置的下一个 token，用 cross-entropy loss 衡量预测好坏。 生成时：给一个开头，模型一个字一个字往后"续写"。

---

二、核心架构（逐层拆解）

2.1 整体结构

输入 token ids → Token Embedding + Position Embedding → Dropout
    → [Transformer Block] × N 层
    → LayerNorm → Linear (lm_head) → logits (词表大小的概率分布)

2.2 Token Embedding 和 Position Embedding

wte = nn.Embedding(vocab_size, n_embd)   # 把 token id 映射成向量
wpe = nn.Embedding(block_size, n_embd)   # 把位置 0,1,2,...,255 映射成向量
x = wte(token_ids) + wpe(positions)       # 直接相加

• 为什么需要位置编码？ Transformer 内部没有"顺序"概念（不像 RNN），所以必须显式告诉它每个 token 在第几个位置
• 这里用的是可学习的位置编码（直接当参数学），GPT-2/3 都是这么做的
• 另一种方案是 sinusoidal 固定编码（原始 Transformer 论文用的），现在不太主流了

2.3 Transformer Block（核心！）

每一层 = 多头因果自注意力 + 前馈网络，都有残差连接和 LayerNorm。

x → LayerNorm → CausalSelfAttention → + x (残差)
  → LayerNorm → MLP                  → + x (残差)

注意这里是 Pre-LN（先 Norm 再计算），而不是原始论文的 Post-LN。Pre-LN 训练更稳定，现代模型基本都用这个。

2.4 因果自注意力 (Causal Self-Attention)

这是 GPT 的灵魂。

# 一次线性变换得到 Q, K, V（效率高）
q, k, v = self.c_attn(x).split(n_embd, dim=2)  # 形状: (B, T, C) 各三份

# 重塑为多头: (B, n_head, T, head_dim)
q = q.view(B, T, n_head, head_dim).transpose(1, 2)

# Scaled dot-product attention with causal mask
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

关键点： - Q (Query), K (Key), V (Value)：来自同一个输入的三个不同线性投影 - Attention 公式：softmax(QK^T / sqrt(d_k)) · V - QK^T：每个位置和其他位置的"相关度"得分 - / sqrt(d_k)：缩放，防止点积太大导致 softmax 梯度消失 - softmax：归一化为概率分布 - · V：用相关度加权聚合信息 - Causal mask（因果掩码）：位置 i 只能看到 0~i，不能偷看未来。这就是"GPT 只能从左往右"的原因 - 多头：把 embedding 拆成多个小空间（head_dim = n_embd / n_head），各自算注意力再拼回来。不同的头可以关注不同类型的关系

2.5 MLP（前馈网络）

x → Linear(n_embd, 4*n_embd) → GELU → Linear(4*n_embd, n_embd) → Dropout

• 先升维 4 倍，过激活函数，再降回来
• GELU 比 ReLU 更平滑，GPT 系列标配
• 直觉：注意力负责"信息交换"，MLP 负责"信息加工"

2.6 Weight Tying

self.transformer.wte.weight = self.lm_head.weight  # 共享权重！

输入的 token embedding 矩阵和最后的输出投影矩阵是同一个矩阵。这样做： - 减少参数量（对于大 vocab 效果显著） - 语义上合理：embedding 空间里相似的 token，输出概率也应该相近

---

三、训练细节

3.1 我们的配置

参数	值	说明
n_layer	6	Transformer 层数
n_head	6	注意力头数
n_embd	384	嵌入维度
block_size	256	上下文窗口长度
vocab_size	65	字符级，Shakespeare 语料的不同字符数
dropout	0.1	正则化
总参数量	~10M	小模型，RTX 3070 轻松跑

对比：GPT-2 Small 是 12 层 / 12 头 / 768 维 / 117M 参数，GPT-3 是 96 层 / 96 头 / 12288 维 / 175B 参数。

3.2 Tokenizer 选择

我们用了字符级 tokenizer（每个字符 = 一个 token）。

为什么不用 tiktoken / BPE？ - tiktoken (cl100k_base) 的 vocab_size = 100,277 - Embedding 层大小 = vocab_size × n_embd = 100277 × 384 ≈ 38M 参数，光这一层就比整个模型还大 - 在 8GB VRAM 上直接 OOM（我们实际踩了这个坑）

字符级的代价： - 序列变长（一个中文字 = 1 个 token，但英文一个词可能 5-10 个 token） - 对于大型模型不高效，但用于学习完全够了

实际工业界用的 BPE (Byte Pair Encoding)：折中方案，常见词整个编码为一个 token，罕见词拆成子词。GPT-2/3/4 都用 BPE。

3.3 数据集准备

# 文本 → token ids
tokens = tokenizer.encode(text)  # [23, 15, 42, 8, ...]

# 切成训练样本：输入 = tokens[i:i+256], 目标 = tokens[i+1:i+257]
# 模型学习：给定前 256 个字符，预测第 2~257 个字符

关键概念：一个训练样本里其实包含 256 个预测任务（位置 0 预测位置 1，位置 0-1 预测位置 2，...）。这就是为什么语言模型训练效率高。

3.4 优化器和学习率

AdamW(lr=1e-3, weight_decay=0.1, betas=(0.9, 0.95))

• AdamW：Adam + 正确的 weight decay。几乎所有现代 LLM 都用这个
• Weight Decay = 0.1：L2 正则化，防止权重过大
• Cosine Learning Rate Schedule with Warmup：
  • 前 200 步：学习率从 0 线性升到 1e-3（warmup，避免初期梯度不稳定）
  • 之后：余弦衰减到 1e-4
  • 这个调度策略是 GPT-3 论文确立的标准做法

3.5 混合精度训练 (FP16)

scaler = torch.amp.GradScaler('cuda')
with torch.amp.autocast('cuda'):
    logits, loss = model(xb, yb)
scaler.scale(loss).backward()

• 前向和部分反向用 FP16（半精度），显存减半，速度翻倍
• GradScaler：FP16 的数值范围小，梯度可能 underflow → scaler 先放大 loss 再缩回来
• 权重更新仍然在 FP32，保证精度

3.6 梯度裁剪

torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

如果梯度范数 > 1.0，就按比例缩小。防止训练不稳定（gradient explosion）。

---

四、训练结果和分析

4.1 实际训练曲线

Step     Train Loss    Val Loss     观察
  250      2.12         2.15       两个 loss 很接近，还在学基础模式
  500      1.53         1.56       快速下降，学会了字母频率和常见词
 1000      1.27         1.31       继续收敛
 1500      1.14         1.26       开始出现 gap！val 下降变慢了
 2000      1.02         1.28       val 开始上升 → 过拟合信号
 3000      0.73         1.49       差距加大，模型在"背"训练数据
 5000      0.46         1.85       严重过拟合

最佳验证 loss = 1.26（在 step 1500 附近）

4.2 过拟合分析

为什么过拟合这么严重？ - 数据太少：Shakespeare 只有 ~1MB / 30 万 tokens。10M 参数的模型完全有能力"背下来" - 模型相对太大：参数量 / 数据量比例失衡 - 对比：GPT-3 175B 参数用了 300B tokens 训练（数据量是参数量的 ~1700 倍），我们是 30 万 tokens / 10M 参数 = 0.03 倍

解决方案（如果要认真训）： 1. 更多数据（最有效）：至少 10-100MB 文本 2. Early stopping：在 val loss 开始上升时停止（~step 1500） 3. 增大 dropout：从 0.1 → 0.2 或 0.3 4. 减小模型：3 层 / 4 头 / 256 维 5. 数据增强：随机 crop 不同的窗口起点（我们已经在做了）

4.3 生成效果

模型学会了 Shakespeare 的格式和风格：

KING RICHARD II:
That would he die, my lord for her dead?

DUCHESS OF YORK:
Thanks, noble prince. Farewell: be patient.
And shall I think my most humble presently.

能观察到： - ✅ 学会了"角色名: 台词"的格式 - ✅ 学会了 Early Modern English 的用词（thee, shall, noble 等） - ✅ 语法基本通顺 - ❌ 语义是胡编的（10M 参数学不到真正的理解力）

---

五、核心代码（完整可运行）

5.1 模型定义 (model.py)

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        y = F.scaled_dot_product_attention(
            q, k, v, dropout_p=self.attn_dropout.p if self.training else 0, is_causal=True
        )
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.gelu = nn.GELU()
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),
            wpe=nn.Embedding(config.block_size, config.n_embd),
            drop=nn.Dropout(config.dropout),
            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f=nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.vocab_size, config.n_embd, bias=False)
        self.transformer.wte.weight = self.lm_head.weight  # weight tying

    def forward(self, idx, targets=None):
        B, T = idx.size()
        pos = torch.arange(T, device=idx.device)
        x = self.transformer.drop(self.transformer.wte(idx) + self.transformer.wpe(pos))
        for block in self.transformer.h:
            x = block(x)
        logits = self.lm_head(self.transformer.ln_f(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=40):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            if top_k:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            idx = torch.cat((idx, torch.multinomial(F.softmax(logits, -1), 1)), dim=1)
        return idx

5.2 生成时的采样策略

logits = logits[:, -1, :] / temperature  # temperature 控制随机性

• temperature = 1.0：按原始概率采样
• temperature < 1.0（如 0.8）：概率分布更尖锐 → 输出更确定、更保守
• temperature > 1.0：更平坦 → 输出更随机、更有创意
• temperature → 0：退化为 greedy decoding（总选概率最高的）

v, _ = torch.topk(logits, top_k)       # 只保留概率最高的 top_k 个
logits[logits < v[:, [-1]]] = -float('Inf')  # 其余设为 -inf（概率变 0）

Top-k 采样：只从最可能的 k 个 token 中选，避免采到奇怪的低概率 token。

还有一种常见策略是 Top-p (nucleus) 采样：不固定数量，而是保留累积概率 >= p 的最少 token。GPT-4 等实际用的是 top-p。

---

六、关键概念深入

6.1 为什么 Transformer 比 RNN 好？

	RNN/LSTM	Transformer
处理长距离依赖	困难（信息需要一步步传递，容易遗忘）	容易（注意力直接连接任意两个位置）
并行化	不行（必须串行处理序列）	可以（所有位置同时计算）
训练速度	慢	快（GPU 并行优势巨大）
缺点	梯度消失/爆炸	显存 O(T²)，长序列吃力

6.2 残差连接为什么重要？

x = x + self.attn(self.ln_1(x))  # 而不是 x = self.attn(self.ln_1(x))

• 没有残差：6 层网络，梯度要穿过 6 个非线性变换 → 容易消失
• 有残差：梯度可以直接"走捷径"回到早期层 → 深层网络可训练
• 也叫 "skip connection"，ResNet 的核心思想

6.3 LayerNorm 做了什么？

# 对每个样本的 embedding 维度做归一化
mean = x.mean(dim=-1, keepdim=True)
std = x.std(dim=-1, keepdim=True)
x = (x - mean) / (std + eps) * gamma + beta  # gamma, beta 可学习

为什么需要？每一层的输出分布会变化（internal covariate shift），归一化让每层输入稳定，训练更顺畅。

6.4 VRAM 占用分析

8GB VRAM 的主要消耗： 1. 模型参数：10M × 4 bytes (FP32) = 40MB（FP16 时 20MB） 2. 梯度：和参数等大 = 40MB 3. 优化器状态：AdamW 存 m 和 v = 80MB 4. 激活值（中间计算结果，用于反向传播）：这才是大头！ - batch_size × seq_len × n_embd × n_layer ≈ 64 × 256 × 384 × 6 × 2 bytes ≈ 72MB 每层 5. Attention 矩阵：batch × heads × T × T = 64 × 6 × 256 × 256 × 2 = 48MB 每层

实际踩坑：用 tiktoken (vocab=100K) 时，光 embedding 层就要 100K × 384 × 2 = 76MB，加上 lm_head 的梯度和优化器状态，直接爆显存。

---

七、从小 GPT 到 ChatGPT 的差距

我们训练的只是 GPT 的"地基"。真正的 ChatGPT 还需要：

1. 规模：175B+ 参数，数万亿 token 训练 → 涌现出推理能力
2. SFT (Supervised Fine-Tuning)：用人工标注的"问题→高质量回答"数据微调
3. RLHF (Reinforcement Learning from Human Feedback)：
   • 训练一个 Reward Model，学习人类偏好
   • 用 PPO 算法优化模型输出，使其更符合人类期望
4. 更好的 Tokenizer：BPE，几万级别的 vocab
5. 更长的上下文：RoPE 位置编码支持 128K+ tokens
6. 各种工程优化：Flash Attention、KV Cache、分布式训练等

但核心的 Transformer decoder 架构？和我们实现的一模一样。

---

八、动手练习建议

如果你想基于这份代码继续实验：

1. 调参实验：改 n_layer、n_head、n_embd，观察参数量和效果变化
2. 换数据集：用中文小说、诗词、代码等不同语料训练，对比生成效果
3. 实现 BPE tokenizer：手写一个 Byte Pair Encoding，理解子词分割
4. 加 early stopping：val loss 连续 3 次不下降就停止
5. 可视化注意力：把 attention weight 画成热力图，看模型在"关注"什么
6. 对比实验：去掉残差连接 / 去掉 LayerNorm / 去掉位置编码，看模型会怎样崩

---

九、推荐资源

• Andrej Karpathy "Let's build GPT from scratch" (YouTube) — 本项目的灵感来源
• "Attention Is All You Need" (2017) — Transformer 原始论文
• "Language Models are Few-Shot Learners" (2020) — GPT-3 论文
• The Illustrated Transformer (Jay Alammar) — 最好的可视化讲解
• nanoGPT GitHub repo (karpathy/nanoGPT) — 完整参考实现

---

实践环境：RTX 3070 8GB / 16GB RAM / Ubuntu / PyTorch 2.10 / CUDA 13.1 训练时间：~7 分钟 / 5000 步 生成于 2026-03-02