ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

从零手搓AI工程:手写自动微分、注意力与KV Cache的底层原理

从零手搓AI工程:手写自动微分、注意力与KV Cache的底层原理 1. 从零手搓AI工程为什么我不建议你直接调包第一次看到ai-engineering-from-scratch这个项目名的时候我正坐在工位上啃一个调了三天还没收敛的推荐模型。说实话那一瞬间我是有点不屑的——这年头谁还从零写AI工程HuggingFace 上pipeline()一行代码就能跑推理LangChain 拖几个节点就能搭个 RAG连微调都有 LoRA 这种省显存的神器何必自讨苦吃去手写矩阵乘法但后来我改主意了。原因很简单调包调得越爽出问题的时候死得越惨。我见过太多这样的情况——线上推理延迟突然从 80ms 飙到 2s团队里没人说得清 KV Cache 到底是怎么占显存的RAG 检索召回率掉了 15 个点大家对着 embedding 模型干瞪眼不知道是分块策略的问题还是向量归一化没做模型输出开始胡言乱语排查半天才发现是 tokenizer 的 padding 方向搞反了。这些坑光靠调 API 是永远填不平的。ai-engineering-from-scratch这个项目标题本质上说的就是这么一件事把 AI 工程里那些被封装层藏起来的核心环节一个一个拆开用最朴素的方式重新实现一遍。它不是让你造一个比 PyTorch 更快的框架而是让你在亲手写一遍的过程中真正理解每一层抽象下面到底发生了什么。这篇文章适合谁看三类人。第一类是有一定 Python 基础、用过 PyTorch 或 TensorFlow 但只会调model.fit()的开发者想搞清楚训练循环里到底发生了什么第二类是准备面试 AI 工程岗位、被问到“手写 attention”就卡壳的人第三类是像我一样被线上问题折磨过、决心把底层原理补回来的工程师。如果你只是想快速搭个 demo那这篇文章可能不适合你——但如果你想在 AI 工程这条路上走得更远从零实现一遍是绕不过去的坎。我打算按我自己实际踩坑的顺序来写先讲整体设计思路和为什么这么选型再拆核心模块的实现细节然后是完整的实操流程和参数计算最后是我踩过的那些坑和排查技巧。全程不藏私能抄的代码直接给你。2. 整体设计与技术选型为什么是这套组合2.1 从零实现的边界在哪里“从零”这个词很容易被误解。有人觉得从零就是连numpy都不能用得用纯 Python 列表做矩阵运算——那是自虐不是学习。我的理解是从零的边界应该划在“核心算法逻辑”这一层。也就是说张量存储、自动微分、优化器更新、注意力计算、分词、采样这些 AI 工程的核心环节你得自己写但底层的 BLAS 加速、CUDA 调度、内存分配这些交给numpy和torch的底层去处理没必要重复造轮子。这个边界划定的逻辑很实在你手写一遍softmax和layer_norm是为了理解数值稳定性的重要性比如为什么要减去最大值你手写一遍反向传播是为了理解计算图和梯度累积你手写一遍 KV Cache是为了理解推理时显存到底被什么吃掉了。但你没必要手写一个矩阵乘法内核因为那属于系统编程的范畴跟 AI 工程的核心能力关系不大。所以这个项目的技术栈我最终定成了Python 3.10 NumPy 做数值计算底座 PyTorch 只用来做梯度校验的对照 纯手写的核心模块。为什么用 NumPy 而不是直接用 PyTorch 的 tensor因为 NumPy 没有自动微分你被迫要自己推导每一层的反向传播公式这个过程才是真正长功力的地方。等你手推完一遍再回头看 PyTorch 的autograd会有一种“原来你帮我干了这些”的顿悟感。2.2 模块拆解与依赖顺序整个项目我拆成了六个核心模块依赖关系是层层递进的不能跳着做模块核心内容依赖预计耗时张量基础多维数组、广播、矩阵运算无1天自动微分计算图、反向传播、梯度累积张量基础2天神经网络层Linear、LayerNorm、Softmax、激活函数自动微分2天注意力机制Self-Attention、Multi-Head、KV Cache神经网络层3天分词器BPE、WordPiece、特殊token处理无2天训练与推理训练循环、优化器、采样策略全部2天这个顺序不是随便排的。自动微分必须建立在张量基础之上因为计算图的节点就是张量操作注意力机制必须建立在神经网络层之上因为 attention 本质就是几个线性层加 softmax 的组合分词器可以独立做因为它不依赖数值计算但它是训练和推理的前置环节。我建议你严格按这个顺序来跳步的话很容易在某个环节卡住回头补基础反而更费时间。2.3 为什么不用现成框架的深层考量有人会问既然 PyTorch 已经这么好用了为什么还要手写我拿一个真实场景来说明。之前我遇到一个线上问题模型在 batch size 从 8 调到 16 之后loss 突然变成 NaN。团队里用 PyTorch 的同事排查了两天最后发现是 LayerNorm 在某个特定维度上的数值溢出。如果他对 LayerNorm 的实现细节有概念知道它内部做了(x - mean) / sqrt(var eps)这个操作就会立刻意识到var在 batch 变大时可能因为浮点精度问题变成极小值导致除零。这个问题的根因不在 PyTorch而在于使用者对底层数值稳定性的无知。手写一遍的价值就在这里你会被迫面对每一个eps、每一个clip、每一个detach你会知道它们不是可有可无的装饰而是保证训练不崩的生命线。这种认知调包是调不出来的。另外还有一个很现实的原因面试。现在 AI 工程岗位的面试手写 attention 几乎是标配。我面过不少人简历上写着“精通 Transformer”让他手写一个 scaled dot-product attention结果连Q K.T / sqrt(d_k)里的sqrt(d_k)是干嘛的都说不清。从零实现一遍这些都会变成肌肉记忆。3. 核心模块实现细节与实操要点3.1 张量基础广播机制是万恶之源张量模块是整个项目的地基我建议你不要一上来就写复杂的先把shape、stride、broadcast这三个概念吃透。我见过太多 bug 都是广播机制引起的——你以为两个张量是逐元素相乘结果 NumPy 悄悄给你广播成了外积shape 从(32, 128)变成了(32, 32, 128)然后一路错到 loss 计算才报错。手写张量类的时候核心要处理的是形状推断和广播规则。广播的规则其实就三条从右往左对齐维度维度相等或其中一个为 1 才能广播结果维度取两者中的较大值。我当时的实现是这样的def broadcast_shapes(shape1, shape2): # 从右往左对齐 result [] for d1, d2 in zip_longest(reversed(shape1), reversed(shape2), fillvalue1): if d1 ! d2 and d1 ! 1 and d2 ! 1: raise ValueError(f无法广播: {shape1} 和 {shape2}) result.append(max(d1, d2)) return tuple(reversed(result))这段代码看起来简单但它是后面所有矩阵运算的基础。实操心得写完广播逻辑后一定要用几个边界 case 测一下比如(3, 1)和(1, 4)广播成(3, 4)(5,)和(5, 1)广播成(5, 5)。我当初就是漏测了最后一个 case导致后面 attention 的 mask 处理错了整整一天。注意广播虽然方便但在 AI 工程里要慎用。尤其是涉及 batch 维度的运算隐式广播很容易掩盖 shape 错误。我的建议是在关键路径上显式写出reshape或expand让 shape 变化可见。3.2 自动微分计算图的构建与反向传播自动微分是 AI 工程的心脏。PyTorch 的autograd帮你做了这件事但你知道它内部是怎么工作的吗简单说前向传播时每个操作都会记录自己是怎么算出来的形成一个计算图反向传播时从 loss 出发沿着计算图反向应用链式法则算出每个参数的梯度。手写自动微分的核心是设计一个Tensor类它要保存三个东西data数值、grad梯度、_backward反向传播函数。每次做运算时除了算出结果还要定义这个结果的反向传播逻辑。我拿加法举例class Tensor: def __init__(self, data, requires_gradFalse): self.data np.array(data, dtypenp.float32) self.grad None self.requires_grad requires_grad 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_gradself.requires_grad or other.requires_grad) def _backward(): if self.requires_grad: self.grad (self.grad or 0) out.grad if other.requires_grad: other.grad (other.grad or 0) out.grad out._backward _backward out._prev {self, other} return out这里有个关键细节梯度要累加而不是覆盖。因为一个张量可能在计算图中被多次使用比如残差连接它的梯度是多个路径回传的总和。我当初就是用了覆盖导致残差连接那块的梯度永远只有一半训练 loss 降不下去排查了好久才发现。反向传播的触发逻辑是拓扑排序加逆序执行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()实操心得写完自动微分后一定要用数值梯度校验。方法很简单对某个参数加上一个极小的eps算两次 loss 的差值除以eps跟反向传播算出来的梯度对比。如果误差在1e-5量级说明实现是对的。我当初校验的时候发现某个激活函数的梯度差了 10 倍查了半天发现是链式法则里漏乘了一项。3.3 注意力机制从公式到代码的鸿沟Attention 的公式大家都背过Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V。但把这个公式变成能跑的代码中间隔着好几个坑。第一个坑是mask 的处理。训练时要用 causal mask 防止看到未来 token推理时要用 padding mask 忽略填充位。这两个 mask 的 shape 和加法方式都不一样。causal mask 是一个下三角矩阵padding mask 是一个跟序列长度相关的向量。我当初把两个 mask 搞混了导致模型在训练时能“偷看”未来loss 降得特别快但推理时效果一塌糊涂。第二个坑是softmax 的数值稳定性。直接算exp(x) / sum(exp(x))在 x 很大时会溢出。正确做法是先减去最大值def softmax(x, axis-1): x_max np.max(x, axisaxis, keepdimsTrue) exp_x np.exp(x - x_max) return exp_x / np.sum(exp_x, axisaxis, keepdimsTrue)这个x - x_max的操作数学上不改变结果因为分子分母同乘了一个常数但数值上能救命。这是 AI 工程里最经典的数值稳定技巧没有之一。第三个坑是KV Cache 的实现。推理时为了不重复计算要把之前算过的 K 和 V 缓存起来。这个逻辑听起来简单但涉及到一个关键问题cache 的 shape 怎么管理。我当初的实现是每次把新的 K、V 拼接到 cache 后面结果发现显存占用随着序列长度线性增长长文本推理直接 OOM。后来改成预分配固定长度的 cache用索引来标记有效位置显存才稳定下来。class KVCache: def __init__(self, max_len, n_heads, head_dim): self.k_cache np.zeros((max_len, n_heads, head_dim)) self.v_cache np.zeros((max_len, n_heads, head_dim)) self.pos 0 def update(self, k, v): seq_len k.shape[0] self.k_cache[self.pos:self.posseq_len] k self.v_cache[self.pos:self.posseq_len] v self.pos seq_len return self.k_cache[:self.pos], self.v_cache[:self.pos]提示KV Cache 是推理加速的核心但它也是显存杀手。一个 7B 模型序列长度 4096KV Cache 能占到好几个 G。理解它的计算方式2 * n_layers * n_heads * head_dim * seq_len * batch_size * dtype_size你就能明白为什么长上下文推理这么贵。3.4 分词器被低估的工程环节很多人觉得分词器不重要不就是把文本切成 token 吗但实际工程里分词器是很多诡异问题的源头。我遇到过模型输出重复、截断位置不对、特殊 token 被当成普通文本处理追根溯源都是分词器的锅。BPEByte Pair Encoding的核心思想是从字符级别开始不断合并出现频率最高的相邻对直到词表达到目标大小。手写 BPE 的关键在于合并规则的有序性——合并的顺序会影响最终的分词结果所以必须严格按照训练时学到的合并顺序来应用。def apply_bpe(self, text): tokens list(text) while len(tokens) 1: pairs [(tokens[i], tokens[i1]) for i in range(len(tokens)-1)] # 找到优先级最高的可合并对 mergeable [(self.merge_ranks[p], p) for p in pairs if p in self.merge_ranks] if not mergeable: break _, best_pair min(mergeable) # 执行合并 new_tokens [] i 0 while i len(tokens): if i len(tokens)-1 and (tokens[i], tokens[i1]) best_pair: new_tokens.append(tokens[i] tokens[i1]) i 2 else: new_tokens.append(tokens[i]) i 1 tokens new_tokens return tokens实操心得分词器写完一定要做 round-trip 测试也就是decode(encode(text)) text。我当初就是漏了这个测试导致某些 Unicode 字符在编解码后变了样模型训练时输入和推理时输入不一致效果直接崩了。另外特殊 token如[PAD]、[UNK]、[CLS]的处理要格外小心它们不应该参与 BPE 合并也不应该被普通文本匹配到。4. 完整实操流程从零到能跑的训练4.1 环境准备与依赖安装环境这块我不建议搞太复杂Python 3.10 加几个基础库就够了python -m venv venv source venv/bin/activate # Windows 用 venv\Scripts\activate pip install numpy matplotlib tqdm为什么不用装 PyTorch因为我们的核心模块是手写的PyTorch 只在最后做梯度校验的时候用一下。如果你实在想装装个 CPU 版本就行没必要折腾 CUDA。这个项目的目的是理解原理不是追求性能用 CPU 跑小规模数据完全够用。我建议你建一个这样的目录结构ai-engineering-from-scratch/ ├── tensor.py # 张量基础 ├── autograd.py # 自动微分 ├── nn.py # 神经网络层 ├── attention.py # 注意力机制 ├── tokenizer.py # 分词器 ├── train.py # 训练循环 └── tests/ # 测试用例每个模块独立成文件方便单独测试。不要把所有代码堆在一个文件里那样调试起来会疯掉。4.2 参数初始化与数值稳定性处理参数初始化是很多人忽略的环节但它直接决定训练能不能收敛。我当初用np.random.randn直接初始化结果第一轮 loss 就是 NaN。原因很简单标准正态分布的方差是 1但经过多层网络累积后激活值的方差会爆炸。正确的做法是用 Xavier 或 Kaiming 初始化。Xavier 的核心思想是让每一层的输出方差等于输入方差def xavier_init(fan_in, fan_out): limit np.sqrt(6.0 / (fan_in fan_out)) return np.random.uniform(-limit, limit, (fan_in, fan_out)) def kaiming_init(fan_in): std np.sqrt(2.0 / fan_in) return np.random.randn(fan_in, fan_in) * stdXavier 适合 tanh 和 sigmoid 激活函数Kaiming 适合 ReLU。为什么因为 ReLU 会把一半的神经元置零实际方差减半所以需要更大的初始方差来补偿。这个细节在论文里可能就一句话但实际工程里搞错了就是训练不收敛。还有一个数值稳定性的关键是梯度裁剪。训练 RNN 或深层 Transformer 时梯度爆炸是家常便饭。裁剪的逻辑很简单def clip_gradients(params, max_norm): total_norm np.sqrt(sum(np.sum(p.grad ** 2) for p in params if p.grad is not None)) if total_norm max_norm: scale max_norm / (total_norm 1e-6) for p in params: if p.grad is not None: p.grad * scalemax_norm一般设 1.0 或 5.0。我实测下来1.0 对 Transformer 比较稳5.0 对 RNN 更合适。这个值不是拍脑袋定的而是根据梯度范数的分布来调的——你可以先跑几百步打印梯度范数看看它的分布范围再决定裁剪阈值。4.3 训练循环的完整实现训练循环看起来简单但里面有很多细节。我把它拆成几个关键步骤第一步是前向传播。把输入数据喂给模型算出预测值。这一步要注意的是训练时和推理时的行为可能不同比如 dropout 只在训练时生效。第二步是计算 loss。分类任务用交叉熵回归任务用 MSE。交叉熵的实现要注意数值稳定性def cross_entropy(logits, targets): # logits: (batch, n_classes) # targets: (batch,) 整数标签 log_probs logits - np.max(logits, axis-1, keepdimsTrue) log_probs log_probs - np.log(np.sum(np.exp(log_probs), axis-1, keepdimsTrue)) batch_size logits.shape[0] return -np.mean(log_probs[np.arange(batch_size), targets])第三步是反向传播。调用loss.backward()算出所有参数的梯度。第四步是参数更新。用优化器SGD 或 Adam更新参数。Adam 的实现要注意动量项的偏差修正class Adam: def __init__(self, params, lr1e-3, betas(0.9, 0.999), eps1e-8): self.params params self.lr lr self.beta1, self.beta2 betas self.eps eps self.m [np.zeros_like(p.data) for p in params] self.v [np.zeros_like(p.data) for p in params] self.t 0 def step(self): self.t 1 for i, p in enumerate(self.params): if p.grad is None: continue self.m[i] self.beta1 * self.m[i] (1 - self.beta1) * p.grad self.v[i] self.beta2 * self.v[i] (1 - self.beta2) * (p.grad ** 2) m_hat self.m[i] / (1 - self.beta1 ** self.t) v_hat self.v[i] / (1 - self.beta2 ** self.t) p.data - self.lr * m_hat / (np.sqrt(v_hat) self.eps)实操心得训练循环里一定要加日志记录 loss、梯度范数、学习率。我当初就是没加梯度范数的日志训练崩了都不知道是梯度爆炸还是学习率太大。另外第一个 batch 的 loss 要特别关注如果它明显偏离预期比如分类任务应该是ln(n_classes)左右说明初始化或数据有问题。4.4 推理与采样策略训练完之后推理环节也有讲究。最基础的是贪心解码每次选概率最大的 tokendef greedy_decode(model, prompt, max_len): tokens tokenizer.encode(prompt) for _ in range(max_len): logits model.forward(tokens) next_token np.argmax(logits[-1]) tokens.append(next_token) if next_token tokenizer.eos_token_id: break return tokenizer.decode(tokens)但贪心解码有个问题输出会重复、很无聊。所以实际工程里常用的是温度采样和 top-k/top-p 采样。温度采样是在 softmax 之前把 logits 除以温度系数def sample_with_temperature(logits, temperature1.0): logits logits / temperature probs softmax(logits) return np.random.choice(len(probs), pprobs)温度小于 1 会让分布更尖锐更确定大于 1 会让分布更平坦更随机。top-k 采样是只从概率最高的 k 个 token 里采样top-p 是从累积概率达到 p 的最小集合里采样。我实测下来top-p 比 top-k 更自然因为它是动态调整候选集的不会在概率分布很尖锐时还硬塞 k 个候选。注意采样策略的选择要看任务。代码生成适合低温加 top-p创意写作适合高温加 top-p翻译适合贪心或低温。没有万能参数得根据实际效果调。5. 常见问题与排查技巧实录5.1 训练不收敛的排查清单训练不收敛是新手最常遇到的问题我整理了一个排查顺序按这个顺序查基本能定位到问题现象可能原因排查方法loss 一直是 NaN学习率太大、初始化有问题、除零打印每层激活值检查是否有 inf/nanloss 不下降学习率太小、梯度消失、数据有问题打印梯度范数检查数据标签是否正确loss 震荡学习率太大、batch size 太小降低学习率增大 batch sizeloss 下降后反弹过拟合、学习率调度不当加正则化检查验证集 loss训练 loss 降但验证 loss 不降过拟合加 dropout、weight decay、数据增强我当初遇到过一个很隐蔽的问题loss 在前 100 步正常下降然后突然变成 NaN。排查了半天发现是某个 batch 里有一个样本的输入全是 padding导致 LayerNorm 的方差为 0除零了。解决方案是在 LayerNorm 里加一个极小的 eps并且在数据预处理时过滤掉全 padding 的样本。5.2 显存溢出的定位与优化显存溢出在推理长文本时特别常见。定位方法很简单打印每一步的张量 shape 和显存占用。我一般会在关键位置加这样的日志def log_memory(tag): import psutil process psutil.Process() mem process.memory_info().rss / 1024 / 1024 print(f[{tag}] 内存占用: {mem:.1f} MB)优化显存的手段有几个KV Cache 用固定长度预分配、及时释放中间变量、用 float16 代替 float32、梯度检查点用时间换空间。其中梯度检查点最有效它不保存中间激活值反向传播时重新计算能把显存降到原来的 1/3 左右代价是训练速度慢 20% 到 30%。5.3 推理结果异常的调试思路推理结果异常比如输出重复、乱码、答非所问排查思路是从后往前查。先看采样策略是不是有问题温度太高、top-p 太大再看 logits 的分布是不是正常有没有某个 token 概率异常高再看 attention 的 mask 是不是对最后看分词器的编解码是不是一致。我遇到过一个很典型的问题模型输出总是重复同一句话。查了半天发现是 KV Cache 的索引没更新每次推理都在用同一段 cache相当于模型一直在看同样的上下文。这种问题在调包的时候根本不会遇到因为框架帮你处理了但手写的时候每一个索引都得自己管。5.4 梯度校验的实操方法梯度校验是验证自动微分实现是否正确的黄金标准。方法是对每个参数用数值梯度跟解析梯度对比def check_gradient(func, x, eps1e-5): # 数值梯度 num_grad np.zeros_like(x) for i in range(x.size): x_flat x.flatten() x_flat[i] eps f_plus func(x_flat.reshape(x.shape)) x_flat[i] - 2 * eps f_minus func(x_flat.reshape(x.shape)) num_grad.flatten()[i] (f_plus - f_minus) / (2 * eps) # 解析梯度通过反向传播算出来的 analytic_grad compute_analytic_grad(func, x) # 对比 diff np.abs(num_grad - analytic_grad) / (np.abs(num_grad) np.abs(analytic_grad) 1e-8) return np.max(diff)如果最大相对误差小于1e-5说明实现是对的。如果某个参数的误差特别大重点检查那个参数相关的操作。实操心得校验的时候要用双精度float64单精度的数值误差会干扰判断。另外eps不能太大也不能太小1e-5是个比较稳的值。6. 我踩过的那些坑与经验总结6.1 关于“从零”的度我一开始有点走极端想连矩阵乘法都手写结果写出来的版本比 NumPy 慢了几百倍调试的时候等得想砸键盘。后来想明白了从零的目的是理解原理不是重新发明轮子。矩阵乘法、卷积、BLAS 这些底层运算用 NumPy 就行没必要自己写。真正要手写的是那些体现 AI 工程核心思想的环节自动微分、注意力、分词、采样。这个度怎么把握我的标准是如果一个操作涉及“算法逻辑”而非“数值计算”就手写如果只是纯粹的数值计算就用库。比如 softmax 涉及数值稳定性处理手写矩阵乘法只是数值计算用 NumPy。按这个标准整个项目的代码量大概在 2000 行左右既能覆盖核心环节又不会陷入无意义的重复劳动。6.2 调试工具与技巧手写 AI 工程最痛苦的就是调试因为没有框架的报错信息帮你定位。我总结了几条实用的调试技巧第一单元测试要覆盖边界情况。比如 shape 为 1 的维度、全零输入、极大极小值。我当初就是没测全零输入导致 LayerNorm 除零的问题拖了两天才发现。第二可视化中间结果。attention 的权重矩阵、loss 曲线、梯度分布能画图就画图。人眼对图形的敏感度远高于数字很多异常一眼就能看出来。第三用 PyTorch 做对照。手写的模块跑出来的结果跟 PyTorch 对应模块的结果对比误差在1e-5以内就说明实现是对的。这个方法特别适合验证 attention 和 LayerNorm 这种容易写错的模块。第四小规模数据先跑通。不要一上来就用完整数据集训练先用几个样本跑通整个流程确认没有 shape 错误、没有 NaN再扩大规模。我当初就是急着上大数据集结果一个 shape 错误导致跑了半小时才报错浪费了大量时间。6.3 性能优化的取舍手写实现的性能肯定不如框架但这不是重点。不过有几个优化还是值得做的用 NumPy 的向量化操作代替循环、避免不必要的内存拷贝、用 in-place 操作减少内存分配。我实测下来向量化能带来几十倍的加速而 in-place 操作能省 30% 左右的内存。但要注意in-place 操作在自动微分里要慎用。因为反向传播需要用到前向传播的中间值如果你 in-place 修改了反向传播时拿到的就是被修改后的值梯度就错了。我当初为了省内存在激活函数里用了 in-place结果梯度全错查了一整天才发现。教训是前向传播的中间值除非你确定反向传播不需要否则不要 in-place 修改。6.4 后续可以扩展的方向这个项目跑通之后可以往几个方向扩展。一是加更多的层比如残差连接、dropout、位置编码把 Transformer 的完整结构搭出来。二是加更多的优化器比如 AdamW、Lion对比它们在不同任务上的表现。三是加分布式训练虽然手写分布式比较复杂但理解数据并行和模型并行的原理对工程能力提升很大。四是加量化把 float32 量化到 int8理解推理加速的底层逻辑。我个人最推荐的是先加位置编码和残差连接把 Transformer 的编码器搭完整然后在一个小规模的语言建模任务上训练看看能不能生成通顺的文本。这个过程会让你对整个 AI 工程的流程有完整的体感比看十篇论文都管用。最后分享一个我自己的体会手写一遍最大的收获不是代码本身而是对“为什么”的理解。为什么 LayerNorm 要减均值除标准差为什么 attention 要除以sqrt(d_k)为什么 softmax 要减最大值这些问题调包的时候你不会问但手写的时候你不得不回答。而正是这些回答构成了 AI 工程师的核心竞争力。
返回列表