ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

GPT-2模型结构拆解与最小实现:从数值对齐到推理优化

GPT-2模型结构拆解与最小实现:从数值对齐到推理优化

最近半年陆陆续续有人来问同一个问题:现在大模型迭代这么快,为什么还要回头去看GPT-2 模型结构?我的回答一直是——因为它是目前所有主流自回归语言模型里,结构最干净、最容易在一张消费级显卡上完整跑通、并且权重还是公开可下载的那一档。你想搞明白 FlashAttention 在优化什么、想搞明白 KV Cache 为什么能省算力、想搞明白 RoPE 到底替换掉了什么东西,最省事的路径就是先把 GPT-2 从头到尾拆一遍,跑通一次前向,做一次数值对齐。

这篇就按这个思路写。前半段讲结构,把十二层里每一块的作用、参数账本、初始化细节讲清楚,顺便把几张流传很广的结构图里画错或者画省事的地方指出来;后半段讲实践,包含一份两百行左右能跑通的最小实现、加载已有权重做数值对齐的具体做法(这里有一个极容易被踩的转置坑)、生成采样参数的取舍、显存估算公式,还有我在实际调试里遇到过的几类典型故障和排查顺序。

不管你是刚入门想搞懂注意力机制到底怎么算的,还是已经调过不少模型但一直没亲手写过一版前向、想知道c_attn这种命名背后是什么东西,这篇应该都能给你省几个小时。

1. 先把 GPT-2 的结构骨架摆清楚

1.1 为什么是纯 Decoder:一个被追问最多的选型问题

GPT-2 用的是一摞 Decoder 层,没有 Encoder,也没有跨注意力。这个选择放到今天看很自然,但当年其实是一道分叉路口。要理解它,得先想清楚目标任务是什么:给定前面的若干 token,预测下一个 token。这个任务的约束是单向的——预测第 t 个位置的时候,第 t+1 个位置的内容在推理阶段根本不存在。

原始的 Encoder-Decoder 架构里,Encoder 做双向编码,Decoder 做单向生成,中间用交叉注意力把两边接起来。这套东西在机器翻译那种"输入一段完整句子、输出另一段完整句子"的场景里非常合适,因为输入端天然是完整的。但一旦任务变成"续写一段文本",输入端和输出端就是同一段序列,你既不需要第二份完整输入,也不需要交叉注意力那一层额外的参数和计算。把 Encoder 砍掉、把交叉注意力砍掉,剩下的就是 Decoder 堆叠。

这个取舍带来的直接好处有三点。参数效率上,同样的层数能堆更多有效容量在生成路径上。推理效率上,每一步只需要处理新进来的那一个 token,前面算过的键值可以直接缓存复用,这就是后面 KV Cache 能成立的结构前提。训练目标简单上,整段序列只需要一个交叉熵损失,每个位置都是监督信号,不存在"只对输出部分算 loss"这种需要额外对齐的东西。

代价也很明确:它是单向的,任何需要"看完整段内容再做判断"的任务——比如给一句话做情感分类、从一段文本里抽实体——都需要额外的适配方式,最常见的做法是取最后一个位置的隐状态去做分类头,或者干脆在末尾拼一个特殊 token 让模型把答案"生成"出来。这点在你决定要不要用纯 Decoder 做下游任务时,必须先想明白。

1.2 十二层里到底装了什么

小尺寸这版的配置是:层数 12,隐藏维度 768,注意力头数 12,每头维度 64,前馈中间层维度 3072,最大位置数 1024,词表大小 50257。把这几个数字记牢,后面所有形状推导都靠它们。

每一层内部是标准的两段式子:先过一个 LayerNorm,进多头因果自注意力,加残差;再过一个 LayerNorm,进前馈网络,加残差。注意归一化是放在子层之前的,也就是所谓的 Pre-LN。这个位置很关键,Post-LN(归一化放在残差相加之后)在深层网络里如果不做精心的学习率预热,很容易训崩,梯度在反向传播时会被反复的归一化缩放搅乱;Pre-LN 让残差通路上是一条干干净净的恒等映射,梯度可以直通到底,训练稳定性提升非常明显。代价是表达能力上略微弱一点,实践中通常通过最后再加一个ln_f来补偿——这也是为什么 GPT-2 在第十二层之后还有一层 LayerNorm,很多人第一次看结构图会漏掉它。

前馈网络是两层线性加一个 GELU,中间维度扩到 4 倍再压回来。这个 4 倍是个经验值,来自于早期 Transformer 的设定,小尺寸这里就是 768 → 3072 → 768。这两层线性承担了整层里约三分之二的参数,是模型存储知识的主要位置,后面做 LoRA 之类的低秩适配,一般都挂在这两个矩阵上。

注意力部分有一个实现上的特殊设计:查询、键、值三个投影矩阵没有分成三个独立的线性层,而是拼成了一个c_attn层,一次算出三倍维度的输出再切开。这么做在数学上完全等价,但在 GPU 上把三次小矩阵乘法合并成一次大矩阵乘法,kernel 启动次数少了,显存访问也更连续,实测能有一点点吞吐收益。你如果自己按论文公式写成三个nn.Linear,结果是对的,但加载现成权重的时候要额外处理切分。

1.3 参数账本:1.24 亿是怎么凑出来的

很多人看论文表格记的是 117M,但实际把权重拉下来数一遍,是 1.24 亿左右。差异来自统计口径——词表大小、位置嵌入算不算进去,得出的数会差好几个百分点。与其记那个数,不如自己会算。下面这张表是按小尺寸配置逐项拆出来的:

模块计算式参数量
词嵌入50257 × 76838,597,376
位置嵌入1024 × 768786,432
每层注意力(含偏置)768×2304 + 2304 + 768×768 + 7681,772,352
每层前馈(含偏置)768×3072 + 3072 + 3072×768 + 7684,722,432
每层两个 LayerNorm768×2×23,072
单层小计—约 7,087,872
12 层合计7,087,872 × 1285,054,464
末尾 LayerNorm768×21,536
输出层与词嵌入共享0
总计—约 1.244 亿

从这个表里能读出两件在实践中有用的事。第一,词嵌入占了将近三成参数,而 50257 这个偏大的词表里,有相当一部分低频词对应的嵌入向量在通用语料上其实训得并不充分,做领域微调时这部分是典型的"占显存不干活"的区域。第二,输出层和词嵌入共享权重,这一手省掉了 3859 万参数,差不多相当于省掉三分之一的模型体积,同时也让输入表示和输出打分处在同一个语义空间里,收敛会更快一些。代价是 embedding 矩阵同时要满足两个方向的梯度,早期训练时偶尔会出现更新拉扯,不过在小模型上一般不成为问题。

1.4 三处容易被结构图误导的细节

网上流传的结构图大多为了排版好看做了简化,有三处我见过不少人被带偏。

第一处,位置编码不是正弦函数。原始 Transformer 用的是固定公式生成的位置编码,而 GPT-2 用的是一张可学习的嵌入表,长度 1024,每个位置一个 768 维向量,跟着训练一起更新。这意味着它没法直接外推到超过 1024 的长度——超出部分根本没有对应的行。这一点在做长文本时是硬约束,不是调参能解决的。

第二处,注意力掩码的填充值不是负无穷。理论上是把未来位置的概率压到 0,实现上要用一个极大的负数去覆盖 softmax 的输入。原作者用的是 -1e4 而不是-inf,原因是混合精度下-inf参与后续运算容易产生 NaN,尤其在低精度累加的时候。你自己实现时如果图省事写float('-inf'),fp32 下一般没事,切到 fp16 就可能在某几个 batch 上突然炸成 NaN,而且这种偶发故障极难定位。

第三处,残差投影层的初始化被特意缩小了。注意力和前馈各自的输出投影层,权重标准差不是常规的 0.02,而是 0.02 除以层数两倍的平方根。理由是残差通路上每层都会叠加一次输出,层数越深,累加后的方差越大,把每个子层的输出幅度按层数缩放回去,深层模型的激活值分布才不会随深度膨胀。你如果是加载现成权重,这条无所谓;但如果打算从零训一个小模型,漏掉这个缩放,十几层之后就很难收敛了。这个细节在论文正文里没有,藏在参考实现里。

2. 一条完整数据通路:从 token id 到 logits

2.1 词表、BPE 与可学习位置嵌入

输入第一步是把文本变成整数。GPT-2 用的是字节级 BPE,词表 50257,拆开来看是 256 个基础字节 token、50000 个合并出来的子词、外加 1 个文本结束符。字节级的好处是永远不会出现未登录词——任何 Unicode 字符最终都能拆成字节序列表示,哪怕是一门模型完全没见过的语言,也只会被拆得很碎,但不会报错。这个设计在当年是很务实的选择,代价是非拉丁语系的文本 token 效率偏低,中文一句话往往要拆出比英文多不少的 token 数。

拿到 token id 序列之后,查嵌入表得到[批大小, 序列长度, 768]的稠密表示,再加上位置嵌入。这里有个实现细节值得注意:位置嵌入的索引是从 0 开始按序列内偏移取的,不管你实际在做第几轮对话、上下文拼了多长,只要当前喂进去的序列长度是 T,位置就是 0 到 T-1。做多轮会话拼接时如果搞错了偏移,模型会以为每句话都从位置 0 开始,输出会明显变得不连贯。这也是后来各种位置编码改进方案想解决的核心痛点之一。

加完嵌入之后通常会接一个 dropout,训练时开、推理时关。这个参数在小模型微调时经常被忽略,默认值 0.1 在数据量小的时候其实是有效的正则手段,数据量上来了再调到 0 或者 0.05。

2.2 合并 QKV 的线性层

进入注意力子层,先做 LayerNorm,然后一次线性变换得到三倍维度的张量,形状是[B, T, 2304]。切开的方式是按最后一维三等分,前 768 是查询,中间 768 是键,后 768 是值。切完之后要做一个 reshape 加转置,把隐藏维度拆成"头数 × 每头维度",变成[B, 12, T, 64]。

这个拆分顺序是有讲究的。常见的错误写法是先 reshape 成[B, T, 12, 64]再直接转置成[B, 12, T, 64],看起来对,但如果切分时的维度顺序和权重排布不一致,结果会错得很隐蔽——损失能降,生成也像模像样,就是质量比预期差一截。判断方法很简单:跟参考实现做一次数值对比,最大绝对误差应该在 1e-5 量级(fp32),对不上就是这里出问题了。

缩放因子是每头维度的平方根,这里是 8。用每头维度而不是总维度,是因为 softmax 的输入尺度取决于单个头的点积长度,跟头数无关。写错成 768 的平方根,注意力分布会变得极其平缓,几乎等价于均匀加权,模型基本学不到东西。

2.3 因果掩码与前缀复用

打分矩阵形状是[B, 12, T, T],第 i 行第 j 列表示第 i 个位置对第 j 个位置的关注强度。因果掩码是一个下三角矩阵,把 j 大于 i 的位置全部压掉。这一步骤在训练时是必需的,因为你一次性喂进整段序列,如果不加掩码,第 3 个位置就能看到第 10 个位置的内容,等于作弊,训练损失会降得很好看但推理时完全不能用。

推理时的情况不一样。自回归生成每次只进来一个新 token,T 等于 1,理论上不需要掩码。但为了复用前面算过的键和值,实践中会把历史键值缓存起来,这时候新 token 的查询要跟全部历史键做点积,掩码就退化成"全部可见",反而不需要了。所以真正高效的实现里,带缓存的那条路径是跳过掩码计算的——少一次[1, 12, 1, T]的 masked_fill,长序列下省下来的时间不算小。

缓存本身占多少显存,值得单独算一下。每一层要存键和值两份,形状都是[B, 12, T, 64],fp16 下每层大约B × T × 768 × 2 × 2字节。批大小 1、长度 1024 的时候,单层约 3MB,12 层加起来约 36MB,完全不是负担;但如果批大小开到 32,长度拉到 1024,就是 1.1GB 左右,这时候缓存反而成了显存大头。这个增长是线性的、跟批大小和长度成正比的关系,规划并发时必须提前算进去。

2.4 FFN、GELU 与残差缩放初始化

注意力输出经过投影层回到 768 维,加上残差,进入前馈子层。前馈这里用的是 GELU 而不是 ReLU。区别在于 ReLU 在负半轴直接截断成 0,梯度在那里是断的;GELU 是一条平滑曲线,负值区间保留了一小部分非零输出,梯度也更连续。工程上感知最明显的是:用 GELU 的模型在同样的学习率下收敛更稳,不容易出现某些神经元彻底"死掉"再也不更新的情况。

还有一个容易忽略的点:GELU 有两个版本,精确版用误差函数算,近似版用双曲正切拟合。两者数值上差异很小,但不是完全相等。加载现成权重时如果版本用错,单层误差看不出来,十二层累加之后 logits 的偏差可能会到 1e-3 量级,做数值对齐验证时会误以为自己哪里写错了。稳妥的做法是明确指定用近似版,跟参考实现保持一致。

前馈的输出投影层同样要套那个按层数缩放的初始化。整层的残差顺序是:x = x + attn(ln_1(x)),然后x = x + mlp(ln_2(x))。注意两个 LayerNorm 的输入都是相加之前的 x,不是相加之后的,写反了在浅层还勉强能跑,深层直接不收敛。

2.5 权重共享的输出层与形状总表

十二层走完,过末尾的 LayerNorm,然后接输出投影。这一层没有独立参数,直接用词嵌入矩阵的转置,把 768 维映射回 50257 维。得到 logits 之后,训练时接交叉熵,推理时接采样。

把所有中间形状整理成一张表,写代码时对着看能省不少调试时间:

阶段张量形状说明
输入 token id[B, T]整数,取值 0~50256
词嵌入 + 位置嵌入[B, T, 768]两者直接相加
LayerNorm 输出[B, T, 768]eps 取 1e-5
QKV 合并输出[B, T, 2304]一次线性得到
拆分并分头[B, 12, T, 64]三个张量同形状
注意力打分[B, 12, T, T]除以 8 后过掩码
注意力输出[B, 12, T, 64]→[B, T, 768]合并头后投影
前馈中间层[B, T, 3072]GELU 激活
单层输出[B, T, 768]残差相加
最终 logits[B, T, 50257]与词嵌入共享权重

最后一行值得多看一眼:序列长度 1024、批大小 8 的时候,logits 张量是8 × 1024 × 50257,fp32 下约 1.6GB。这个中间结果在很多实现里是显存峰值的主要来源,尤其是训练时还要保留它算梯度。优化手段要么是分块算损失,要么是换成低精度。

3. 动手复现:两百行代码跑通前向

3.1 环境与依赖

只需要 PyTorch 和 transformers 两个包,前者做张量运算,后者用来拉现成权重和分词器做对齐验证。版本上建议 PyTorch 2.x 配 transformers 4.3x 以上,主要是新版在scaled_dot_product_attention上有融合优化,虽然我们自己的实现用不上,但对照组的推理速度会快很多。

pip install torch transformers

显卡不是必需品,小尺寸这版在 CPU 上跑单条推理也就几百毫秒,做数值验证完全够用。但如果要跑批、要试生成,还是建议有一张 8GB 以上的卡,fp16 下权重只占 250MB 左右,剩下的显存留给缓存和中间激活。

3.2 核心模块实现

先写注意力。关键点都在注释里标出来了:

import math import torch import torch.nn as nn import torch.nn.functional as F class CausalSelfAttention(nn.Module): def __init__(self, d_model=768, n_head=12, max_len=1024): super().__init__() assert d_model % n_head == 0 self.n_head = n_head self.d_head = d_model // n_head self.c_attn = nn.Linear(d_model, 3 * d_model) self.c_proj = nn.Linear(d_model, d_model) # 注意用 -1e4 而不是 -inf mask = torch.tril(torch.ones(max_len, max_len)).view(1, 1, max_len, max_len) self.register_buffer("bias", (1.0 - mask) * -1e4) def forward(self, x, kv_cache=None, use_cache=False): B, T, C = x.shape qkv = self.c_attn(x) # [B, T, 3C] q, k, v = qkv.split(C, dim=2) q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2) k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2) v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2) if kv_cache is not None: pk, pv = kv_cache k = torch.cat([pk, k], dim=2) v = torch.cat([pv, v], dim=2) att = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_head) if T > 1: L = att.size(-1) att = att + self.bias[:, :, L - T:L, :L] att = F.softmax(att, dim=-1) y = att @ v y = y.transpose(1, 2).contiguous().view(B, T, C) out = self.c_proj(y) return (out, (k, v)) if use_cache else out

然后是整层和整个模型:

class Block(nn.Module): def __init__(self, d_model=768, n_head=12, d_ff=3072): super().__init__() self.ln_1 = nn.LayerNorm(d_model, eps=1e-5) self.attn = CausalSelfAttention(d_model, n_head) self.ln_2 = nn.LayerNorm(d_model, eps=1e-5) self.mlp = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(approximate="tanh"), # 必须和参考实现一致 nn.Linear(d_ff, d_model), ) def forward(self, x, kv_cache=None, use_cache=False): if use_cache: a, cache = self.attn(self.ln_1(x), kv_cache, True) x = x + a else: x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return (x, cache) if use_cache else x class GPT2(nn.Module): def __init__(self, vocab=50257, n_layer=12, d_model=768, n_head=12): super().__init__() self.wte = nn.Embedding(vocab, d_model) self.wpe = nn.Embedding(1024, d_model) self.blocks = nn.ModuleList([Block(d_model, n_head) for _ in range(n_layer)]) self.ln_f = nn.LayerNorm(d_model, eps=1e-5) def forward(self, idx): B, T = idx.shape pos = torch.arange(T, device=idx.device) x = self.wte(idx) + self.wpe(pos) for blk in self.blocks: x = blk(x) x = self.ln_f(x) return x @ self.wte.weight.T # 权重共享

写完这段,最直观的感受应该是:抛开工程优化,GPT-2 的计算图短得惊人。真正复杂的不是结构本身,而是围绕它做训练稳定性、显存、吞吐的一整套配套工程。

3.3 加载已有权重做数值对齐

这一步是整个实践里最容易卡住的地方,核心是一个转置问题。参考实现是从另一个深度学习框架迁移过来的,里面的线性层沿用了那个框架的Conv1D写法,权重形状是[输入维度, 输出维度];而 PyTorch 的nn.Linear是[输出维度, 输入维度]。名字都叫 weight,形状正好反过来。

所以加载权重的时候,除了 LayerNorm 和嵌入层,所有线性层的权重都要转置:

import torch from transformers import GPT2LMHeadModel ref = GPT2LMHeadModel.from_pretrained("gpt2").eval() mine = GPT2().eval() sd = ref.state_dict() mapping = { "wte.weight": "wte.weight", "wpe.weight": "wpe.weight", "ln_f.weight": "ln_f.weight", "ln_f.bias": "ln_f.bias", } new_sd = {k: sd[v].clone() for k, v in mapping.items()} for i in range(12): p = f"h.{i}." q = f"blocks.{i}." for name in ["ln_1", "ln_2"]: new_sd[q + name + ".weight"] = sd[p + name + ".weight"].clone() new_sd[q + name + ".bias"] = sd[p + name + ".bias"].clone() for src, dst in [("attn.c_attn", "attn.c_attn"), ("attn.c_proj", "attn.c_proj"), ("mlp.c_fc", "mlp.0"), ("mlp.c_proj", "mlp.2")]: w = sd[p + src + ".weight"] new_sd[q + dst + ".weight"] = w.t().contiguous() # 关键转置 new_sd[q + dst + ".bias"] = sd[p + src + ".bias"].clone() mine.load_state_dict(new_sd, strict=True)

加载完做对比,取一段固定输入,看两边 logits 的最大绝对误差。fp32、CPU 上正常情况下应该在 1e-5 到 1e-4 之间,浮点运算顺序不同必然有微小差异,不用追求完全相等。如果误差在 1e-2 量级,八成是 GELU 版本用错了;如果误差在 1 以上,基本可以确定是转置没做对,或者c_attn切开后的顺序搞反了。

提示:验证时务必用torch.no_grad()并且把两边都设成eval()。忘记切换训练模式的话,dropout 会随机失活一部分神经元,每次跑出来的误差都不一样,会让你误以为实现有 bug。

3.4 生成循环与采样参数

对齐通过之后就可以写生成循环了。核心逻辑很简单:每步取最后一个位置的 logits,处理成概率分布,采一个 token,拼回序列,直到遇到结束符或到达长度上限。真正决定输出质量的是中间那几个参数。

参数典型取值作用与副作用
temperature0.7 ~ 1.0除以 logits 后再 softmax,小于 1 让分布更尖锐,大于 1 更随机
top_k40 ~ 50只保留概率最高的 k 个候选,k 太大等于没截断
top_p0.9 ~ 0.95按累积概率截断,动态候选数,通常比 top_k 更稳
repetition_penalty1.0 ~ 1.2对已出现 token 的对数概率做惩罚,超过 1.3 会明显影响通顺度
no_repeat_ngram_size2 ~ 3硬性禁止重复的 n 元组,长文本生成里很管用
max_new_tokens按场景定注意包含结束符,设太小会被硬截断在半句话上

实际调的时候,temperature和top_p一起改往往会互相打架。我的经验是先固定temperature=1.0,只用top_p调节随机性,等整体感觉对了,再微调温度去处理"太死板"或者"太发散"的问题。repetition_penalty属于救急手段,能不开就不开,它对语法结构的破坏比对内容重复的修复更明显。

还有一个批处理时的坑:GPT-2 没有专门的填充 token,分词器默认的pad_token_id是空的。批量生成时短句要补齐到同一长度,补齐位置如果不设成结束符,模型会把填充内容当成真实上下文,输出会莫名其妙地带上前一句的尾巴。正确做法是把填充 token 设成结束符 id,同时传好 attention mask,让模型知道哪些位置是补出来的。

4. 微调与部署:显存、吞吐和实际取舍

4.1 显存怎么估

显存估算这件事,掌握一个粗略公式就够用了。训练时的占用大致是:模型参数、梯度、优化器状态、激活值四块。

参数按 fp32 算是1.24 亿 × 4 字节 ≈ 500MB;梯度同量级,再来 500MB;如果用 Adam,动量和方差各一份,又是 1GB。这三块加起来已经 2GB 左右,跟批大小无关,是硬开销。所以哪怕批大小设成 1,训练也要准备 2GB 以上的余量。

激活值才是随配置线性增长的部分。粗略估法是批大小 × 序列长度 × 隐藏维度 × 层数 × 系数,系数取决于要不要重算、用了多少融合算子,经验值在 10 到 20 之间。小尺寸、批大小 8、长度 512 的配置下,实际占用大约 1GB 出头。再加上注意力打分矩阵那一块批大小 × 头数 × 长度²× 2 字节,长度 512 时这一项约 50MB,还不太显眼;长度拉到 2048,就变成 800MB 了,平方增长是长序列训练最先撞墙的地方。

推理侧就宽松得多,只剩参数和缓存两块。fp16 权重 250MB,缓存按前面算的公式加。一张 8GB 的卡跑小尺寸推理,批大小开到 16、长度 512,占用不到 2GB,余量很充裕。

4.2 小样本微调的关键参数

在几千条领域数据上做微调,几个参数的经验值可以直接抄:

  • 学习率从5e-5起步,比从头训练小一个量级以上。太大容易把预训练学到的通用能力冲掉,表现为通用任务上突然变差。
  • 批大小在显存允许范围内尽量大,配合梯度累积凑到等效 32 左右比较稳。
  • 训练轮数控制在 3 到 5 轮。小数据上多轮训练过拟合非常快,第三轮之后验证损失往往就回升了。
  • 学习率预热比例 6% 左右,后面接线性衰减。这个配置对 Pre-LN 结构不是必需的,但能平滑早期的不稳定。
  • 权重衰减给 0.01,LayerNorm 和偏置项不参与衰减。这条几乎每个实现都会处理,但自己写循环时经常漏。

判断是否过拟合,别只看训练损失,那个数字一定会降。看验证集上的困惑度,或者更实用的做法是定期让模型生成几条样本人工看——几百条数据训 5 轮之后,模型往往会开始复读训练集里的原句,这是过拟合最直观的信号。

注意:微调之后词嵌入层会被大幅更新,但输出层和它共享权重,这意味着原来词表里那些低频词的表征会被一起带动。如果你的领域语料很小、词汇又集中,一个副作用是模型对通用词的处理能力下降。缓解办法是给嵌入层单独设一个更小的学习率,或者干脆冻结它只训练上层。

4.3 推理侧加速

不换模型的前提下,能拿到明显收益的手段按投入产出比排序大致是这样。

第一是 KV Cache,前面写过实现。没有它的话,生成长度 200 的文本要重复计算前面所有位置,复杂度是平方级;加上之后变成线性。这是单项收益最大的一步,而且不改任何数值结果,属于纯赚。

第二是低精度推理,fp16 或 bf16 直接砍掉一半权重显存和一半带宽。小尺寸模型对精度不敏感,实测生成质量肉眼几乎看不出差异。注意的是 bf16 动态范围大,比 fp16 更不容易溢出,如果硬件支持优先选它。

第三是批处理,把多条请求拼成一个 batch 送进去。GPU 在小 batch 下利用率很低,批大小从 1 提到 8,吞吐往往能翻好几倍,而单条延迟只增加一点点。代价是必须处理变长补齐和缓存管理,工程复杂度上来了。

再往下的量化、算子融合、编译优化,收益开始取决于具体场景,投入产出比要自己算。我的一般建议是先把缓存和批处理做扎实,这两块没做好之前,折腾量化的收益往往被前者的浪费抵消掉了。

4.4 上线前必须确认的几件事

第一,长度上限要显式校验。位置嵌入只有 1024 行,输入超了不会报错,会直接索引越界或者悄悄回绕,行为不可预期。在接口层做硬校验,超长直接拒掉或者截断。

第二,结束符要正确配置。分词器的结束符 id 是 50256,生成时必须把它加进终止条件,否则模型会一直生成到长度上限才停,白白浪费算力,输出结尾还经常是半句废话。

第三,并发下的缓存管理。多条请求共享一个模型实例时,缓存必须按请求隔离,别图省事用全局变量。我见过一个线上故障就是缓存串了,A 用户的问题被拼到了 B 用户的回答开头。

第四,输入预处理要和训练时一致。标点、空格、大小写处理如果和预训练语料的处理方式差太远,模型表现会明显下滑。这一条听起来理所当然,但实践中出问题的频率相当高。

5. 排查实录:那些生成异常和数值对不上的坑

5.1 生成重复、跑飞、提前截断

输出整段复读同一句,最常见的原因不是模型坏了,而是解码策略太保守。temperature设得极低、top_k设成 1(等价于贪心解码)的时候,模型会陷入一个自我强化的循环:生成了某个短语,下一位置的最高概率又指向同一个短语,如此往复。解法是把温度抬到 0.7 以上,同时启用top_p。如果还是复读,再加一点重复惩罚,但别超过 1.2。

输出到一半突然变成乱码或者无关语言的字符,通常是数值溢出。fp16 下 logits 有可能超出表示范围变成 inf,softmax 之后整行变成 NaN,采样就取到了词表里完全无关的 token。切 bf16 或者保持 fp32 能解决大部分。

生成被硬截断在半句话上,检查max_new_tokens是否包含结束符本身,以及终止条件里是不是把结束符漏掉了。还有一种情况是流式返回时提前关闭了连接,这个得看服务端日志才能定位。

同一个 prompt 每次输出差异极大,先确认是不是temperature太高,再确认模型是否还在训练模式。忘记eval()的话 dropout 会一直生效,每次前向的隐状态都不一样,输出自然发散得厉害。

5.2 数值对不齐的排查顺序

自己实现和参考实现对比对不上,别急着怀疑整体架构,按这个顺序查基本能定位:

  1. 先看转置。所有线性层的权重是否做了转置,这是出错概率最高的一项,症状是误差大到离谱但形状完全对得上,load_state_dict(strict=True)也不会报错。
  2. 再看 GELU 版本。近似版和精确版的差异会累积,症状是误差在 1e-3 量级,看起来像浮点误差但其实不是。
  3. 查掩码填充值。用-inf的时候如果输入里有全被屏蔽的行(补齐位置常见),softmax 会输出 NaN,症状是某些 batch 突然全 NaN。
  4. 查c_attn切分顺序。是不是 q、k、v 的顺序切反了。切反之后模型还能跑,损失也能降一点,但生成质量明显差,这个最难发现。
  5. 查 LayerNorm 的 eps。默认值和参考实现可能不同,1e-5 和 1e-6 在数值上会有可观测的差异。
  6. 查位置索引的起点。是不是从 1 开始取的,差一位在小序列上影响不大,长序列上会明显劣化。

5.3 速查表

现象最可能原因处理方式
加载权重不报错但输出乱码线性层未转置对四个线性层权重做.t()
误差稳定在 1e-3GELU 用了精确版改成approximate="tanh"
偶发 NaN掩码用了负无穷换成 -1e4
生成复读温度过低或贪心解码温度 0.7 + top_p 0.9
批量生成串味填充 token 未设置pad 设为结束符 id 并传 mask
长文本后半段质量崩塌超出位置嵌入范围硬截断到 1024 以内
微调后通用能力骤降学习率过大或轮数过多降到 5e-5、3 轮以内
显存随长度暴涨注意力矩阵平方增长缩短长度或换分块注意力

我自己的习惯是,每次改完实现,先跑一个固定种子的最小对比脚本,把 logits 的最大误差和几个位置的 argmax 一起打出来。花了五分钟写这个脚本,后面每次重构都能省下几小时的瞎猜。踩过几次坑之后你会发现,这类数值问题基本都能靠一次严格的对齐验证拦在门外,真正难缠的反而都是配置层面的东西——填充、截断、模式切换,这些不会报错,只会让结果悄悄变差。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进