从零手写Transformer：自动微分、注意力机制与AI工程实战全解析-育师

1. 从零构建AI工程：为什么要做这件事

先说一句大实话：过去两三年里，我见过太多人把“AI工程”理解成“调库”。写两行transformers的代码、跑通一个GPT接口、用LangChain串几个Prompt，就觉得自己是AI工程师了。不能说完全错，但这些事情更像“AI应用开发”，离“AI工程”还有一段实实在在的距离。

我发起ai-engineering-from-scratch这个项目的初衷很简单：不带任何框架、不调任何现成的推理库，从最基本的数学原理和代码开始，一步步把现代大模型的核心部件手写出来。项目名称里的from-scratch不是噱头，而是真的从张量操作、反向传播、注意力机制这些最底层的东西写起。目标读者不是已经在大厂做算法的人，而是那些想真正理解AI底层原理、希望从“会用”进阶到“懂为什么”的工程师和学生。

这个项目能解决的问题非常具体：当你在面试中被问到“BatchNorm在推理阶段和训练阶段有什么不同”，当你在调参时面对“Loss不下降但不知道是学习率问题还是梯度消失问题”，当你想给团队解释“为什么这块GPU显存不够但理论计算量明明不大”——如果你亲手写过这些模块，你根本不需要背答案，你脑子里就有那幅图景。

从内容设计上，我把整个项目分成了四个递进阶段：数学基础与自动微分、经典神经网络架构、注意力机制与Transformer、分布式训练与推理优化。每个阶段都配套完整的Python代码实现、数值实验和理论推导笔记。下面我把每一层的设计思路、实现关键点和踩过的坑详细拆开讲，这份经验可以直接作为你自建学习路线的参考。

2. 基础层：自动微分和训练循环，是所有上层建筑的基石

2.1 为什么第一步是手写自动微分，而不是直接用PyTorch

很多人一开始不理解：既然最终要学Transformer，为什么不直接讲注意力机制？我的回答是：注意力机制里最核心的Softmax(QK^T / sqrt(d))V一旦需要反向传播，你就必须理解雅可比矩阵如何通过链式法则传播梯度。如果用的是PyTorch，你看到的是一个out.backward()，但梯度到底怎么流经Q、K、V，你仍然是一团黑。

所以我在项目的第一阶段，用纯NumPy实现了一个微型的自动微分引擎，支持标量、向量和矩阵三种张量类型，实现了加法、乘法、矩阵乘、转置、求和、reshape等常用算子。每个算子都定义了forward和backward两个方法，用计算图记录反向传播路径。

这里有一个关键设计决策：我用的是动态计算图，而不是静态图。动态图意味着每次前向传播都会重新构建计算图，这对调试极其友好——你可以在任意一个中间节点打印梯度。静态图（像早期TensorFlow）虽然优化空间更大，但实现起来复杂得多，对初学者不友好。我的建议是先用动态图理解核心概念，后续需要性能再研究图优化。

# 极简张量类的核心：forward和backward class Tensor: def __init__(self, data, requires_grad=False): self.data = np.array(data, dtype=np.float64) self.requires_grad = requires_grad self.grad = None self._backward = lambda: None self._prev = set() def __add__(self, other): other = other if isinstance(other, Tensor) else Tensor(other) out = Tensor(self.data + other.data, requires_grad=self.requires_grad or other.requires_grad) out._prev = {self, other} def _backward(): if self.requires_grad: self.grad = self.grad + out.grad if self.grad is not None else out.grad.copy() if other.requires_grad: other.grad = other.grad + out.grad if other.grad is not None else out.grad.copy() out._backward = _backward return out def backward(self): topo = [] visited = set() def build_topo(v): if v not in visited: visited.add(v) for child in v._prev: build_topo(child) topo.append(v) build_topo(self) self.grad = np.ones_like(self.data) for v in reversed(topo): v._backward()

这里有个细节容易踩坑：梯度的累加问题。当一个张量被多个算子共用时（比如x同时参与了y1 = x * 2和y2 = x + 1），反向传播时x的梯度是两个分支梯度的和。如果每次直接赋值而不是累加，后一个分支的梯度就会覆盖前一个。我在实现时用self.grad = self.grad + ...而不是self.grad = ...，就是为了解决这个问题。这个问题在真实框架里也存在——PyTorch默认就是累加梯度，所以每次optimizer.zero_grad()是必须的。

2.2 训练循环细节：BatchSize、学习率和梯度裁剪的相互作用

有了自动微分引擎之后，我实现了第一个完整的训练循环——在合成数据集上训练两层的MLP做二分类。这一步看似简单，但它把整个机器学习训练的所有关键环节都串起来了：参数初始化、前向传播、Loss计算、反向传播、参数更新。

这里我特别强调一个容易被忽略的环节：参数初始化方式对收敛的影响。我用两种方式做了对比实验——标准正态分布初始化和Xavier初始化。实验结果显示，在同样的学习率（0.1）和同样的迭代次数（2000轮）下，Xavier初始化的模型最终准确率可以达到96%以上，而标准正态初始化的模型只能到88%。原因在于Xavier初始化考虑了前一层和后一层的神经元数量，使得每层的输入方差和输出方差保持一致，从而避免了信号在前向传播过程中被放大或衰减得太厉害。

# Xavier初始化 def xavier_init(fan_in, fan_out): limit = np.sqrt(6.0 / (fan_in + fan_out)) return np.random.uniform(-limit, limit, size=(fan_in, fan_out))

训练循环里另一个重要细节是梯度裁剪的时机。很多入门教程根本不提梯度裁剪，但在实际训练中，尤其是使用深度网络时，梯度范数很容易爆炸——尤其是Loss曲面比较陡峭的时候。我的做法是设置一个阈值（比如max_norm=1.0），在每轮参数更新之前计算整个梯度的L2范数，如果超过阈值就对所有参数的梯度做等比缩放：

# 梯度裁剪的核心逻辑 total_norm = 0.0 for p in parameters: if p.grad is not None: total_norm += np.sum(p.grad.data ** 2) total_norm = np.sqrt(total_norm) clip_coef = max_norm / (total_norm + 1e-6) if clip_coef < 1: for p in parameters: if p.grad is not None: p.grad.data *= clip_coef

这个操作的物理意义可以理解为：当梯度异常大的时候，它往往携带的不仅仅是真实信号，还包含大量噪声。直接把它等比例缩小，保留方向信息、削弱幅度噪声，参数更新就更稳定。我在项目中专门跑了一组梯度裁剪开关的对比实验，结果显示裁剪后训练曲线的震荡幅度显著减小。

3. 核心突破层：从零手写Transformer，模型并行与注意力掩码的工程细节

3.1 手写多头注意力时，最容易搞错的四个维度

整个项目难度最高的部分就是手写Transformer。我放弃PyTorch的nn.MultiheadAttention，直接用NumPy实现了完整的多头注意力模块。这个过程中最折磨人的不是算法理解，而是维度调换时的张量形状管理。

多头注意力的标准流程是：输入的(batch, seq_len, d_model)经过三个不同的权重矩阵得到Q、K、V，然后把d_model维度切成num_heads份，变成(batch, num_heads, seq_len, d_k)，分别做注意力计算后再拼接回去。

我踩过最大的一个坑是：缩放因子除以的是d_k而不是d_model。d_k = d_model / num_heads，比如d_model=512、num_heads=8，那么每个头的维度是64，缩放因子就是sqrt(64)=8。如果用sqrt(512)来缩放，点积结果会被过度压缩，Softmax之后概率分布会变得过于平滑，注意力几乎等于均匀分布，模型学不到有效的关联关系。

# 多头注意力核心计算（NumPy实现） def forward(self, Q, K, V, mask=None): batch_size = Q.shape[0] d_k = self.d_model // self.num_heads # 线性投影 Q = Q @ self.W_Q K = K @ self.W_K V = V @ self.W_V # 重塑为多头：从 (batch, seq, d_model) -> (batch, heads, seq, d_k) Q = Q.reshape(batch_size, -1, self.num_heads, d_k).transpose(0, 2, 1, 3) K = K.reshape(batch_size, -1, self.num_heads, d_k).transpose(0, 2, 1, 3) V = V.reshape(batch_size, -1, self.num_heads, d_k).transpose(0, 2, 1, 3) # 计算注意力分数 scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(d_k) if mask is not None: scores = np.where(mask == 0, -1e9, scores) attn_weights = softmax(scores, axis=-1) output = attn_weights @ V output = output.transpose(0, 2, 1, 3).reshape(batch_size, -1, self.d_model) output = output @ self.W_O return output

关于mask的处理，我用的是-1e9而不是0。原因非常明确：Softmax内部有exp()操作，如果被掩码的位置分数设为0，exp(0)=1，这并不能让它完全消失——它只是和其他正常位置一样参与Softmax概率分配。只有设为负无穷或绝对值足够大的负数，exp(-1e9)才趋近于0，Softmax出来的结果才约等于0。这个细节决定了模型训练时能不能正确忽略未来位置的信息。

3.2 为什么我坚持用NumPy而不用PyTorch，以及什么时候该换

这是个值得展开的问题。项目早期，我用NumPy实现所有算子，因为NumPy在接口设计上强迫你显式管理每一次张量运算、每一个维度变换，你看得见数据流动的每一步。PyTorch把所有底层细节都封装好了，你反而看不到梯度是怎么流动的。

但学到后面我承认，纯NumPy实现训练大模型是不可能的。尤其是当我推到LLaMA级别的模型（虽然只是缩小版），显存管理和计算图优化已经不是手写能搞定的事了。所以我在项目后期做了一个明确的切换声明：前几章坚持手写，后几章切换到PyTorch实现，但保持“不调用高层封装”的原则——比如手写KL散度损失函数、手写Rotary Position Embedding、手写SwiGLU激活函数，只使用PyTorch的自动微分和GPU加速能力。

这里给读者一个可复现的操作建议：

理解阶段：用NumPy手写，配合小规模数据（比如序列长度32、模型维度64）验证接口正确性。
效率阶段：用PyTorch但禁用torch.nn.TransformerEncoderLayer，所有模块自己实现。
最终阶段：对比NumPy和PyTorch实现的速度差异，量化工程优化的收益。

这三个阶段衔接起来，认知是最完整的。

3.3 LayerNorm的坑：训练和推理阶段行为不一致

实现Transformer时，每个子层后面都要跟LayerNorm。这个操作看起来只是“减均值除方差再乘缩放”，但有一个极容易被忽略的细节：BatchNorm在训练时用当前batch的统计量，在推理时用训练阶段的滑动平均值；但LayerNorm永远用当前样本的统计量，因为它是逐样本做的归一化。这就是为什么LayerNorm更适合Transformer而BatchNorm不适合的底层原因——序列长度变化频繁，batch统计量不稳定。

在NumPy里实现LayerNorm很容易犯错的地方是归一化维度的选择。正确的做法是沿着特征维度（最后一维）做均值方差计算，而不是沿着序列维度。如果是图像数据，BatchNorm沿着通道维度跨Batch做；但文本数据里每个token的语义是独立的，必须逐token归一化。写代码时我会明确np.mean(x, axis=-1, keepdims=True)，这个keepdims=True也是常见的坑——如果忘了，后面的广播就会出问题。

def layer_norm(x, gamma, beta, eps=1e-5): mean = np.mean(x, axis=-1, keepdims=True) var = np.var(x, axis=-1, keepdims=True) x_hat = (x - mean) / np.sqrt(var + eps) return gamma * x_hat + beta

eps的选择也很讲究。太小（比如1e-8）会导致数值不稳定，特别是当特征维度很小、方差接近0时除出来的数会极大；太大会让归一化不够彻底。Transformer里通常取1e-5或1e-6，我实测下来1e-5最稳。

4. 性能优化层：显存管理、模型并行与推理加速的落地经验

4.1 大模型推理时，KV Cache为什么能省这么多计算

训练好一个Transformer之后，下一步就是让它做推理生成。这里我引入了一个在工业界极其重要但很多科普文章不讲的概念：KV Cache（键值缓存）。它的核心思想是：生成式模型在自回归解码时，每一步只生成一个新的token，前面所有token的K和V其实已经完全计算过了，不需要重新算。

我第一次测试KV Cache的效果时，生成512个token的耗时从大约3.8秒降到1.2秒，加速比接近3倍。这个提升幅度跟序列长度之间是线性增长的关系——序列越长，节省的计算越多。原因很好理解：如果不带Cache，每一步都要把前面所有token重新过一遍注意力计算，复杂度是O(n^2)的（n是序列长度）；而带Cache之后，每一步只计算当前token与历史KV的注意力，总复杂度降到O(n)。

# KV Cache的伪代码，实际实现时请结合你的张量存储格式适配 def generate_with_kv_cache(model, input_ids, max_new_tokens=100): # 初始输入过一遍模型，得到第一个新token和对应的KV K_cache = [] V_cache = [] for layer_idx, layer in enumerate(model.layers): K, V = layer.compute_kv(input_ids) K_cache.append(K) V_cache.append(V) next_token = model.next_token(input_ids, K_cache, V_cache) generated = [next_token] for _ in range(max_new_tokens - 1): # 只需输入最新的token new_input = generated[-1] for layer_idx, layer in enumerate(model.layers): new_K, new_V = layer.compute_kv(new_input) K_cache[layer_idx] = np.concatenate([K_cache[layer_idx], new_K], axis=1) V_cache[layer_idx] = np.concatenate([V_cache[layer_idx], new_V], axis=1) next_token = model.next_token_with_cache(new_input, K_cache, V_cache) generated.append(next_token) return generated

但KV Cache不是没有代价的——它牺牲显存换速度。512序列长度、8个注意力头、d_k=64、FP16存储，每一层的KV缓存大约占用2 * batch_size * seq_len * num_heads * d_k * 2 bytes。算下来一层模型在batch=1、seq_len=512时大约占用1MB，听起来不多，但如果是70B模型70层，显存占用就变得非常可观了。所以工业界还有KV Cache量化（比如8bit或4bit存储K和V）的技术，就是在这个空间和精度的权衡上做优化。

4.2 五维张量并行：当我用多块GPU训练一个大模型时在做什么

在项目的收尾阶段，我引入了一个真实的分布式训练场景：把一个小规模的Transformer（1亿参数级别）通过张量并行和多卡数据并行跑在4块GPU上。这里我最想强调的不是代码，而是并行策略选择的判断思路。

数据并行是最简单的：每张卡有完整的模型副本，但处理不同的batch数据，每轮训练后同步梯度。它的通信量跟模型参数量成正比，但在模型较小、数据量大的时候效率最高。张量并行则是把同一个层的权重切分到多张卡上（比如把W_Q按行切成两块，每块卡算一半的输出维度），然后通过AllReduce合并结果。通信量更大，但单卡显存占用更小。

实际工程中几乎没有人只用一种并行策略。现代大模型训练普遍采用混合并行：在数据维度用数据并行扩大吞吐，在模型维度用张量并行或流水线并行突破单卡显存限制。我最终在项目里实现的方案是：2路数据并行 × 2路张量并行，通过NCCL的AllReduce做梯度同步。

# 模拟张量并行：把线性层拆成两块 # 假设 W 原本是 [out_dim, in_dim] # GPU0 持有 W[:out_dim//2], GPU1 持有 W[out_dim//2:] def parallel_matmul(x, local_W, is_master): # 本地计算 local_out = x @ local_W.T # AllReduce: 本质上是对序列维度做拼接 # 各卡把自己的输出全量拼起来 full_out = all_gather(local_out, dim=-1) return full_out

一个我在实际并行训练中踩过的坑：AllReduce之后的梯度归一化。如果使用数据并行，每张卡算出的梯度是基于自己那部分batch的，平均值才是真正的全局梯度。如果忘了把梯度除以world_size，学习率等效放大了world_size倍，训练很容易震荡甚至发散。这个问题在TensorFlow时代经常遇到，PyTorch的DistributedDataParallel已经帮你处理了，但手写时特别容易漏掉。

4.3 推理优化的三板斧：量化、剪枝、蒸馏的落地优先级

从我实操的体感来看，推理优化不是越复杂越好，而是先做见效快的、再做锦上添花的。我处理过几个真实项目之后，总结出一个性价比排序：

第一步：FP16/INT8量化。对模型权重做量化，内存直接减半再减半。很多显卡在FP16下的计算吞吐远高于FP32，这几乎是白捡的性能。具体操作时，最简单的方案是把权重从FP32转成FP16，精度损失通常可以忽略；想进一步压缩就做INT8量化，但需要校准过程以避免某些敏感层的精度严重下降。

第二步：算子融合。把LayerNorm的均值方差计算和后面的缩放偏移融合成一个kernel，把Q@K^T和sqrt(d_k)融合，减少kernel启动开销。这个优化在GPU上能明显看到效果，因为小kernel的启动时间往往比kernel本身的计算时间还要长。

第三步：蒸馏或剪枝。这些方法需要额外的训练或结构调整，周期长、风险高。我一般只在前面几步还不够用时才考虑。

为什么量化放在最前面？因为量化是纯粹的“端口”优化——模型结构不变、训练不重跑、代码改动最小，收益却最大。我在一个移动端模型项目中做过实测，从FP32压缩到INT8后，模型文件从120MB降到30MB，端侧推理耗时从200ms降到80ms。同类项目里，这个性价比是其他优化手段很难比的。

5. 避坑实录：手写AI工程时，我吃过的那些教训

5.1 数值稳定性：Softmax溢出问题，我差点以为模型写错了

在实现Softmax时，最容易遇到的一个问题是上溢出（overflow）。这在Dot-Product Attention里尤其常见——当序列长度很长、Q@K^T的值很大时（比如超过709，也就是exp能表达的极限），np.exp(x)会变成inf，然后inf / inf就是nan，整个梯度就崩了。

解决方案是减去最大值。在计算exp之前先找出行内最大值，然后每个元素减去它。这个数学上等价，因为：

softmax(x_i) = exp(x_i) / sum(exp(x_j)) = exp(x_i - max) / sum(exp(x_j - max))

因为exp(x_i)/sum(exp(x_j))对整体加减一个常数不变。我在实现时写成：

def stable_softmax(x): x = x - np.max(x, axis=-1, keepdims=True) exp_x = np.exp(x) return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

后来我把这个稳定的版本写进所有算子中。这个经验教训是：在数值计算中，算法等价性和数值稳定性是两码事，一定要优先保证数值稳定。类似的坑还有在对数域计算交叉熵时直接对概率取log，如果概率被四舍五入到0，log(0)直接负无穷。更稳妥的方式是在log里面加一个小数clip或eps。

5.2 训练不收敛时，按什么顺序排查

我调试过非常多次训练不收敛的问题，总结出一个系统的排查顺序，这也是整个项目中最实用的经验之一：

第一步：检查Loss数量级。如果是分类任务，初始Loss大约在-log(1/类别数)附近，太高太低都说明初始化有问题。

第二步：检查梯度范数。如果超过100，基本可以断定是梯度爆炸；如果低于1e-8，基本是梯度消失。两个方向的处理方式完全不同——前者加梯度裁剪或降低学习率，后者要检查是否有过深的网络或激活值饱和（比如sigmoid激活到极端区域）。

第三步：用小规模数据过拟合。这个技巧非常有用。选10条训练样本，让模型反复跑几百轮，如果Loss能降到极小值（接近0），说明模型有足够的容量学习，问题大概率出在数据处理或学习率上；如果连10条都过拟合不了，那问题就在模型结构或输入输出编码上。这一步能迅速缩小排查范围。

第四步：检查数据处理。很多看似模型的问题其实是数据问题。标签和输入错位了吗？归一化做对了吗？有没有NaN混进来？我至少有一半的排查时间花在这里。

那我在手写实现时常见的一个愚蠢错误：更新参数时忘了减学习率，或者反过来对参数做了+=而不是-=。这种低级错误在手动实现训练循环时真的会出现，检查时我会把参数更新前后打印几个值对比一下，基本能马上定位。

5.3 评估指标陷阱：Loss降了，不代表模型变好了

训练稳定之后还有一个隐藏陷阱：监控的指标不对。很多人在训练日志里只盯train_loss，但train_loss和模型质量之间并不总是正相关。我训练语言模型时遇到过一种情况：train_loss一直在降，但BLEU分数纹丝不动。后来查出来是模型在疯狂复读高频词，这些词对Loss贡献大，但生成的句子几乎没有可用性。

所以我的项目里从第一天起就同时监控几个指标：train_loss、val_loss、梯度范数、学习率以及任务级别的指标（如BLEU、Accuracy）。训练过程中我会每隔50步打一次这些值，看它们是否有一个合理的协同变化趋势。如果train_loss降但val_loss升，那就是过拟合；如果两个都降但任务指标不升，那就是优化目标和任务目标不一致，需要检查是不是Loss设计有问题。

6. 复盘与延伸：这条路能走多远

6.1 我踩过“框架黑盒”的坑，才明白from-scratch的真正价值

项目做到现在，我最深的体会是：从零手写不是目的，而是手段。手写代码的过程本质上是在做“认知卸载”——把那些原本模糊的、只存在于概念层面的“应该大概是这样”的东西，变成精确的、可执行的、能调试的代码逻辑。当你亲手写过一次注意力机制之后，再看论文里那些变体（Mamba的线性注意力、GQA的分组查询、MLA的潜在多头注意力），你会立刻理解它们在改什么、为什么改。这种感觉和只读过论文摘要完全不同。

我自己就有这样的经历：以前看RoPE（Rotary Position Embedding）总是觉得玄乎，直到我按照公式一步步实现，发现它只是在Q和K的每个二维子空间里做了旋转操作，大小不变、角度按位置编码改变，然后用点积的性质自然让相对位置信息融入注意力分数。那一刻我才真正理解这是“通过旋转让位置信息以角度形式参与点积”，而不仅仅是一堆奇怪的三角函数公式。

6.2 后续可以延伸的方向

项目目前完成了基础四阶段，接下来我有几个明确的扩展计划：

方向一：多模态扩展。在纯文本Transformer的基础之上，加上视觉编码器（比如用卷积层把图像切片成patch），实现一个简化版的CLIP。跨模态对齐是当前工业界最热门的方向之一，而理解了Transformer之后，跨模态的核心挑战从“模型结构”变成了“数据对齐策略”，这个认知转变非常有价值。

方向二：强化学习对齐。训练出基础模型之后，加上PPO微调做RLHF。很多人以为RLHF很难，但其实核心技术在于奖励模型和策略梯度的实现，而这两者恰恰是前面自动微分引擎已经解决的问题。

方向三：推理系统完善。加入更多的推理加速手段（比如PagedAttention的简化实现、连续批处理调度），这是从“模型能跑”到“服务能上线”的必要一步。