看不懂这些代码？

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
optimizer.zero_grad(set_to_none=True)

这是 MiniMind 训练循环的核心五行。如果你觉得像天书，没关系——这一章把它们逐行拆开。

本章路线图

Part 1：Python 速查

MiniMind 全部用 Python 编写。这一节只覆盖项目中反复出现的语法，不追求面面俱到。

数据类型速查

# MiniMind 配置文件 config.py 中实际出现的数据类型
dim = 768                       # int: 隐藏层维度
n_layers = 8                    # int: Transformer 层数
lr = 5e-4                       # float: 学习率
data_path = "dataset/pretrain"  # str: 数据路径
vocab = [i for i in range(6400)]  # list: 词表索引
msg = {"role": "user", "content": "你好"}  # dict: 对话

控制流与函数

if / for / 函数是读懂任何代码的基础

# ---- 条件分支：MiniMind 中的典型模式 ----
if use_cache:                          # bool 控制是否缓存 KV
    xk, xv = past_kv[layer_idx]
else:
    xk, xv = self.k_proj(x), self.v_proj(x)

# ---- 循环：逐层过 Transformer ----
for layer in self.layers:
    x = layer(x, start_pos, mask)

# ---- 函数：封装可复用的逻辑 ----
def get_lr(step, total_steps, base_lr):
    """余弦学习率调度——训练中动态调整学习率"""
    return base_lr * (0.1 + 0.45 * (
        1 + math.cos(math.pi * step / total_steps)
    ))

Lambda 是一行函数的简写。项目中偶尔出现：key=lambda x: x[1]，意思是"以第二个元素排序"。

推导式与格式化

两种 Pythonic 写法，项目中几乎每页都出现

# ---- 列表推导：一行创建列表 ----
layers = [TransformerBlock(config) for _ in range(n_layers)]
# 等价于：
# layers = []
# for _ in range(n_layers):
#     layers.append(TransformerBlock(config))

# ---- f-string：在字符串中嵌入变量 ----
epoch, total, loss, lr = 1, 10, 3.1416, 0.0005
print(f'Epoch:[{epoch+1}/{total}] loss:{loss:.4f} lr:{lr:.8f}')
# 输出: Epoch:[2/10] loss:3.1416 lr:0.00050000
#   {loss:.4f}  → 保留 4 位小数
#   {lr:.8f}   → 保留 8 位小数

面向对象编程

class 把数据和操作绑在一起，MiniMind 全是 class

nn.Module                    # PyTorch 所有网络层的基类
  ├── parameters()            # 收集所有可学习参数
  ├── forward(x)              # 定义前向计算
  ├── state_dict()            # 导出参数字典（保存模型）
  └── load_state_dict()       # 加载参数字典（加载模型）

MiniMindModel(nn.Module)     # 继承 nn.Module
  ├── embed_tokens            # 词嵌入层
  ├── layers                  # Transformer Block 列表
  ├── norm                    # 最终归一化
  └── forward(input_ids)      # 前向传播

Attention(nn.Module)          # 也继承 nn.Module
  ├── q_proj: Linear          # Query 投影
  ├── k_proj: Linear          # Key 投影
  ├── v_proj: Linear          # Value 投影
  └── forward(x, mask)        # 注意力计算

继承与 self

三步理解面向对象：定义蓝图 → 创建对象 → 调用方法

class Attention(nn.Module):          # ① 继承 nn.Module
    def __init__(self, dim):          # ② 构造函数：买积木
        super().__init__()            # 调用父类构造（必须）
        self.q_proj = nn.Linear(dim, dim, bias=False)  # 用 self 存储层

    def forward(self, x):             # ③ 前向计算：搭积木
        return self.q_proj(x)         # self = "我自己"

# 使用：创建对象 → 调方法
attn = Attention(dim=768)             # 触发 __init__
out = attn.forward(x)                 # 触发 forward

• 继承：class Child(Parent)，子类获得父类全部能力
• self：对象自身的引用，类似日常说的"我的"
• super().__init__()：先让父类完成初始化

装饰器与生成器

两个 Python 进阶语法，在项目中反复出现

# ---- 装饰器：给函数"穿外套"，不改动原代码 ----
@torch.no_grad()          # 推理时不计算梯度，节省显存
def generate(self, ...):
    ...

@abstractmethod           # 强制子类必须实现此方法
def rollout(self, ...):
    pass

# ---- 生成器 yield：一次只吐一行，不把文件全读进内存 ----
def get_training_texts(path):
    with open(path, 'r') as f:
        for line in f:
            yield line.strip()   # 暂停，把这一行交给调用方

for text in get_training_texts("data.txt"):
    tokens = tokenizer(text)     # 逐行处理，内存友好

Python 常见陷阱

练习：Python

任务：阅读以下代码，回答三个问题。

class MyModel(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.linear = nn.Linear(dim, dim)

    @torch.no_grad()
    def predict(self, x):
        return self.linear(x)

model = MyModel(dim=768)
print(sum(p.numel() for p in model.parameters()))

问题

1. super().__init__() 如果删掉，会怎样？
2. @torch.no_grad() 在这里的效果是什么？
3. 最后一行会输出什么数字？（提示：nn.Linear(768, 768) 有 bias）

Part 2：PyTorch 核心

MiniMind 使用 PyTorch 框架。这一部分覆盖模型定义、训练循环、数据加载中反复出现的 API。

Tensor 张量

定义：多维数组 | 白话：GPU 上的 NumPy | 用途：模型中所有数据都是 Tensor

import torch
x = torch.tensor([1.0, 2.0, 3.0])   # 从列表创建
print(x.shape)    # torch.Size([3])
print(x.device)   # cpu 或 cuda:0

# MiniMind 中的维度标注
# input_ids:  (batch, seq_len)       例如 (4, 512)
# embeddings: (batch, seq_len, 768)  例如 (4, 512, 768)
# logits:     (batch, seq_len, 6400) 例如 (4, 512, 6400)

MiniMind 数据流

从输入文本到输出概率，维度如何一步步变化

每个方框对应模型中的一个 nn.Module。数据的 shape 在每个阶段都标注出来——这是调试 shape 错误时最重要的参考。

# 在 MiniMind 中你可以这样打印每一步的 shape
def forward(self, input_ids):
    x = self.embed_tokens(input_ids)  # (B, seq, 768)
    for layer in self.layers:
        x = layer(x)                  # 每层保持 (B, seq, 768)
    x = self.norm(x)                  # 最后归一化
    logits = self.lm_head(x)          # (B, seq, 6400)
    return logits

Tensor 高频操作

这些操作在 MiniMind 源码中反复出现

# ===== 形状操作（MiniMind 中最频繁） =====
x.view(batch, seq_len, 768)       # 改形状，不复制数据
x.unsqueeze(1)                    # 加维度: (B,D) → (B,1,D)
x.transpose(1, 2)                 # 交换两个维度
torch.cat([a, b], dim=1)          # 沿维度1拼接（KV-Cache）

# ===== 索引（取特定位置的数据） =====
x[:, -1, :]                       # 取最后一个 token 的表示
x[mask == 1]                      # 按条件筛选

# ===== 数学运算（神经网络的核心） =====
torch.softmax(x, dim=-1)          # 转概率分布（生成时采样用）
torch.gather(logits, 2, idx)      # 按索引取值（提取 log 概率）
torch.topk(x, k=50)               # 取最大的 k 个（Top-K 采样）

自动求导

定义：PyTorch 自动追踪运算并计算梯度 | 白话：你说"算"，它帮你求导 | 用途：训练的核心机制

import torch
w = torch.tensor([2.0], requires_grad=True)  # 标记：需要对 w 求导
y = w ** 2 + 3 * w       # 前向传播：y = w² + 3w
y.backward()             # 反向传播：自动算 dy/dw
print(w.grad)            # tensor([7.]) → 2w+3=2×2+3=7

# 关键：只有 requires_grad=True 的张量才会计算梯度
# 推理时用 torch.no_grad() 可以完全关闭追踪，节省显存

训练循环五步

MiniMind 中 train_pretrain.py 的核心——记住这五步就够了

# ① 前向传播：输入数据，得到预测和损失
res = model(input_ids, labels=labels)
loss = res.loss

# ② 反向传播：计算每个参数的梯度
scaler.scale(loss).backward()

# ③ 梯度裁剪：防止梯度爆炸
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

# ④ 更新参数：用梯度调整权重
scaler.step(optimizer)

# ⑤ 梯度清零：为下一步准备
optimizer.zero_grad(set_to_none=True)

这五步在每一个训练 step 中重复执行。scaler 的作用见混合精度幻灯片。

nn.Module

定义：神经网络层的基类 | 白话：所有层都是它的子类 | 用法：__init__ 买积木，forward 搭积木

class Attention(nn.Module):
    def __init__(self, dim):
        super().__init__()
        # __init__：定义需要哪些层（买积木）
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim // 2, bias=False)
        self.v_proj = nn.Linear(dim, dim // 2, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)

    def forward(self, x, mask=None):
        # forward：定义计算流程（搭积木）
        q = self.q_proj(x)          # (B, seq, dim)
        k = self.k_proj(x)          # (B, seq, dim//2)
        v = self.v_proj(x)
        # ... 注意力计算（阶段 2 详解）
        return self.out_proj(attn_output)

关键特性：在 __init__ 中赋值给 self 的 nn.Module 会被自动追踪。model.parameters() 能递归收集所有参数。

比喻卸甲：上面"买积木/搭积木"只是为了建立直观感受。在 PyTorch 内部，nn.Module 实际上是通过一个 _modules 字典注册子层、在 forward 中按你写的计算图顺序调用它们——并不存在物理意义上的"拼接"，子层之间是函数调用关系，不是结构连接。

优化器与学习率

AdamW 控制参数更新方向和步长

import torch.optim as optim

optimizer = optim.AdamW(
    model.parameters(),  # 需要优化的参数
    lr=5e-4,             # 学习率：每步更新多大
    weight_decay=0.1     # 权重衰减：防止过拟合
)

# MiniMind 使用余弦学习率调度
def get_lr(step, total, lr):
    return lr / (1 + step * 1e-5) * \
           (0.1 + 0.45 * (1 + math.cos(
               math.pi * step / total)))

学习率太大 → 训练发散；太小 → 收敛太慢。余弦调度是实践中最好用的策略之一。

DataLoader

定义：数据加载管道 | 白话：自动帮你分批、打乱、并行加载 | 用途：喂给模型训练数据

from torch.utils.data import Dataset, DataLoader

class PretrainDataset(Dataset):
    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        text = self.samples[idx]
        tokens = self.tokenizer(text)
        return torch.tensor(tokens[:-1]), torch.tensor(tokens[1:])
        # 输入是 [0..n-1]，目标是 [1..n]（预测下一个 token）

loader = DataLoader(dataset, batch_size=32, shuffle=True)
for input_ids, labels in loader:  # 每次取一个 batch
    loss = model(input_ids, labels=labels).loss

混合精度训练

定义：用 BF16 算前向，FP32 存权重 | 白话：算得快、省显存、精度不掉 | 用途：MiniMind 默认开启

# MiniMind 中的混合精度写法
scaler = torch.amp.GradScaler('cuda')

with torch.amp.autocast('cuda', dtype=torch.bfloat16):
    # 前向用 bf16（每个数 2 字节，比 fp32 的 4 字节省一半）
    res = model(input_ids, labels=labels)
    loss = res.loss

scaler.scale(loss).backward()    # 反向传播（自动缩放梯度）
scaler.step(optimizer)           # 更新参数（自动转回 fp32）
optimizer.zero_grad(set_to_none=True)

模型保存与加载

训练完了要存盘，推理时再读回来

# ---- 保存模型 ----
torch.save(model.state_dict(), "model.pth")
# state_dict() 返回一个 dict，key=层名, value=参数张量

# ---- 加载模型 ----
model.load_state_dict(torch.load("model.pth"))
model.eval()  # 切换到评估模式（关闭 Dropout 等）

# ---- MiniMind 中的实际用法 ----
# 训练：每 N 步自动保存 checkpoint
# 推理：加载 checkpoint → generate() 生成文本
# HuggingFace 格式：save_pretrained() / from_pretrained()

PyTorch 排错指南

练习：PyTorch

任务：在终端运行以下代码，观察 Tensor 形状变化。

>>> import torch
>>> x = torch.randn(2, 3, 4)
>>> print("原始:", x.shape)            # torch.Size([2, 3, 4])
>>> print("view:", x.view(2, 12).shape) # torch.Size([2, 12])
>>> print("unsqueeze:", x.unsqueeze(1).shape)  # torch.Size([2, 1, 3, 4])
>>> print("transpose:", x.transpose(1,2).shape) # torch.Size([2, 4, 3])
>>> print("最后token:", x[:, -1, :].shape)      # torch.Size([2, 4])

挑战

1. 把 view(2, 12) 改成 view(2, 11)，看报什么错？为什么？
2. unsqueeze(1) 和 unsqueeze(-1) 的结果有什么不同？
3. 对 transpose 后的张量直接 .view(2, 12)，能成功吗？为什么？

Part 3：数学最低限度

不追求数学严谨，只记住结论和直觉。后面遇到具体公式时再回来查。

线性代数

神经网络的所有计算本质上是矩阵乘法

import torch

# 矩阵乘法：神经网络层的本质
x = torch.randn(1, 512, 768)   # 输入: (batch, seq_len, dim)
W = torch.randn(768, 768)      # 权重矩阵
y = x @ W.T                    # 线性变换: y = xW^T
print(y.shape)                  # (1, 512, 768)

# 点积：Attention 中衡量"相关性"
q = torch.randn(1, 512, 64)    # Query 向量
k = torch.randn(1, 512, 64)    # Key 向量
score = q @ k.transpose(-2,-1)  # 点积 → 注意力分数
print(score.shape)               # (1, 512, 512)

• 向量：一列数字，如 [0.1, 0.5, -0.3]，可以理解为"方向"
• 矩阵乘法：nn.Linear(768, 768) 本质就是 y = xW^T + b
• 点积：衡量两个向量的"相似度"。Attention 中 Q·K 就是用点积算相关性

概率与 Softmax

语言模型的核心任务：预测下一个 token 的概率

import torch

# Softmax：把任意数字变成概率分布（加起来等于 1）
logits = torch.tensor([2.0, 1.0, 0.0])
probs = torch.softmax(logits, dim=-1)
print(probs)  # tensor([0.6654, 0.2447, 0.0900])
# 最大的 logits 对应最大的概率

# MiniMind 生成时：从概率分布中采样下一个 token
next_token = torch.multinomial(probs, num_samples=1)

• 概率分布：所有可能性之和 = 1，softmax 保证这一点
• 条件概率：语言模型就是 P(下一个词 | 前面的词)
• 训练目标：让正确 token 的概率尽可能大

交叉熵与损失

衡量预测与真实差距的标尺——Loss 越小，模型越准

• 交叉熵：越小 = 预测越准。公式 -Σ y_i × log(p_i)。MiniMind 中 CrossEntropyLoss 就是它
• KL 散度：衡量两个分布的差异。蒸馏和 RL 中用来约束模型不要跑偏

梯度与链式法则

训练能work的数学原理——你不需要手算，但需要直觉

import torch

# 梯度：告诉参数"往哪个方向调、调多少"
w = torch.randn(3, requires_grad=True)
loss = (w ** 2).sum()     # 前向：计算损失
loss.backward()            # 反向：自动求梯度
print(w.grad)              # tensor([2w1, 2w2, 2w3]) → ∂loss/∂w

# 链式法则：复合函数的导数 = 各层导数相乘
# 这就是"反向传播"的数学基础
# PyTorch 的 autograd 自动帮你算，你不需要手写

• 导数：函数在某点的变化率。正 = 上升，负 = 下降
• 梯度：多维导数，指向函数值增长最快的方向。训练沿梯度的反方向走
• 链式法则：复合函数的导数 = 各层导数相乘。反向传播就是链式法则的自动执行

学完自测

推荐学习资源