从 MicroGPT 最小实现迁移到 PyTorch（Tensor 版）实战指南

如果你已经读过前一篇“MicroGPT 最小实现”，你大概率会遇到下一个问题：

我看懂了原理，但怎么把这套“教学代码”迁移成 PyTorch 的工程写法？

这篇文章的目标就是解决这个问题。

我们不会一上来就“全部重写”，而是采用可验证的分阶段迁移，每一步都保证“能跑、可对齐、好定位错误”。

---

1. 先明确：我们要迁移什么，不迁移什么

典型最小版（标量计算图）通常有这些特征：

• 参数是自定义 Value 节点；
• 前向是 Python 循环 + 标量运算；
• 反向由手写 backward() 完成；
• 优化器是简单 SGD。

迁移到 PyTorch 后，我们的目标是：

• 参数变成 nn.Parameter；
• 运算全部张量化（batch 并行）；
• 用 autograd 自动求导；
• 保持“任务定义不变”（仍是 next token prediction）。

不变的是建模思想，变化的是实现载体。

---

2. 第 0 步：冻结行为，先写“对齐基线”

迁移前最容易踩坑的点不是 API，而是“你不知道哪里和原版不一致”。

建议先固定下面四件事：

1. 固定随机种子；
2. 固定词表与编码映射；
3. 固定一批训练样本（例如前 64 条）；
4. 记录基线指标（初始 loss、100 step 后 loss、生成样例）。

有了基线，你才能判断迁移后是“实现差异”还是“正常随机波动”。

---

3. 第 1 步：先只迁移参数容器，不迁移训练循环

第一阶段建议最小改动：

• 把参数放进 nn.Module；
• 前向逻辑暂时仍按原来写法；
• loss 先跑通并能 backward()。

示意代码：

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

class TinyCharLM(nn.Module):
    def __init__(self, vocab_size, n_embd, block_size):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        tok = self.token_emb(idx)                              # [B, T, C]
        pos = self.pos_emb(torch.arange(T, device=idx.device)) # [T, C]
        x = tok + pos
        logits = self.lm_head(x)                               # [B, T, V]

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(B * T, -1),
                targets.view(B * T)
            )
        return logits, loss

这一步结束标准：

• loss.backward() 可以正常执行；
• 参数的 .grad 非空；
• 一个 step 后 loss 有下降趋势。

---

4. 第 2 步：把“单样本循环”改成 batch 张量训练

很多最小实现是逐样本训练，这在教学上直观，但在工程上效率很低。

你需要把数据组织成：

• x: [B, T]，输入 token id；
• y: [B, T]，右移一位的目标。

核心收益有三个：

• GPU/向量化吞吐明显提升；
• 梯度估计更稳定；
• 训练循环更贴近真实项目。

常见错误：

• targets 展平维度写错；
• 忘了把 idx、targets 搬到同一个 device；
• 仍在 Python for-loop 内部重复做可向量化操作。

---

5. 第 3 步：补齐 Transformer 模块（注意力 + FFN + 残差）

如果你的最小版还没有完整 block，可以按下面顺序渐进增加：

1. 单头 self-attention；
2. 因果掩码（不能看未来）；
3. 多头并行；
4. 前馈网络（FFN）；
5. 残差连接与归一化。

建议每加一层结构都做一次“小回归”：

• 输入固定 batch；
• 打印 logits 的统计量（mean/std）；
• 检查 loss 没有 NaN；
• 跑 20~50 step 看是否持续下降。

这比“一次性全部合并后再排错”省很多时间。

---

6. 第 4 步：训练工程化（优化器、调度、混合精度）

当模型跑通后，再做工程增强：

• torch.optim.AdamW 替代纯 SGD；
• 学习率 warmup + cosine decay；
• 梯度裁剪（例如 clip_grad_norm_）；
• 混合精度（torch.cuda.amp）。

建议的最小训练骨架：

model = TinyCharLM(vocab_size, n_embd=256, block_size=128).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)

for step in range(max_steps):
    x, y = get_batch()  # [B, T], [B, T]
    x, y = x.to(device), y.to(device)

    optimizer.zero_grad(set_to_none=True)
    _, loss = model(x, y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

---

7. 如何验证“迁移成功”而不是“刚好能跑”

推荐用三层验证标准：

7.1 数值级验证

• 固定种子后，初始 loss 在同一数量级；
• 前 100 step 的下降趋势一致；
• 不出现 NaN/Inf。

7.2 功能级验证

• 可以稳定采样生成文本；
• temperature/top-k 等采样参数生效；
• 保存/加载 checkpoint 后生成行为一致。

7.3 工程级验证

• 吞吐（tokens/s）明显高于标量版本；
• 显存占用可控；
• 训练脚本可重复运行。

---

8. 常见迁移坑位清单（建议直接收藏）

1. view/reshape 误用：非连续内存时 view 可能报错，先 contiguous() 或用 reshape。
2. mask 设备不一致：attention mask 在 CPU，logits 在 GPU，直接崩。
3. logits 与 targets 维度错配：cross_entropy 期待 [N, C] + [N]。
4. 忘记 model.train() / model.eval()：dropout、norm 行为会不同。
5. 采样时没关梯度：生成阶段忘记 torch.no_grad()，显存飙升。

---

9. 一条推荐迁移路线图

如果你希望最稳妥地完成迁移，可以按这条路线执行：

• Day 1：参数容器迁移 + autograd 跑通；
• Day 2：batch 化 + loss 对齐；
• Day 3：attention block 完整化；
• Day 4：训练工程化 + checkpoint；
• Day 5：采样策略与评估脚本。

核心原则只有一句话：

每次只改一层抽象，每次都要可验证。

---

10. 小结

从最小版迁移到 PyTorch，不是“推翻重写”，而是把同一套思想放到更高效的执行框架里。

你真正应该保留的是：

• 对 next token prediction 的理解；
• 对 loss 与梯度流的感知；
• 对训练行为可解释、可验证的习惯。

当你能稳定完成这次迁移，后续再上中文语料、分词器、可视化分析，就会顺畅很多。

下一篇我们就讲：如何把字符级模型改成中文分词版本。