1. 从零手搓AI工程为什么我不建议你直接调包第一次看到ai-engineering-from-scratch这个项目名的时候我正坐在工位上啃一个调了三天的模型部署脚本。那会儿我的日常就是pip install transformers、from_pretrained、model.generate一套组合拳打得飞起但真要我解释清楚一个 token 是怎么从输入文本变成向量、再经过注意力层、最后吐出下一个词的我脑子里其实是一团浆糊。这个项目标题像一记闷棍敲在我后脑勺上——从零开始做AI工程不是从零调包是从零理解每一个矩阵乘法的意义。这个项目适合谁如果你已经会用 PyTorch 或 TensorFlow 搭个分类器但每次遇到维度不匹配就靠print(shape)硬试如果你能跑通 HuggingFace 的示例代码但换个模型结构就不知道从哪改起如果你面试时被问到“手写一个注意力机制”就手心冒汗——那这个项目就是给你准备的。它不教你调参技巧不教你刷榜策略它教你的是把AI工程拆解到最底层用最朴素的方式重新组装一遍。我花了大概三周时间利用下班和周末把这个项目从环境搭建一路推到能跑通一个简化版 Transformer 的完整训练流程。踩过的坑包括但不限于梯度爆炸到 loss 变成 nan、位置编码写错导致模型完全学不到顺序信息、学习率设太大直接把权重更新成乱码。这篇文章就是我这三周的血泪总结我会把每个环节的为什么这么做讲清楚把参数选择的计算过程摆出来把那些文档里不会写的避坑经验全倒出来。你不需要有博士学位但你需要有耐心——因为从零开始这件事本身就是一场和自己的耐心博弈。2. 项目整体设计与思路拆解2.1 为什么选择“从零实现”而不是“调包微调”市面上教AI的教程分两类一类是“三行代码调用预训练模型”另一类是“从数学公式推导反向传播”。前者让你快速出结果但永远是个黑盒操作员后者让你理解原理但离工程落地隔了十万八千里。ai-engineering-from-scratch走的是第三条路用工程化的方式重新实现核心组件但不过度追求数学严谨性。具体来说这个项目的设计哲学是每个模块先用最直观的方式实现一版能跑的然后再逐步优化到接近生产可用的状态。比如实现注意力机制第一版可能就是一个双重循环时间复杂度 O(n²) 但逻辑清晰第二版改成矩阵运算利用 GPU 并行第三版才加入多头、掩码、dropout 这些工程细节。这种渐进式实现的好处是你始终知道每一行代码在干什么而不是复制粘贴一堆自己都看不懂的 tensor 操作。我选择这个项目的一个核心原因是面试造火箭工作拧螺丝。但如果你连螺丝的螺纹方向都搞不清楚造火箭的图纸摆在你面前你也看不懂。从零实现一遍之后再看 HuggingFace 的源码你会发现那些nn.Module的嵌套关系、那些forward函数里的 shape 变换突然变得有迹可循了。2.2 技术栈选型为什么是 PyTorch 而不是 TensorFlow项目默认用 PyTorch这不是随便选的。从工程角度看PyTorch 的动态计算图对调试友好得多。你可以在forward函数里随便print中间变量的 shape可以在任意位置打断点可以动态改变网络结构。TensorFlow 1.x 的静态图时代调试一个 shape 错误可能要花半小时在tf.Print和 session.run 之间来回折腾。另一个原因是生态。HuggingFace 的transformers、datasets、accelerate全是 PyTorch 优先你想从零实现过渡到工业级工具链PyTorch 是更顺滑的路径。而且 PyTorch 的 API 设计更符合 Python 程序员的直觉——torch.nn.Linear就是一个可调用对象loss.backward()就是自动求导没有那么多tf.Session和placeholder的概念负担。提示如果你之前只用过 TensorFlow建议先花半天时间过一遍 PyTorch 的官方 60 分钟入门教程。重点理解autograd机制和nn.Module的继承方式这是后续所有实现的基础。2.3 项目模块划分与依赖关系整个项目我把它拆成了五个核心模块每个模块都可以独立运行和测试模块名称核心功能依赖关系预计耗时数据管道文本清洗、分词、词表构建、批次生成无1天基础层线性层、嵌入层、层归一化、激活函数数据管道2天注意力机制自注意力、多头注意力、掩码机制基础层3天模型组装Transformer Block、编码器、解码器注意力机制2天训练循环损失计算、反向传播、优化器、学习率调度模型组装2天这个划分方式的好处是可测试性。每完成一个模块你都可以写一个简单的单元测试来验证它的行为是否符合预期。比如基础层写完后你可以构造一个随机输入检查线性层的输出维度是否正确、层归一化的均值是否接近零。这种增量式验证能帮你快速定位问题——如果模型训练不收敛你可以从后往前逐个模块排查而不是面对一个几千行的脚本发呆。3. 核心细节解析与实操要点3.1 数据管道从原始文本到模型可吃的张量数据管道是整个项目的地基地基没打好后面全是空中楼阁。我见过太多人直接拿torchtext或者datasets库一把梭结果遇到自定义数据集就傻眼。从零实现一遍之后你会发现文本到张量的转换其实就三步分词、映射、批处理。分词这步英文可以用空格加标点分割中文就需要更复杂的策略。项目里为了简化先用字符级分词——每个字符就是一个 token。这样做的好处是词表小、不会遇到 OOV未登录词坏处是序列长度爆炸。一个 100 字的句子变成 100 个 token注意力机制的计算量直接起飞。但作为学习项目字符级分词能让你把注意力集中在模型结构本身而不是分词算法的细节上。词表构建就是统计所有出现过的字符给每个字符分配一个唯一整数 ID。这里有个坑一定要预留特殊 token 的位置。pad用于填充、unk用于未知字符、bos和eos用于序列起始和结束。我一开始忘了留pad结果批次里长短不一的序列没法对齐只能手动在 collate 函数里补零代码写得极其丑陋。批处理的核心是动态填充。假设一个批次里有三个句子长度分别是 5、8、3你需要把它们都填充到 8然后生成一个掩码矩阵告诉模型哪些位置是真实的、哪些是填充的。这个掩码矩阵在后续的注意力计算里会用到——填充位置的注意力分数要设成负无穷这样 softmax 之后这些位置的权重就是零。def collate_fn(batch): # batch 是 (input_ids, labels) 的列表 max_len max(len(x[0]) for x in batch) padded_inputs [] padded_labels [] masks [] for input_ids, labels in batch: pad_len max_len - len(input_ids) padded_inputs.append(input_ids [pad_token_id] * pad_len) padded_labels.append(labels [-100] * pad_len) # -100 在交叉熵里被忽略 masks.append([1] * len(input_ids) [0] * pad_len) return torch.tensor(padded_inputs), torch.tensor(padded_labels), torch.tensor(masks)注意-100这个值是 PyTorchCrossEntropyLoss默认忽略的标签值。如果你用别的损失函数需要查一下对应的 ignore_index 参数。3.2 基础层实现线性层、层归一化与残差连接线性层看起来简单——不就是y xW b吗但工程实现上有几个细节值得注意。第一是权重初始化。如果你直接用torch.randn初始化方差太大会导致梯度爆炸方差太小会导致梯度消失。PyTorch 的nn.Linear默认用 Kaiming 初始化具体公式是std sqrt(2 / fan_in)其中fan_in是输入维度。这个系数的选择是为了让前向传播时每一层的输出方差保持一致。第二是偏置项的处理。理论上线性层可以没有偏置因为后面的层归一化会重新中心化。但实践中保留偏置通常收敛更快因为模型多了一点自由度。我试过去掉偏置loss 下降速度明显变慢后来查资料发现是因为没有偏置时模型需要靠权重矩阵的特定行来补偿增加了学习难度。层归一化LayerNorm是 Transformer 的稳定器。它的作用是对每个样本的特征维度做归一化使其均值为零、方差为一。公式是(x - mean) / sqrt(var eps) * gamma beta其中gamma和beta是可学习的缩放和平移参数。这里的关键是在哪个维度上计算均值和方差。对于形状为(batch, seq_len, hidden_dim)的张量LayerNorm 是在hidden_dim维度上计算的也就是说每个位置的特征向量独立归一化。残差连接是深层网络的标配。output x sublayer(x)这个简单的加法解决了梯度消失问题让梯度可以直接通过跳跃连接回传。但有个细节残差连接要求输入输出维度一致。如果你在子层里改变了维度就需要在跳跃连接上加一个线性投影。我在实现时偷懒没加投影结果维度不匹配报错排查了半天才发现是这个问题。3.3 注意力机制从直觉到矩阵运算注意力机制的核心思想是每个位置根据与其他位置的相关性加权聚合信息。用生活类比来说你在读一句话的时候理解某个词的含义会参考上下文——注意力机制就是把这个过程数学化。最基础的自注意力计算分三步把输入映射成 Query、Key、Value 三个矩阵计算 Query 和 Key 的点积得到注意力分数用 softmax 归一化分数然后对 Value 加权求和用矩阵形式写出来就是Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。这里的sqrt(d_k)是缩放因子目的是防止点积结果太大导致 softmax 梯度消失。假设d_k 64点积的方差大约是 64标准差是 8softmax 之后最大的那个值会接近 1其他接近 0梯度几乎为零。除以 8 之后方差回到 1 左右softmax 的输出分布更平滑梯度更健康。多头注意力就是把上面的过程并行做多次每次用不同的 Q、K、V 投影矩阵最后把结果拼接起来再投影一次。这样做的好处是模型可以同时关注不同位置的不同特征——一个头关注语法关系另一个头关注语义相似度。头的数量通常是 8 或 16每个头的维度是hidden_dim / num_heads。class MultiHeadAttention(nn.Module): def __init__(self, hidden_dim, num_heads): super().__init__() self.num_heads num_heads self.head_dim hidden_dim // num_heads self.q_proj nn.Linear(hidden_dim, hidden_dim) self.k_proj nn.Linear(hidden_dim, hidden_dim) self.v_proj nn.Linear(hidden_dim, hidden_dim) self.out_proj nn.Linear(hidden_dim, hidden_dim) def forward(self, x, maskNone): batch, seq_len, _ x.shape Q self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) K self.k_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) V self.v_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) out torch.matmul(attn, V).transpose(1, 2).contiguous().view(batch, seq_len, -1) return self.out_proj(out)提示transpose之后一定要用contiguous()否则view会报错。这是因为 transpose 只改变了张量的步长信息内存布局没变view 要求内存连续。4. 实操过程与核心环节实现4.1 环境搭建与依赖安装我用的环境是 Python 3.10 PyTorch 2.1 CUDA 12.1。如果你没有 GPUCPU 也能跑只是训练速度会慢很多。一个简化版 Transformer 在 CPU 上跑一个 epoch 大概要 10 分钟GPU 上只要 30 秒。conda create -n ai-scratch python3.10 conda activate ai-scratch pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 pip install numpy tqdm matplotlib这里有个坑PyTorch 版本和 CUDA 版本必须匹配。我一开始装了 CUDA 11.8 的 PyTorch但系统驱动是 12.1结果torch.cuda.is_available()一直返回 False。后来查了官方兼容性表格才发现问题。如果你不确定直接去 PyTorch 官网用它的安装命令生成器选好系统和 CUDA 版本它会给你一条正确的 pip 命令。4.2 数据准备与预处理实战我用的是一个小型中文对话数据集大概 5000 条问答对。预处理流程如下读取原始文本去掉空行和异常字符在每句话前后加上bos和eos统计字符频率保留出现次数大于 5 的字符其余映射到unk构建字符到 ID 的映射表保存为 JSON 文件把所有句子转成 ID 序列保存为 PyTorch 的.pt文件这里的关键决策是词表大小的选择。词表太小很多字符变成unk模型学不到东西词表太大嵌入矩阵参数量爆炸容易过拟合。我的数据集字符种类大概 3000 个去掉低频后剩 2000 左右嵌入矩阵就是2000 x 256参数量 50 万完全可以接受。另一个细节是序列长度的截断。我统计了所有句子的长度分布95% 的句子在 50 个字符以内最长的有 200 多个字符。我把最大长度设为 64超出的截断不足的填充。这样每个批次的张量形状是(batch_size, 64)计算量可控。4.3 模型组装与参数初始化模型组装就是把前面实现的组件像搭积木一样拼起来。一个标准的 Transformer Block 包含多头自注意力层残差连接 层归一化前馈网络两个线性层 激活函数残差连接 层归一化前馈网络的隐藏维度通常是4 * hidden_dim这是原论文里的设置。我试过2 * hidden_dim和8 * hidden_dim发现 4 倍确实是个甜点——再小表达能力不足再大参数量增加但效果提升不明显。参数初始化我用了 PyTorch 默认的初始化但做了一点调整嵌入层的权重乘以sqrt(hidden_dim)。这是因为嵌入层的输出会和位置编码相加如果嵌入值太小位置编码会主导信号如果太大位置编码又不起作用。乘以sqrt(hidden_dim)是为了让两者的量级匹配。位置编码我用的是正弦余弦函数这是原论文的方案。公式是PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中pos是位置索引i是维度索引。这个设计的巧妙之处在于不同位置的编码可以通过线性变换相互表示这让模型更容易学到相对位置关系。我试过用可学习的位置嵌入效果差不多但参数量多了max_len * hidden_dim个对于小数据集来说容易过拟合。4.4 训练循环与学习率调度训练循环的骨架很简单前向传播、计算损失、反向传播、更新参数。但魔鬼在细节里。损失函数我用的是交叉熵但加了标签平滑label smoothing。标签平滑把硬标签[0, 0, 1, 0]变成软标签[0.01, 0.01, 0.97, 0.01]防止模型对某个类别过度自信。这个技巧在翻译和文本生成任务里几乎是标配能提升 1-2 个 BLEU 点。优化器用 AdamW学习率设成 3e-4。AdamW 和 Adam 的区别在于权重衰减的实现方式——AdamW 把权重衰减从梯度更新里解耦出来更符合理论预期。学习率 3e-4 是 Transformer 类模型的常用值太大容易震荡太小收敛慢。学习率调度用余弦退火加线性预热。预热的意思是前 1000 步学习率从 0 线性增加到 3e-4这是为了防止训练初期梯度不稳定。余弦退火则是让学习率按照余弦曲线从 3e-4 降到 1e-6帮助模型在后期精细调整。def get_lr(step, warmup_steps, max_steps, max_lr): if step warmup_steps: return max_lr * step / warmup_steps progress (step - warmup_steps) / (max_steps - warmup_steps) return max_lr * 0.5 * (1 math.cos(math.pi * progress))梯度裁剪是另一个必备技巧。我把梯度的 L2 范数限制在 1.0超过就按比例缩放。这能防止偶尔出现的梯度爆炸把权重更新到离谱的值。我试过不裁剪结果训练到第 3000 步左右 loss 突然变成 nan加了裁剪之后就稳定了。5. 常见问题与排查技巧实录5.1 训练不收敛从 loss 曲线诊断问题训练不收敛是最常见的问题但 loss 曲线的形状能告诉你很多信息。我整理了一个速查表loss 曲线形状可能原因排查方法解决方案一直震荡不下降学习率太大打印每步的梯度范数降低学习率 10 倍下降很慢但稳定学习率太小检查学习率调度器提高学习率或减少预热步数先降后升过拟合对比训练集和验证集 loss加 dropout、权重衰减、早停突然变成 nan梯度爆炸检查梯度范数加梯度裁剪、降低学习率完全不下降数据或标签有问题用随机权重跑一遍看 loss 是否接近 ln(词表大小)检查数据预处理和标签对齐我遇到过一次 loss 完全不下降的情况排查了半天发现是标签错位。在语言模型里输入是[bos, 你, 好, eos]标签应该是[你, 好, eos, pad]也就是输入向右移一位。我一开始忘了移位模型在学预测自己loss 当然降不下去。5.2 显存溢出批次大小与序列长度的权衡显存溢出是另一个高频问题。Transformer 的显存占用主要来自注意力矩阵形状是(batch, num_heads, seq_len, seq_len)。假设batch32, heads8, seq_len64这个矩阵就有32 * 8 * 64 * 64 100 万个元素float32 下占 4MB。看起来不大但反向传播时还要存中间激活值实际占用可能是前向的 3-5 倍。我的经验是先固定序列长度再调批次大小。序列长度对显存的影响是平方级的批次大小是线性的。如果显存不够优先降低序列长度。另外混合精度训练能省一半显存PyTorch 的torch.cuda.amp用起来很简单scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意混合精度下有些操作会溢出比如 softmax 里的 exp。PyTorch 的 autocast 会自动处理这些但如果你手写了自定义的 CUDA 核函数需要自己加类型转换。5.3 模型过拟合正则化手段的组合拳小数据集上训练 Transformer过拟合几乎是必然的。我的数据集只有 5000 条模型参数量 200 万训练集 loss 降到 0.5 的时候验证集 loss 已经开始上升了。我用了三招组合拳第一招是dropout。在嵌入层、注意力层、前馈网络里都加了 dropout概率 0.1。dropout 的本质是每次前向传播随机丢弃一部分神经元相当于训练了多个子网络然后集成。推理时关闭 dropout用全部神经元相当于取平均。第二招是权重衰减。AdamW 的权重衰减系数设成 0.01让权重倾向于小值。权重越小模型越简单越不容易过拟合。第三招是早停。每训练一个 epoch 就在验证集上评估一次如果验证集 loss 连续 3 个 epoch 不下降就停止训练。这个策略简单但有效能省下大量无用的训练时间。5.4 推理阶段从模型输出到可读文本训练完之后推理阶段也有坑。模型输出的是每个位置对词表中所有 token 的 logits形状是(batch, seq_len, vocab_size)。要生成文本需要逐个 token 采样。最简单的策略是贪心搜索每个位置选概率最大的 token。但这样生成的文本往往很无聊因为模型总是选最安全的词。更好的策略是温度采样把 logits 除以温度参数T然后 softmax。T1是原始分布T1让分布更尖锐更保守T1让分布更平坦更随机。我一般用T0.8在多样性和连贯性之间取平衡。还有一个技巧是top-k 采样只从概率最高的 k 个 token 里采样其余概率设为零。这能避免模型偶尔抽风选到完全不相关的词。k50是个常用值。def generate(model, input_ids, max_new_tokens50, temperature0.8, top_k50): model.eval() with torch.no_grad(): for _ in range(max_new_tokens): logits model(input_ids)[:, -1, :] / temperature top_k_logits, top_k_indices torch.topk(logits, top_k) probs torch.softmax(top_k_logits, dim-1) next_token top_k_indices.gather(-1, torch.multinomial(probs, 1)) input_ids torch.cat([input_ids, next_token], dim-1) if next_token.item() eos_token_id: break return input_ids6. 从零实现之后我学到了什么三周时间从零实现了一个简化版 Transformer跑通了训练和推理。回头看最大的收获不是代码本身而是对 shape 的敏感度。以前看到(batch, seq_len, hidden_dim)就是一个抽象概念现在我能在大脑里模拟每一步的维度变换知道哪个维度该转置、哪个维度该广播。另一个收获是对超参数的理解。以前调学习率就是试 1e-3、1e-4、1e-5现在我知道学习率和批次大小是耦合的——批次越大梯度估计越准可以用更大的学习率。线性缩放规则说批次大小翻倍学习率也翻倍。虽然实际中不会这么激进但至少有了理论依据。最后分享一个小技巧每次修改代码后先用一个极小的数据集跑通全流程。我习惯用 100 条数据、2 个 epoch、批次大小 4确保代码没有 shape 错误、没有 NaN、loss 能下降。这个“冒烟测试”能帮你快速发现低级错误避免在完整数据集上浪费几个小时才发现一个拼写错误。这个项目后续还可以这样扩展加入 Beam Search 提升生成质量、实现 KV Cache 加速推理、尝试不同的位置编码方案如 RoPE、ALiBi、把模型规模扩大到能处理真实任务。但那是下一步的事了先把从零实现的这条路走通后面的路会顺很多。