简介：面向自然语言处理学习者与PyTorch实践者的BERT实现代码包，对应Google AI 2018年提出的深度双向Transformer预训练模型论文，适合希望从源码层面理解BERT结构、在本地复现语言模型训练流程的中级开发者参考。压缩包共33个文件，以27个Python脚本为主体，覆盖模型定义、数据集处理与训练器模块，另含Makefile、requirements.txt、配置文件及README说明，整体仅28KB，轻量易读，代码附带简单注释，便于逐行梳理。项目当前仍在开发中，代码尚未经过完整验证，读者可将其作为学习Transformer编码器与预训练思路的入门素材，结合语料准备与快速开始说明动手调试，理解自注意力机制在语言模型中的应用方式。已有879人学习关注，适合作为论文阅读之外的代码对照参考。

1. BERT 的“双向”到底双向在哪：一个被误读最多的结论

很多人第一次接触 BERT，会以为“双向”就是把句子倒过来再跑一遍，或者把左右两个 LSTM 拼在一起。真正翻开源码才会发现，BERT 的双向不是结构上的双向，而是训练目标上的双向：它把一句话里 15% 的 token 挖掉，让模型同时利用左右两侧的上下文去猜这些被挖掉的位置。这个设计直接绕开了自回归语言模型“只能看左边”的限制，也让 Transformer 的并行能力第一次在预训练阶段被彻底释放。搞懂这一点，才能理解为什么 BERT 的输入要三段拼接、为什么要有 mask、为什么位置编码是学出来的而不是正余弦。这份 PyTorch 实现把论文里的每个模块拆成了可读的类，从 Embedding 到 Transformer 堆叠再到两个预训练头，逐层对应原论文的公式，是少见的能把“机制”和“代码”对上号的资源。

2. 输入表示：为什么一条句子要拆成三段 Embedding

2.1 Token、Segment、Position 三路相加的动机

BERT 的输入不是简单地把词转成 id 就完事。源码里BertEmbeddings这个类做了三件事：查词表得到 token embedding、查句子段得到 segment embedding、按位置查 position embedding，然后把三者逐元素相加。为什么要加而不是拼接？因为拼接会让维度膨胀，后续每一层注意力的计算量都跟着涨；相加则保持维度不变，同时让三种信息在同一个向量空间里叠加。这个选择背后有个隐含假设：三种信息是正交的、可以线性叠加的。实际训练中这个假设并不严格成立，但实验证明够用，这也是深度学习里常见的“工程上先跑通再说”的取舍。

class BertEmbeddings(nn.Module): def __init__(self, vocab_size, hidden_size, max_position, type_vocab_size): super().__init__() # 词表嵌入：把 token id 映射到 hidden_size 维向量 self.word_embeddings = nn.Embedding(vocab_size, hidden_size) # 位置嵌入：可学习参数，长度固定为 max_position self.position_embeddings = nn.Embedding(max_position, hidden_size) # 段嵌入：区分句子 A 和句子 B，通常只有 0/1 两种 self.token_type_embeddings = nn.Embedding(type_vocab_size, hidden_size) self.LayerNorm = nn.LayerNorm(hidden_size, eps=1e-12) self.dropout = nn.Dropout(0.1) def forward(self, input_ids, token_type_ids): seq_len = input_ids.size(1) # 位置 id 从 0 到 seq_len-1，固定生成，不参与梯度 position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device) position_ids = position_ids.unsqueeze(0).expand_as(input_ids) words_emb = self.word_embeddings(input_ids) pos_emb = self.position_embeddings(position_ids) type_emb = self.token_type_embeddings(token_type_ids) # 三路相加后统一做 LayerNorm，稳定数值分布 embeddings = words_emb + pos_emb + type_emb return self.dropout(self.LayerNorm(embeddings))

这段代码里有个容易被忽略的细节：position_ids是用torch.arange现场生成的，而不是作为输入传进来。这意味着位置信息完全由序列长度决定，和具体是哪个句子无关。好处是实现简单，坏处是模型学到的位置是“绝对位置”，换一个更长的序列就得重新学。后来的 RoPE 之类方案就是冲着这个痛点去的，但在这份实现里，绝对位置编码已经够用，因为预训练的最大长度是固定的 512。

2.2 特殊 token 的布局与 mask 的生成

输入序列的布局是[CLS] + 句子A + [SEP] + 句子B + [SEP]。[CLS]这个位置的最终隐状态会被拿去做下一句预测任务，[SEP]用来分隔两个句子。源码里BertModel的 forward 会同时返回sequence_output和pooled_output，后者就是[CLS]位置经过一个线性层加 tanh 的结果。attention mask 的生成逻辑是：真实 token 位置为 1，padding 位置为 0，然后在注意力计算时把 0 的位置填成负无穷，softmax 之后这些位置的权重就趋近于零。这个机制保证了 padding 不会污染真实 token 的表示，是批处理能跑起来的前提。

提示：mask 的维度要和注意力分数的维度对齐，常见错误是把[batch, seq_len]的 mask 直接加到[batch, heads, seq_len, seq_len]的分数上，广播规则不对就会静默出错，loss 看起来在降但模型学不到东西。

3. Transformer 编码层：自注意力的实现细节与数值稳定性

3.1 多头注意力为什么要除以 sqrt(d_k)

BertSelfAttention里最核心的一行是attention_scores = attention_scores / math.sqrt(self.attention_head_size)。这个缩放因子不是拍脑袋定的。点积的方差随维度线性增长，维度越大，点积结果的绝对值越大，softmax 之后分布越尖锐，梯度越容易消失。除以sqrt(d_k)相当于把方差拉回 1 附近，让 softmax 工作在梯度比较健康的区间。这是 Transformer 能堆深的前提之一，也是很多复现翻车的地方——忘了除或者除错维度，训练前期 loss 会卡住不动。

class BertSelfAttention(nn.Module): def __init__(self, hidden_size, num_heads): super().__init__() self.num_heads = num_heads self.head_dim = hidden_size // num_heads # 一个线性层同时产出 Q、K、V，减少参数和计算 self.query = nn.Linear(hidden_size, hidden_size) self.key = nn.Linear(hidden_size, hidden_size) self.value = nn.Linear(hidden_size, hidden_size) def forward(self, hidden_states, attention_mask): batch, seq_len, _ = hidden_states.size() # 拆头：[batch, seq, hidden] -> [batch, heads, seq, head_dim] q = self.query(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = self.key(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v = self.value(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 缩放点积，防止 softmax 饱和 scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim) # mask 位置填极小值，softmax 后权重趋零 scores = scores + attention_mask probs = torch.softmax(scores, dim=-1) context = torch.matmul(probs, v) # 拼回头：[batch, heads, seq, head_dim] -> [batch, seq, hidden] context = context.transpose(1, 2).contiguous().view(batch, seq_len, -1) return context

view加transpose这套组合是 PyTorch 里拆头的标准写法，但要注意transpose之后张量在内存里不连续，直接view会报错，所以代码里用了contiguous()。这个细节在调试时经常被忽略，报错信息也不直观，属于典型的血泪经验。

3.2 残差连接与 LayerNorm 的位置

每个子层外面都包了一层残差：output = LayerNorm(x + Sublayer(x))。这个顺序叫 Post-LN，是原论文的写法。残差的作用是给梯度一条高速公路，让深层网络不至于梯度消失；LayerNorm 放在残差之后，负责把数值拉回稳定区间。后来的研究指出 Pre-LN（先 Norm 再进子层）训练更稳定，不需要 warmup 也能收敛，但这份实现忠于原论文，用的是 Post-LN。理解这个顺序的差别，对排查“训练到一半 loss 突然爆炸”这类问题很有帮助——Post-LN 在深层堆叠时对学习率更敏感，warmup 不是可选项而是必需品。

3.3 前馈网络与激活函数的选择

BertIntermediate和BertOutput组成前馈网络，中间维度是hidden_size * 4，激活函数用的是 GELU 而不是 ReLU。GELU 在零点附近是平滑的，梯度不会像 ReLU 那样在负半轴直接归零，实验上对语言模型更友好。这个选择在 2018 年还算新鲜，现在已经成了标配。前馈网络的作用是给每个位置独立做一次非线性变换，和注意力负责的“跨位置信息交换”形成互补：注意力管横向，前馈管纵向。

4. 预训练任务：MLM 与 NSP 的损失怎么算

4.1 Masked Language Model 的 15% 与 80/10/10

MLM 的流程是：随机选 15% 的 token 位置，其中 80% 替换成[MASK]，10% 替换成随机 token，10% 保持不变。为什么要这么绕？如果全换成[MASK]，模型只在看到[MASK]时才学习预测，而下游微调时根本没有[MASK]，会造成预训练和微调的不一致。加入随机替换和保持不变，是让模型对“输入可能是噪声”这件事有鲁棒性。源码里BertPreTrainingHeads的predictions是一个hidden_size -> vocab_size的线性层，损失只在被选中的 15% 位置上计算，其余位置忽略。这个“只在部分位置算 loss”的实现，是很多人复现时对不齐论文指标的原因——如果把所有位置都算进去，loss 数值会小很多，但模型学到的信号被稀释了。

def mask_tokens(inputs, tokenizer, mlm_prob=0.15): labels = inputs.clone() # 按 15% 概率选出要预测的位置 prob_matrix = torch.full(labels.shape, mlm_prob) masked_indices = torch.bernoulli(prob_matrix).bool() # 未被选中的位置 label 设为 -100，交叉熵会忽略 labels[~masked_indices] = -100 # 80% 换成 [MASK] indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices inputs[indices_replaced] = tokenizer.convert_tokens_to_ids("[MASK]") # 10% 换成随机 token indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long) inputs[indices_random] = random_words[indices_random] # 剩下 10% 保持不变 return inputs, labels

注意这里-100是 PyTorchCrossEntropyLoss的默认忽略值，不是随便选的数字。参数mlm_prob控制掩码比例，论文用的是 0.15，实际复现时有人试过 0.2 也能 work，但偏离太多会改变任务难度。

4.2 Next Sentence Prediction 与它的争议

NSP 任务拿[CLS]的输出去判断两个句子是不是连续的。正样本是原文里相邻的两句，负样本是从别的文档随机抽的一句。源码里seq_relationship是一个二分类线性层。这个任务后来被不少工作质疑，认为它太简单、对下游帮助有限，RoBERTa 直接把它去掉了。但在这份实现里保留 NSP 是有价值的：它让[CLS]这个位置被迫承载句子级别的语义，下游做句子分类时直接拿pooled_output就能用。理解这个设计动机，比单纯记住“NSP 有用/没用”的结论更重要。

5. 避坑与排查：复现 BERT 时最容易翻车的几处

5.1 loss 不降或降得极慢

现象是训练跑了几个 epoch，loss 在 10 附近晃悠。原因通常是学习率没配 warmup，Post-LN 结构在训练初期梯度方差大，直接上大学习率会把参数推到坏区域。解决办法是加线性 warmup，前 10% 的步数把学习率从 0 线性拉到峰值，再线性衰减。另一个常见原因是 mask 的加法方向搞反了，把-10000加到了真实 token 上，模型只能看到 padding，自然学不到东西。

5.2 显存爆掉但 batch size 已经很小

现象是CUDA out of memory，但 batch 只设了 8。原因往往不是 batch 本身，而是注意力分数矩阵的维度是[batch, heads, seq, seq]，seq 长度 512、heads 12 时，单个样本的这一层就要占几十 MB，反向传播还要存一份。解决办法是缩短序列长度、减少 heads，或者用梯度累积模拟大 batch。还有一个隐蔽原因是attention_mask用了 float 而不是 bool，PyTorch 会额外分配内存做类型转换。

5.3 微调时准确率不如随机猜

现象是下游任务微调后指标和瞎猜差不多。原因通常是[CLS]的用法错了：做句子分类应该取pooled_output（已经过线性层和 tanh），而不是sequence_output[:, 0, :]。两者差一个线性变换，直接取后者等于让分类头自己去学这个变换，小数据集上根本学不出来。另一个原因是学习率沿用了预训练的大值，微调阶段学习率应该小一到两个数量级，否则预训练学到的表示会被冲掉。

5.4 位置编码越界

现象是换了一个比预训练更长的输入，直接报 index out of range。原因是position_embeddings的长度固定为max_position，超出就查不到。解决办法是截断输入，或者换用可外推的位置编码方案。这个坑在长文本场景下几乎必踩，提前知道能省不少调试时间。

6. 从这份实现看预训练架构的设计取舍

把这份 PyTorch 实现读完，会发现 BERT 的每个设计选择背后都有一组权衡。双向注意力换来了更强的上下文表示，代价是不能像 GPT 那样做自回归生成；绝对位置编码实现简单，代价是长度外推能力差；Post-LN 忠于原论文，代价是对学习率和 warmup 更敏感；MLM 的 80/10/10 缓解了预训练与微调的分布差异，代价是训练信号只来自 15% 的位置，收敛更慢。这些取舍没有绝对的对错，只有场景匹配与否。

验证一份实现是否真的对，我一般会做三件事：先用极小的模型和造出来的数据过拟合，看 loss 能不能降到接近零，这一步能排除大部分实现 bug；再检查 mask 和 label 的忽略逻辑，确认梯度只来自该来的地方；最后拿预训练权重跑一个下游小任务，看指标是否在合理区间。这三步走完，基本能判断这份代码是“看起来像”还是“真的能跑”。

从那以后我每次读预训练模型的源码，都会先把输入表示和损失计算这两块单独拎出来跑一遍小数据，确认数值对得上再往下看。很多玄学问题，其实在第一步就能暴露。希望这份拆解能帮你在读这份实现时少走几段弯路，把注意力放在机制本身而不是调试细节上。

本文还有配套的精品资源，点击获取