1. 这个项目到底在解决什么问题
第一次看到 "ai-engineering-from-scratch" 这个标题,我脑子里蹦出来的第一个念头是:又是一个教人调包的教程?但仔细琢磨了一下 "from scratch" 这四个字,我意识到它想做的事情可能完全不一样。市面上讲 AI 工程化的内容,绝大多数都是从pip install transformers或者import openai开始的,你跟着敲一遍,模型跑起来了,loss 也降下去了,但真要问你一句“这个张量在显存里到底是怎么排布的”“反向传播的时候梯度是怎么一层层传回去的”,很多人就卡壳了。这个项目标题里的 "from scratch",我理解它的核心诉求就是:把 AI 工程化这条链路上所有被封装起来的黑盒,一层层拆开给你看。
它要解决的问题很具体。现在一个刚入行的算法工程师,日常工作大概率是:拉一个开源仓库,改改配置文件,跑一下训练脚本,然后调 API 做推理。这套流程能出活,但一旦遇到显存溢出、梯度爆炸、推理延迟高、模型精度不达标这些真实问题,没有底层认知就完全无从下手。ai-engineering-from-scratch这类项目的价值,就是帮你建立从数学原理到工程实现的完整映射——你知道矩阵乘法在 GPU 上是怎么被切成 tile 的,你知道 attention 的 QKV 是怎么算出来的,你知道一个 tokenizer 的 BPE 合并规则是怎么一步步构建词表的。这些知识不是学术上的自娱自乐,而是你排查线上问题、做性能优化、甚至自己设计新架构时的底气。
适合谁来参考?我梳理了三类人。第一类是转行进入 AI 领域的工程师,你可能后端、前端或者数据方向做了几年,想切到 AI 工程,但发现面试问的全是底层原理,光会调库根本过不了。第二类是在校学生或者刚毕业的新人,学校课程讲了理论但没讲工程,你懂反向传播的公式但没自己手写过,这个项目能帮你把理论和代码之间的鸿沟填上。第三类是工作了两三年想突破瓶颈的算法工程师,你日常 CRUD 模型已经熟练了,但想往架构或者性能优化方向走,需要重新把底层捡起来。这三类人的共同点是:不满足于“能用”,想搞清楚“为什么能用”以及“怎么让它更好用”。
我自己的经历是,早年做推荐系统的时候,模型训练经常 OOM,当时只会无脑调小 batch size,后来逼着自己去读 PyTorch 的显存管理机制,去手写了一遍简单的反向传播,才真正理解 activation 占显存是怎么回事、checkpointing 为什么能省显存。那种“原来如此”的瞬间,是调包永远给不了的。所以看到这个标题,我特别有共鸣,也特别想把这个项目可能涉及的核心内容、实操路径、以及我踩过的坑,系统地梳理一遍。
2. 从零构建 AI 工程能力的整体设计思路
2.1 为什么“从零”比“调包”更重要
很多人会问,现在工具链这么成熟,为什么还要从零写?这不是重复造轮子吗?我的回答是:造轮子不是为了用,是为了懂。你从零实现一遍矩阵乘法,不是为了替代 cuBLAS,而是为了理解为什么你的模型在 batch size 增大时显存会非线性增长;你从零写一遍 attention,不是为了替代 FlashAttention,而是为了理解为什么 FlashAttention 能省显存、它的 IO 瓶颈在哪里。这种底层认知,决定了你在面对新问题时的解决上限。
从工程角度看,“from scratch” 的路径设计通常遵循一个自底向上的逻辑。最底层是数学和数值计算——向量、矩阵、张量、自动微分。往上一层是神经网络的基本组件——线性层、激活函数、归一化、注意力机制。再往上是模型架构——Transformer、CNN、RNN 的变体。最上层才是训练工程和推理部署——优化器、学习率调度、分布式训练、量化、服务化。这个分层不是随便定的,它对应的是依赖关系:你不理解自动微分,就没法理解为什么 PyTorch 的backward()能自动算梯度;你不理解 attention 的计算复杂度,就没法理解为什么长文本推理会那么慢。
2.2 技术选型背后的考量
如果我来设计这个项目的技术栈,我会做几个关键取舍。语言层面,Python 是必须的,因为 AI 生态几乎全在 Python 上,但纯 Python 做数值计算太慢,所以核心计算部分要用 NumPy 打底,再逐步过渡到 PyTorch。为什么不直接上 PyTorch?因为 PyTorch 的 tensor 已经封装了太多东西,你torch.matmul(a, b)一行就完事了,但你看不到底层的循环、内存分配、并行策略。先用 NumPy 手写一遍,你才能体会到“哦,原来矩阵乘法就是三层循环”,然后再看 PyTorch 怎么用 CUDA 把它加速几千倍,这个对比带来的认知冲击是巨大的。
框架层面,我倾向于从纯 NumPy 开始,然后引入 PyTorch 的 autograd 机制,最后再对比 JAX 或者 TinyGrad 这类更底层的框架。这个顺序的逻辑是:先理解“计算是什么”,再理解“梯度怎么自动算”,最后理解“怎么在不同硬件上高效执行”。每一步都在前一步的基础上增加抽象层,但因为你已经知道了下面一层长什么样,所以上面的抽象对你来说不再是黑盒。
硬件层面,CPU 和 GPU 的对比是绕不开的。你手写的 NumPy 矩阵乘法在 CPU 上跑,和 PyTorch 在 GPU 上跑,性能差距可能是几百倍。这个差距来自哪里?来自 SIMD 指令、来自内存带宽、来自并行度。理解这些,你才能在做推理优化时知道该往哪个方向使劲——是减少计算量,还是减少内存访问,还是提高并行度。
2.3 学习路径的阶段性目标
我把这个项目的学习路径分成四个阶段,每个阶段都有明确的产出物,这样你学起来不会迷茫。
第一阶段:数值计算基础。目标是能手写一个支持自动微分的张量库。产出物是一个几百行的micrograd或者tinytensor,支持加减乘除、矩阵乘法、ReLU、Sigmoid 这些基本操作,并且能对一个简单的两层神经网络做反向传播。这个阶段的核心是理解计算图和链式法则。
第二阶段:神经网络组件。目标是用自己写的张量库实现一个完整的 Transformer。产出物是一个能在小数据集上训练并收敛的迷你 GPT。这个阶段你会遇到很多工程细节:LayerNorm 的 epsilon 放哪里、位置编码用正弦还是可学习、attention mask 怎么加、dropout 在训练和推理时的区别。
第三阶段:训练工程。目标是把训练效率提上去。产出物是一个支持混合精度、梯度累积、梯度裁剪、学习率 warmup 的训练脚本。这个阶段你会理解为什么 batch size 不能无限大、为什么需要 warmup、为什么混合精度能省显存但可能掉精度。
第四阶段:推理与部署。目标是把训练好的模型高效地跑起来。产出物是一个支持 KV Cache、量化、批处理的推理服务。这个阶段你会理解为什么推理和训练对硬件的需求不一样、为什么量化能加速、为什么批处理能提高吞吐。
这四个阶段走下来,你对 AI 工程化的理解就不再是“我会用某个框架”,而是“我知道这个框架为什么这么设计,以及我能在它不够用的时候自己改”。
3. 核心细节拆解与实操要点
3.1 自动微分:从计算图到反向传播
自动微分是整个深度学习框架的基石,但很多人对它的理解停留在“PyTorch 会自动帮我算梯度”。如果你从零实现一遍,你会发现它的核心思想其实非常优雅:前向传播时构建计算图,反向传播时沿着图反向应用链式法则。
具体实现上,每个张量需要记录三个东西:它的数据(data)、它的梯度(grad)、以及它是怎么被算出来的(grad_fn 或者 parents)。当你做c = a * b时,你不仅得到了 c 的值,还记录了 c 依赖于 a 和 b,以及这个操作的局部导数是什么。反向传播时,你从 loss 开始,沿着计算图反向遍历,用链式法则把梯度一层层传回去。
这里有几个实操要点。第一,计算图要在每次前向传播时重建,因为每次的输入可能不同,图的结构也可能不同(比如有控制流)。这就是为什么 PyTorch 是动态图,而早期的 TensorFlow 是静态图。第二,梯度要累加而不是覆盖,因为一个张量可能被多个下游节点使用,它的梯度是所有下游梯度之和。第三,要支持梯度清零,因为 PyTorch 默认会累加梯度,如果你不清零,梯度会越积越大。
我手写自动微分时踩过的一个坑是:忘记处理广播(broadcasting)。比如a的形状是(3, 1),b的形状是(1, 4),a + b的结果是(3, 4)。反向传播时,a的梯度应该是(3, 1),但如果你直接按(3, 4)传回去,形状就对不上。正确的做法是对广播的维度做 sum 归约。这个细节在调包时完全感知不到,但自己实现时如果漏了,梯度就会出错,而且很难 debug。
# 广播反向传播的简化示意 def add_backward(grad_output, a_shape, b_shape): # 对 a 的梯度:如果 a 被广播了,需要对广播维度求和 grad_a = grad_output if a_shape != grad_output.shape: # 找到被广播的维度并求和 for i, dim in enumerate(a_shape): if dim == 1 and grad_output.shape[i] > 1: grad_a = grad_a.sum(dim=i, keepdim=True) # b 同理 grad_b = grad_output if b_shape != grad_output.shape: for i, dim in enumerate(b_shape): if dim == 1 and grad_output.shape[i] > 1: grad_b = grad_b.sum(dim=i, keepdim=True) return grad_a, grad_b注意:广播的反向传播是新手手写自动微分时最容易出错的地方,建议写完一个操作就立刻用数值梯度校验(gradient check),确保解析梯度和数值梯度在误差范围内一致。
3.2 注意力机制:从公式到高效实现
Attention 是 Transformer 的核心,它的公式看起来很简单:Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V。但工程实现上有非常多细节。
首先是缩放因子。为什么要除以sqrt(d_k)?因为当d_k很大时,QK^T的点积结果会很大,softmax 之后会变得非常尖锐,梯度会趋近于零,导致训练困难。除以sqrt(d_k)可以把方差拉回到 1 附近,保持梯度的稳定性。这个细节如果你只是调包,可能永远不会注意到,但自己实现时如果不加,训练很容易发散。
其次是mask 的处理。在 decoder 中,你需要 causal mask,让每个位置只能看到它之前的位置。实现上是在 softmax 之前把未来位置的 logits 设成负无穷,这样 softmax 之后这些位置的权重就是 0。这里有个坑:负无穷不能直接用-inf,因为-inf参与计算会产生 NaN。实践中通常用一个很大的负数,比如-1e9,来代替。
然后是KV Cache。推理时,如果你每次都重新计算所有位置的 K 和 V,计算量会随序列长度平方增长。KV Cache 的思路是把之前算过的 K 和 V 存下来,每次只算新 token 的 K 和 V,然后和缓存拼接。这样推理的计算量就从 O(n^2) 降到了 O(n)。但 KV Cache 会占显存,序列越长、batch 越大,占的显存越多。这就是为什么长文本推理对显存要求很高。
# 带 KV Cache 的 attention 推理简化示意 def attention_with_cache(q, k_cache, v_cache, new_k, new_v): # 拼接缓存 k = torch.cat([k_cache, new_k], dim=-2) v = torch.cat([v_cache, new_v], dim=-2) # 计算 attention attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1)) attn = torch.softmax(attn, dim=-1) output = torch.matmul(attn, v) return output, k, v提示:KV Cache 的显存占用可以用这个公式估算:
2 * batch_size * num_heads * seq_len * head_dim * num_layers * dtype_size。以 LLaMA-7B 为例,fp16 下每个 token 的 KV Cache 大约是 0.5MB,4096 个 token 就是 2GB。这个数字在做长文本推理规划时非常关键。
3.3 训练工程:那些调包时感知不到的细节
训练一个模型,调包时你只需要写model.train()和loss.backward(),但背后有一大堆工程细节决定了你能不能训得起来、训得好。
混合精度训练是现在标配。它的核心思想是:前向和反向用 fp16 算,因为 fp16 计算快、省显存;但参数更新用 fp32,因为 fp16 的精度不够,直接更新会导致参数无法收敛。PyTorch 的autocast和GradScaler帮你自动处理了这些,但你需要理解为什么需要 loss scaling——因为 fp16 的表示范围窄,小梯度会下溢成 0,所以要把 loss 放大,让梯度落在 fp16 能表示的范围内,更新前再缩回去。
梯度累积是另一个常用技巧。当显存不够跑大 batch 时,你可以把一个大 batch 拆成几个小 batch,分别前向反向,把梯度累加起来,最后一起更新。这样等效于大 batch 训练,但显存占用只有小 batch 的水平。这里有个坑:如果你用了 BatchNorm,梯度累积会改变统计量的计算方式,因为 BatchNorm 是在每个小 batch 上算均值和方差的。所以用梯度累积时,要么改用 LayerNorm,要么同步 BN 的统计量。
学习率调度对训练稳定性影响巨大。最常用的是 warmup + cosine decay:前几百步线性增加学习率,避免一开始梯度太大把模型带偏;然后余弦衰减到接近零,让模型在后期精细收敛。warmup 的步数通常设为总步数的 1% 到 5%,具体取决于模型大小和数据量。
| 训练技巧 | 解决的问题 | 代价 | 适用场景 |
|---|---|---|---|
| 混合精度 | 显存不足、训练慢 | 可能掉精度,需要 loss scaling | 几乎所有现代模型训练 |
| 梯度累积 | 显存不足以跑大 batch | 训练速度略慢 | 大模型、小显存 |
| 梯度裁剪 | 梯度爆炸 | 可能减慢收敛 | RNN、深层 Transformer |
| 学习率 warmup | 初期训练不稳定 | 前期收敛慢 | 大 batch、大模型 |
| KV Cache | 推理重复计算 | 占显存 | 自回归生成 |
3.4 推理优化:从能跑到跑得快
训练完模型只是第一步,怎么让它在线上高效运行是另一个大课题。推理优化有几个方向:减少计算量、减少内存访问、提高并行度。
量化是最直接的减少计算量和内存访问的方法。把 fp16 的权重和激活量化成 int8,理论上计算量减少一半,内存占用减少一半。但量化会引入误差,需要校准(calibration)来找到合适的缩放因子。实践中,weight-only 量化(只量化权重,激活保持 fp16)实现简单、精度损失小,适合大多数场景;weight+activation 量化更激进,需要更复杂的校准和可能的微调。
批处理是提高吞吐的关键。GPU 的算力很强,但如果你每次只推理一个样本,GPU 的利用率可能只有 10%。把多个请求攒成一个 batch 一起推理,吞吐能提升几倍到几十倍。但批处理会增加延迟,因为你要等请求攒够。所以线上服务通常会在吞吐和延迟之间做权衡,设置一个最大 batch size 和最大等待时间。
连续批处理(continuous batching)是现在大模型推理的标配。传统的静态批处理要等整个 batch 的所有请求都生成完才能处理下一批,但不同请求的生成长度不一样,短的请求早就结束了,GPU 却在等长的请求。连续批处理让每个请求独立生成,一个请求结束就立刻把新的请求塞进来,GPU 始终处于满载状态。vLLM 和 TensorRT-LLM 都实现了这个机制。
4. 完整实操流程:手写一个迷你 GPT
4.1 环境准备与依赖安装
这个实操的目标是用纯 PyTorch 手写一个迷你 GPT,在字符级语言模型任务上训练,让它能生成一段像样的文本。环境准备很简单,Python 3.10 以上,PyTorch 2.0 以上,再加个 NumPy 做数据处理。
pip install torch numpy matplotlib数据集我用的是莎士比亚全集,字符级建模,词表大概 65 个字符。这个规模的数据集在 CPU 上也能训,但用 GPU 会快很多。如果你没有 GPU,把模型调小一点(4 层、4 头、256 维),在 CPU 上跑几个小时也能看到效果。
import torch import torch.nn as nn import torch.nn.functional as F # 超参数配置 batch_size = 64 block_size = 128 # 上下文长度 n_embd = 384 # 嵌入维度 n_head = 6 # 注意力头数 n_layer = 6 # Transformer 层数 dropout = 0.2 learning_rate = 3e-4 max_iters = 5000注意:
block_size决定了模型能看到的最大上下文长度。字符级模型通常 128 到 256 就够了,但如果你做的是 token 级模型,这个值要设得更大,比如 2048 或 4096。n_embd必须能被n_head整除,因为每个头分到的维度是n_embd // n_head。
4.2 数据准备与字符级 Tokenizer
字符级 tokenizer 是最简单的 tokenizer:把每个字符映射成一个整数。虽然简单,但它能让你专注于模型本身,而不是 tokenizer 的细节。
# 读取数据 with open('input.txt', 'r', encoding='utf-8') as f: text = f.read() # 构建词表 chars = sorted(list(set(text))) vocab_size = len(chars) stoi = {ch: i for i, ch in enumerate(chars)} itos = {i: ch for i, ch in enumerate(chars)} # 编码和解码函数 encode = lambda s: [stoi[c] for c in s] decode = lambda l: ''.join([itos[i] for i in l]) # 划分训练集和验证集 data = torch.tensor(encode(text), dtype=torch.long) n = int(0.9 * len(data)) train_data = data[:n] val_data = data[n:]数据加载的核心是随机采样固定长度的序列。每次从训练集中随机选一个起始位置,取block_size + 1个字符,前block_size个作为输入,后block_size个作为目标。这样每个样本的输入和输出是错开一位的,模型学的是“给定前 n 个字符,预测第 n+1 个字符”。
def get_batch(split): data = train_data if split == 'train' else val_data ix = torch.randint(len(data) - block_size, (batch_size,)) x = torch.stack([data[i:i+block_size] for i in ix]) y = torch.stack([data[i+1:i+block_size+1] for i in ix]) return x, y4.3 模型定义:从注意力到 GPT
模型定义是整个实操的核心。我把它拆成几个模块:单头注意力、多头注意力、前馈网络、Transformer Block、以及最终的 GPT 模型。
单头注意力的实现要注意几个点:Q、K、V 的线性变换、缩放点积、causal mask、softmax、dropout。
class Head(nn.Module): def __init__(self, head_size): super().__init__() self.key = nn.Linear(n_embd, head_size, bias=False) self.query = nn.Linear(n_embd, head_size, bias=False) self.value = nn.Linear(n_embd, head_size, bias=False) self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size))) self.dropout = nn.Dropout(dropout) def forward(self, x): B, T, C = x.shape k = self.key(x) # (B, T, head_size) q = self.query(x) # (B, T, head_size) # 计算注意力分数 wei = q @ k.transpose(-2, -1) * (k.shape[-1] ** -0.5) # 应用 causal mask wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) wei = F.softmax(wei, dim=-1) wei = self.dropout(wei) # 加权聚合 v = self.value(x) out = wei @ v return out多头注意力就是把多个单头注意力的输出拼接起来,再过一个线性层。这里的关键是每个头关注不同的子空间,有的头可能关注语法结构,有的头可能关注语义关联。拼接后再投影,让模型能融合不同头的信息。
class MultiHeadAttention(nn.Module): def __init__(self, num_heads, head_size): super().__init__() self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)]) self.proj = nn.Linear(n_embd, n_embd) self.dropout = nn.Dropout(dropout) def forward(self, x): out = torch.cat([h(x) for h in self.heads], dim=-1) out = self.dropout(self.proj(out)) return out前馈网络是 Transformer 里另一个核心组件。它是一个两层的 MLP,中间层的维度通常是嵌入维度的 4 倍。这个 4 倍不是随便定的,而是经验上能在表达能力和计算量之间取得平衡。
class FeedForward(nn.Module): def __init__(self, n_embd): super().__init__() self.net = nn.Sequential( nn.Linear(n_embd, 4 * n_embd), nn.ReLU(), nn.Linear(4 * n_embd, n_embd), nn.Dropout(dropout), ) def forward(self, x): return self.net(x)Transformer Block 把注意力和前馈网络串起来,每个子层后面加残差连接和 LayerNorm。这里有个细节:Pre-LN 还是 Post-LN。原始 Transformer 用的是 Post-LN(先残差再 Norm),但后来大家发现 Pre-LN(先 Norm 再残差)训练更稳定,不需要 warmup 也能训。现在主流实现基本都是 Pre-LN。
class Block(nn.Module): def __init__(self, n_embd, n_head): super().__init__() head_size = n_embd // n_head self.sa = MultiHeadAttention(n_head, head_size) self.ffwd = FeedForward(n_embd) self.ln1 = nn.LayerNorm(n_embd) self.ln2 = nn.LayerNorm(n_embd) def forward(self, x): x = x + self.sa(self.ln1(x)) x = x + self.ffwd(self.ln2(x)) return x最后是 GPT 模型本身。它包含 token 嵌入、位置嵌入、多个 Transformer Block、最后的 LayerNorm 和输出投影。
class MiniGPT(nn.Module): def __init__(self): super().__init__() self.token_embedding = nn.Embedding(vocab_size, n_embd) self.position_embedding = nn.Embedding(block_size, n_embd) self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)]) self.ln_f = nn.LayerNorm(n_embd) self.lm_head = nn.Linear(n_embd, vocab_size) 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.lm_head(x) if targets is None: loss = None else: B, T, C = logits.shape logits = logits.view(B*T, C) targets = targets.view(B*T) loss = F.cross_entropy(logits, targets) return logits, loss def generate(self, idx, max_new_tokens): for _ in range(max_new_tokens): idx_cond = idx[:, -block_size:] logits, _ = self(idx_cond) logits = logits[:, -1, :] probs = F.softmax(logits, dim=-1) idx_next = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, idx_next), dim=1) return idx4.4 训练循环与参数调优
训练循环看起来简单,但里面有很多可以调的地方。优化器我用 AdamW,学习率 3e-4,这是 Transformer 类模型的经验值。损失函数是交叉熵,因为这是语言建模任务。
model = MiniGPT().to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) for iter in range(max_iters): xb, yb = get_batch('train') xb, yb = xb.to(device), yb.to(device) logits, loss = model(xb, yb) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() if iter % 500 == 0: # 评估验证集损失 model.eval() with torch.no_grad(): xv, yv = get_batch('val') xv, yv = xv.to(device), yv.to(device) _, val_loss = model(xv, yv) model.train() print(f"Step {iter}: train loss {loss.item():.4f}, val loss {val_loss.item():.4f}")参数调优方面,我总结了几条经验。学习率是最重要的超参数,3e-4 是个安全的起点,如果 loss 震荡就调小,如果收敛太慢就调大。batch size影响梯度估计的噪声,太小噪声大,太大泛化可能变差,64 到 256 是比较常用的范围。dropout防止过拟合,小数据集上 0.1 到 0.3 比较合适,大数据集可以设 0。层数和头数决定模型容量,但也要考虑计算资源,6 层 6 头是个不错的平衡点。
实操心得:训练初期 loss 会快速下降,然后进入平台期。如果你发现 loss 一直不降,先检查数据有没有问题(比如标签错位),再检查学习率是不是太大或太小。如果 loss 降但验证 loss 不降,说明过拟合了,加 dropout 或者加数据。
4.5 推理与文本生成
训练完之后,用generate方法生成文本。这里用的是多项式采样,而不是贪心,因为贪心容易生成重复的文本。多项式采样根据概率分布随机选下一个 token,多样性更好。
# 生成文本 context = torch.zeros((1, 1), dtype=torch.long, device=device) generated = model.generate(context, max_new_tokens=500) print(decode(generated[0].tolist()))如果你想让生成更可控,可以加 temperature 参数。temperature 越高,分布越平坦,生成越随机;temperature 越低,分布越尖锐,生成越确定。通常 0.7 到 1.0 是比较好的范围。
def generate_with_temperature(model, idx, max_new_tokens, temperature=1.0): for _ in range(max_new_tokens): idx_cond = idx[:, -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 idx5. 常见问题与排查技巧实录
5.1 训练不收敛或 loss 震荡
这是最常见的问题。排查思路按优先级来:先看数据,再看学习率,最后看模型结构。
数据问题包括:标签和输入没有对齐(比如忘了错位一位)、数据里有大量异常值、词表构建有问题。检查方法是手动打印几个样本,看看输入和输出是不是合理的。
学习率问题包括:太大导致震荡,太小导致收敛慢。可以画 loss 曲线,如果 loss 上下跳动很大,就是学习率太大;如果 loss 几乎不降,就是学习率太小。解决方案是加 warmup 或者用学习率搜索。
模型结构问题包括:LayerNorm 位置不对、残差连接忘了加、初始化有问题。Transformer 类模型对初始化比较敏感,PyTorch 默认的初始化通常够用,但如果你自己改了,要确保方差是合理的。
| 现象 | 可能原因 | 排查方法 | 解决方案 |
|---|---|---|---|
| loss 不降 | 学习率太小、数据有问题 | 打印梯度范数、检查数据 | 调大学习率、修复数据 |
| loss 震荡 | 学习率太大、batch 太小 | 画 loss 曲线 | 调小学习率、增大 batch |
| loss 变 NaN | 梯度爆炸、除零 | 检查梯度、检查分母 | 梯度裁剪、加 epsilon |
| 验证 loss 上升 | 过拟合 | 对比训练和验证曲线 | 加 dropout、加数据、早停 |
| 生成重复文本 | 采样策略问题 | 检查 temperature | 调高 temperature、用 top-k |
5.2 显存溢出(OOM)的排查与解决
OOM 是训练大模型时的家常便饭。排查思路是:先算理论显存,再看实际占用,最后找优化空间。
理论显存包括四部分:模型参数、梯度、优化器状态、激活值。以 fp32 为例,一个 1 亿参数的模型,参数占 400MB,梯度占 400MB,Adam 优化器状态占 800MB(一阶矩和二阶矩),加起来就是 1.6GB。激活值取决于 batch size 和序列长度,通常是最大的一块。
优化手段按性价比排序:混合精度能省一半显存,几乎无损;梯度累积让你用小 batch 模拟大 batch,不省显存但能跑大 batch;梯度检查点用计算换显存,能省很多但训练变慢;模型并行把模型切到多卡上,适合超大模型。
注意:PyTorch 的显存分配器会缓存已分配的显存,所以
nvidia-smi显示的显存占用可能比实际需要的大。如果你想看真实占用,用torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()。
5.3 推理速度慢的优化路径
推理慢通常有三个原因:计算量大、内存访问瓶颈、并行度不够。
计算量大:模型太大或者序列太长。优化方法是量化、剪枝、蒸馏,或者用更高效的注意力实现(比如 FlashAttention)。
内存访问瓶颈:GPU 算得快但等数据。优化方法是算子融合、减少中间张量、用更高效的内存布局。
并行度不够:batch 太小或者 GPU 利用率低。优化方法是增大 batch、用连续批处理、用多卡推理。
我实测下来,对于 7B 级别的模型,用 vLLM 做推理比用 HuggingFace 的generate快 5 到 10 倍,主要就是连续批处理和 PagedAttention 的功劳。如果你要上生产,这些推理框架基本是必选的。
5.4 手写代码时的常见 bug 清单
自己从零写代码,bug 比调包多得多。我整理了一份常见 bug 清单,供你对照排查。
- 形状不匹配:矩阵乘法前先打印形状,确认维度对齐。
- 忘记
zero_grad:梯度累加导致更新错误。 - 忘记
model.train()/model.eval():dropout 和 batchnorm 行为不对。 - mask 用
-inf导致 NaN:用大负数代替。 - 位置编码越界:生成时序列长度超过
block_size。 - softmax 维度搞错:应该在最后一维做 softmax。
- 学习率没 warmup:大模型初期容易发散。
- 数据没 shuffle:模型学到顺序偏差。
- 验证集没
no_grad:显存爆炸。 - 保存模型没保存优化器状态:断点续训出问题。
实操心得:手写代码时,每写一个模块就单独测试一下。比如写完 attention,用一个随机输入跑一遍,检查输出形状对不对、数值范围合不合理。不要等整个模型写完再 debug,那样定位问题会非常痛苦。
6. 从项目到能力:我的个人体会
这个项目最值钱的地方,不是让你学会写一个 GPT,而是让你建立一种**“我能拆开看”的自信**。我以前遇到模型问题,第一反应是搜 issue、问别人、试各种 trick,但自从手写过一遍核心组件之后,我遇到问题的第一反应变成了“让我看看底层发生了什么”。这种思维方式的转变,比任何具体的知识点都重要。
另一个体会是,从零实现和调包不是对立的,而是互补的。你从零实现是为了理解,理解之后再用成熟的框架,你会用得更明白、更放心。你知道torch.nn.Linear背后做了什么,你知道F.cross_entropy的数值稳定性是怎么保证的,你知道AdamW和Adam的区别在哪里。这些认知让你在选型、调参、排查问题时,都有更清晰的判断。
最后分享一个小技巧:如果你觉得从零写一个完整的 GPT 太难,可以从更小的目标开始。先写一个线性回归,手写梯度下降;再写一个两层 MLP,手写反向传播;然后写一个单头 attention;最后拼成完整的 Transformer。每一步都跑通、都理解,比一口气写一个大项目然后 debug 到崩溃要高效得多。我自己就是这么一步步过来的,踩过的坑不少,但每填一个坑,能力就实打实长一分。