简介这份资源是 Google AI 于 2018 年提出的 BERT 模型的 PyTorch 实现面向已具备 Transformer 与自然语言处理基础、希望从代码层面理解预训练语言模型的中高级学习者。项目以简洁注释还原 BERT 的核心结构涵盖模型定义、数据集处理与训练器三大模块可用于复现论文思路、搭建自己的预训练实验或迁移到下游 NLP 任务。压缩包共 33 个文件以 27 个 Python 脚本为主体另含 Makefile、配置文件、依赖清单与说明文档整体约 28KB体量轻巧、便于通读源码。目前已有 879 人学习下载说明其在 BERT 入门与源码研读场景中具有一定参考价值。读者可借此理清自注意力编码器、语料预处理与训练流程的组织方式并对照论文理解深度双向 Transformer 的实现细节适合作为研读 BERT 源码的起点。1. 从一份“能跑起来的 BERT”说起2018 年 Google AI 放出 BERT 之后NLP 圈几乎被重写了一遍。但真到动手时很多人卡在第一步论文看得懂Transformer 结构也画得出来可官方实现是 TensorFlow 的想在自己熟悉的 PyTorch 里复现一遍从 tokenizer 到预训练任务再到微调处处是坑。这份 BERT-pytorch 项目就是干这件事的——用 PyTorch 从零实现 BERT 的完整链路包括 WordPiece 分词、Masked LM 与 NSP 两个预训练任务、Transformer 编码器堆叠以及下游分类任务的微调接口。它适合两类人一类是想把 Transformer 内部结构吃透、愿意逐行读代码的 NLP 学习者另一类是手里有语料、想自己跑一遍预训练或拿预训练权重做文本分类的工程开发者。下面我按“这项目怎么组织、怎么跑、怎么改”的顺序拆一遍。2. 项目结构与技术选型从目录看 BERT 的模块边界2.1 目录布局与各模块职责这类从零实现的 BERT 项目目录通常不会太深核心逻辑集中在几个文件里。我按常见组织方式梳理一遍你拿到源码后可以对照着看BERT-pytorch/ ├── models/ │ ├── bert.py # BERT 主体Embedding Encoder 堆叠 Pooler │ ├── transformer.py # 单层 Transformer Block多头注意力 FFN 残差 │ ├── embedding.py # Token / Position / Segment 三种 Embedding 相加 │ └── attention.py # 缩放点积注意力与多头拆分 ├── dataset/ │ ├── vocab.py # 词表加载与 WordPiece 分词 │ └── pretrain_dataset.py# 预训练数据生成 [MASK] 样本与 NSP 正负样本对 ├── tasks/ │ ├── pretrain.py # 预训练入口MLM NSP 联合损失 │ └── classify.py # 下游文本分类微调入口 ├── config.py # 超参配置层数、头数、隐藏维度、词表大小 └── utils/ ├── optimizer.py # AdamW 学习率 warmup 调度 └── checkpoint.py # 权重保存与加载这个划分的好处是边界清晰models/只管网络结构dataset/只管数据怎么变成张量tasks/是训练脚本config.py集中所有超参。二次开发时你改结构只动models/换数据只动dataset/不会牵一发动全身。常见做法是把config.py写成一个 dataclass字段名和论文里的符号对应比如hidden_size、num_hidden_layers、num_attention_heads、intermediate_size这样读代码时能直接和论文表格对上。2.2 技术栈与依赖技术栈很纯粹PyTorch 做张量计算和自动求导Python 标准库处理文件读写分词部分要么自己实现 WordPiece要么依赖tokenizers这类库。没有花哨的框架封装这正是它的价值——每一层都能看到。依赖大致是pip install torch numpy tqdm # 如果分词用现成库 pip install tokenizers我一般会额外装tensorboard看 loss 曲线预训练这种长任务光看终端输出很难判断模型是不是在学。注意 PyTorch 版本别太老多头注意力的scaled_dot_product_attention在较新版本里才有老版本得手写矩阵乘法性能差一截。2.3 配置对象超参集中管理BERT-base 的经典配置是 12 层、768 隐藏维度、12 个注意力头、词表 30522。项目里通常用一个配置类承载这些from dataclasses import dataclass dataclass class BertConfig: vocab_size: int 30522 # 词表大小WordPiece 词表决定 hidden_size: int 768 # 每层输出维度也是 embedding 维度 num_hidden_layers: int 12 # Transformer Block 堆叠层数 num_attention_heads: int 12# 多头注意力的头数需整除 hidden_size intermediate_size: int 3072# FFN 中间层维度通常是 hidden 的 4 倍 max_position_embeddings: int 512 # 最大序列长度 hidden_dropout_prob: float 0.1 attention_probs_dropout_prob: float 0.1参数说明hidden_size必须能被num_attention_heads整除否则拆头时维度对不上这是新手最容易翻车的地方intermediate_size取 4 倍是论文默认小数据集上可以调小省显存max_position_embeddings决定了位置编码表的大小超过这个长度的输入会被截断。我一般会把这个配置存成 JSON 一起进 checkpoint加载权重时先读配置再建模型避免结构对不上。3. 核心模块拆解一个句子如何走完 BERT 的前向3.1 Embedding 层三种信息的叠加BERT 的输入表示是 Token Embedding、Position Embedding、Segment Embedding 三者相加。Token 表示词本身Position 表示位置因为 Transformer 本身没有顺序感Segment 区分句子 A 和句子 B为 NSP 任务服务。import torch import torch.nn as nn class BertEmbedding(nn.Module): def __init__(self, config): super().__init__() # 三个独立的 embedding 表 self.token_emb nn.Embedding(config.vocab_size, config.hidden_size, padding_idx0) self.pos_emb nn.Embedding(config.max_position_embeddings, config.hidden_size) self.seg_emb nn.Embedding(2, config.hidden_size) # 只有 A/B 两种 self.layer_norm nn.LayerNorm(config.hidden_size, eps1e-12) self.dropout nn.Dropout(config.hidden_dropout_prob) def forward(self, input_ids, segment_ids): seq_len input_ids.size(1) # 位置 id 从 0 到 seq_len-1每次前向动态生成 position_ids torch.arange(seq_len, deviceinput_ids.device).unsqueeze(0) x self.token_emb(input_ids) self.pos_emb(position_ids) self.seg_emb(segment_ids) return self.dropout(self.layer_norm(x))逻辑说明三个 embedding 维度都是hidden_size直接逐元素相加相加后过 LayerNorm 再 dropout这是论文里的顺序。padding_idx0让 padding 位置的 token embedding 恒为零且不更新梯度省得 padding 干扰注意力。参数上要注意segment_ids的取值范围只能是 0 或 1构造数据时别越界否则nn.Embedding会直接报索引错误。3.2 多头注意力拆头、缩放、拼接多头注意力是 BERT 的计算核心。把hidden_size拆成num_heads份每份独立算注意力最后拼回来。缩放因子是1/sqrt(head_dim)防止点积过大导致 softmax 梯度消失。class MultiHeadAttention(nn.Module): def __init__(self, config): super().__init__() assert config.hidden_size % config.num_attention_heads 0 self.num_heads config.num_attention_heads self.head_dim config.hidden_size // config.num_attention_heads # Q/K/V 三个线性投影一次算完再拆头 self.qkv nn.Linear(config.hidden_size, config.hidden_size * 3) self.out nn.Linear(config.hidden_size, config.hidden_size) self.dropout nn.Dropout(config.attention_probs_dropout_prob) def forward(self, x, attention_mask): B, L, _ x.shape qkv self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim) q, k, v qkv.permute(2, 0, 3, 1, 4) # [3, B, heads, L, head_dim] # 缩放点积注意力 scores torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5) # mask 掉 padding 位置用极小值填充 scores scores.masked_fill(attention_mask 0, -1e9) probs self.dropout(torch.softmax(scores, dim-1)) ctx torch.matmul(probs, v) # [B, heads, L, head_dim] ctx ctx.transpose(1, 2).reshape(B, L, -1) # 拼回 hidden_size return self.out(ctx)逻辑说明qkv一次线性变换出三份比三个独立 Linear 少一次 kernel 调用permute把 heads 维度提到前面方便批量矩阵乘masked_fill用-1e9而不是-inf避免 softmax 出现 NaN这是血泪经验。参数上attention_mask的形状要和scores的[B, heads, L, L]广播兼容通常传[B, 1, 1, L]或[B, L]再 unsqueeze。改头数时只动num_attention_headshead_dim自动跟着变但记得hidden_size得整除。3.3 Transformer Block 与整体堆叠单层 Block 是“注意力 残差 LayerNorm FFN 残差 LayerNorm”的标准结构。BERT 用的是 Post-LN先残差后归一化这点和后来的 GPT 系列 Pre-LN 不同训练时对 warmup 更敏感。class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.attn MultiHeadAttention(config) self.norm1 nn.LayerNorm(config.hidden_size, eps1e-12) self.norm2 nn.LayerNorm(config.hidden_size, eps1e-12) self.ffn nn.Sequential( nn.Linear(config.hidden_size, config.intermediate_size), nn.GELU(), # BERT 用 GELU 而非 ReLU nn.Linear(config.intermediate_size, config.hidden_size), ) self.dropout nn.Dropout(config.hidden_dropout_prob) def forward(self, x, attention_mask): x self.norm1(x self.dropout(self.attn(x, attention_mask))) x self.norm2(x self.dropout(self.ffn(x))) return x逻辑说明残差连接保证深层梯度能回传LayerNorm 稳定分布。GELU 是 BERT 的默认激活换成 ReLU 也能跑但和预训练权重不匹配。整个 BERT 就是把TransformerBlock堆num_hidden_layers层最后接一个 Pooler取[CLS]位置过一层 tanh 线性变换供分类任务用。堆叠时我一般用nn.ModuleList而不是nn.Sequential方便以后做层间分析或梯度检查点。3.4 预训练任务MLM 与 NSP 的损失构造预训练有两个目标Masked LM 随机遮词让模型预测Next Sentence Prediction 判断两句是否连续。数据构造是关键80% 换[MASK]、10% 换随机词、10% 保持原词的经典比例要写对。def create_mlm_labels(input_ids, vocab_size, mask_id, pad_id, mask_prob0.15): labels input_ids.clone() prob torch.full(input_ids.shape, mask_prob) prob[input_ids pad_id] 0.0 # padding 不参与预测 masked torch.bernoulli(prob).bool() labels[~masked] -100 # -100 是 CrossEntropy 的忽略索引 rand torch.rand(input_ids.shape) # 80% 换 [MASK]10% 换随机词10% 不变 input_ids[masked (rand 0.8)] mask_id random_words torch.randint(vocab_size, input_ids.shape) input_ids[masked (rand 0.8) (rand 0.9)] random_words[masked (rand 0.8) (rand 0.9)] return input_ids, labels逻辑说明labels只在被遮位置有真实值其余设-100让损失函数忽略-100是 PyTorchCrossEntropyLoss的默认ignore_index别改成别的值。参数mask_prob0.15是论文默认小语料上可以适当调高让模型多见监督信号。NSP 部分则是构造 50% 连续句对、50% 随机句对标签 0/1接一个二分类头。两个损失相加联合优化权重一般各取 1。4. 跑起来从零训练与加载权重两条路4.1 环境准备与数据格式先确认 PyTorch 能用到 GPUpython -c import torch; print(torch.__version__, torch.cuda.is_available())预训练数据一般是纯文本一行一句或一段。项目里的dataset/pretrain_dataset.py会负责分词、截断到max_position_embeddings、构造句对。我一般先把语料清洗成每行一句去掉空行和超长行再跑分词。词表要么用项目自带的 WordPiece 词表要么用自己的语料训一个——后者更贴合领域但词表大小别超过配置里的vocab_size否则 embedding 索引越界。4.2 预训练启动与关键参数启动脚本通常长这样python -m tasks.pretrain \ --data_path ./corpus.txt \ --vocab_path ./vocab.txt \ --batch_size 32 \ --lr 1e-4 \ --epochs 10 \ --warmup_steps 10000 \ --save_dir ./checkpoints参数说明batch_size受显存限制BERT-base 在 16G 卡上大概能到 32lr预训练一般 1e-4 量级配合 warmupwarmup_steps是学习率从 0 线性升到峰值再衰减的步数Post-LN 结构对 warmup 很敏感省了它 loss 容易炸。训练时盯两个数MLM 的 loss 应该稳步下降NSP 的 accuracy 应该从 0.5 往上走如果 NSP 一直卡在 0.5多半是负样本构造有问题。4.3 下游分类微调有了预训练权重分类任务只需在[CLS]上接一个线性层class BertForClassification(nn.Module): def __init__(self, bert, num_labels): super().__init__() self.bert bert self.classifier nn.Linear(bert.config.hidden_size, num_labels) def forward(self, input_ids, segment_ids, attention_mask): # 取 [CLS] 位置索引 0的池化输出 pooled self.bert(input_ids, segment_ids, attention_mask)[:, 0] return self.classifier(pooled)微调时学习率要比预训练小一个量级常见 2e-5 到 5e-5训练轮数 3 到 5 轮就够多了容易过拟合小数据集。加载权重时注意键名匹配预训练模型和分类模型的 BERT 部分参数名要一致不一致就用load_state_dict(strictFalse)再手动核对缺失的键。5. 避坑与排查预训练路上最容易翻车的几处5.1 loss 变成 NaN 或突然爆炸现象训练几百步后 loss 变nan或者突然从 2 点几跳到几十。原因通常是学习率过高、没做 warmup或者注意力 mask 用了-inf导致 softmax 出 NaN。解决把warmup_steps加上峰值学习率降到 5e-5 试mask 填充值统一用-1e9再加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)。我一般训练前几百步会打印梯度范数异常增大就是前兆。5.2 显存不够batch 上不去现象CUDA out of memory调小 batch 后训练又太慢。原因BERT-base 参数量约 1.1 亿加上优化器状态和激活值显存吃得很凶。解决开混合精度torch.cuda.amp激活值显存能省近一半或者用梯度累积小 batch 多步累加等效大 batch再不行就减层数或隐藏维度做个小号 BERT 先验证流程。5.3 分词结果和预期对不上现象中文语料分词后全是[UNK]或者词表里查不到常见词。原因WordPiece 词表是按特定语料训的直接拿英文词表处理中文必然大面积未知。解决用自己的语料重新训词表或者换字符级分词同时检查vocab_size配置和实际词表行数是否一致不一致会导致 embedding 越界或浪费。5.4 微调后效果不如直接训一个小模型现象拿预训练权重微调分类准确率还不如从头训的 TextCNN。原因预训练语料和下游任务领域差太远或者预训练本身没训够loss 还没降下来就停了。解决确认预训练 loss 已收敛再拿去微调领域差异大时先在目标领域语料上继续预训练一段再微调小数据集上冻结底层、只训顶层几层也能缓解过拟合。6. 二次开发改结构、换任务、验证改动的具体手法拿到这份代码后最常见的改造方向有三个换下游任务、改模型规模、接自己的数据管道。换任务最直接——分类头换成序列标注头每个 token 接一个线性层输出标签或者换成句对匹配两个[CLS]拼接再过分类层。改规模就是动config.py里的层数和维度做个小号 BERT 在单卡上快速验证想法验证通了再放大。接数据管道则是重写dataset/里的__getitem__保证返回的input_ids、segment_ids、attention_mask三个张量形状对齐。改完怎么验证没改坏我的习惯是写一个极小的过拟合测试拿 8 条样本关掉 dropout训 200 步看 loss 能不能降到接近 0。能过拟合说明前向反向链路是通的再去跑全量数据。这个测试帮我省过很多次“训练半天发现 mask 传错维度”的时间。还有一个具体技巧改注意力或 FFN 结构时先保持参数量不变做对照实验。比如把 GELU 换 ReLU其他全不动跑同样的步数比 loss 曲线这样能判断改动本身的影响而不是被参数量变化干扰。我一般会在config.py里加个variant字段不同变体存不同 checkpoint 目录方便回溯。最后说个教训。早期我做预训练时图省事checkpoint 只存模型权重不存配置和优化器状态结果中断后想续训学习率调度器状态丢了续上的曲线和原来完全对不上白跑两天。从那以后我每次存 checkpoint 都强制把 config、optimizer state、scheduler state、当前 step 一起打包续训前先assert一遍字段齐全。希望这些能帮到你少走点我踩过的弯路。本文还有配套的精品资源点击获取
简介这份资源是 Google AI 于 2018 年提出的 BERT 模型的 PyTorch 实现面向已具备 Transformer 与自然语言处理基础、希望从代码层面理解预训练语言模型的中高级学习者。项目以简洁注释还原 BERT 的核心结构涵盖模型定义、数据集处理与训练器三大模块可用于复现论文思路、搭建自己的预训练实验或迁移到下游 NLP 任务。压缩包共 33 个文件以 27 个 Python 脚本为主体另含 Makefile、配置文件、依赖清单与说明文档整体约 28KB体量轻巧、便于通读源码。目前已有 879 人学习下载说明其在 BERT 入门与源码研读场景中具有一定参考价值。读者可借此理清自注意力编码器、语料预处理与训练流程的组织方式并对照论文理解深度双向 Transformer 的实现细节适合作为研读 BERT 源码的起点。1. 从一份“能跑起来的 BERT”说起2018 年 Google AI 放出 BERT 之后NLP 圈几乎被重写了一遍。但真到动手时很多人卡在第一步论文看得懂Transformer 结构也画得出来可官方实现是 TensorFlow 的想在自己熟悉的 PyTorch 里复现一遍从 tokenizer 到预训练任务再到微调处处是坑。这份 BERT-pytorch 项目就是干这件事的——用 PyTorch 从零实现 BERT 的完整链路包括 WordPiece 分词、Masked LM 与 NSP 两个预训练任务、Transformer 编码器堆叠以及下游分类任务的微调接口。它适合两类人一类是想把 Transformer 内部结构吃透、愿意逐行读代码的 NLP 学习者另一类是手里有语料、想自己跑一遍预训练或拿预训练权重做文本分类的工程开发者。下面我按“这项目怎么组织、怎么跑、怎么改”的顺序拆一遍。2. 项目结构与技术选型从目录看 BERT 的模块边界2.1 目录布局与各模块职责这类从零实现的 BERT 项目目录通常不会太深核心逻辑集中在几个文件里。我按常见组织方式梳理一遍你拿到源码后可以对照着看BERT-pytorch/ ├── models/ │ ├── bert.py # BERT 主体Embedding Encoder 堆叠 Pooler │ ├── transformer.py # 单层 Transformer Block多头注意力 FFN 残差 │ ├── embedding.py # Token / Position / Segment 三种 Embedding 相加 │ └── attention.py # 缩放点积注意力与多头拆分 ├── dataset/ │ ├── vocab.py # 词表加载与 WordPiece 分词 │ └── pretrain_dataset.py# 预训练数据生成 [MASK] 样本与 NSP 正负样本对 ├── tasks/ │ ├── pretrain.py # 预训练入口MLM NSP 联合损失 │ └── classify.py # 下游文本分类微调入口 ├── config.py # 超参配置层数、头数、隐藏维度、词表大小 └── utils/ ├── optimizer.py # AdamW 学习率 warmup 调度 └── checkpoint.py # 权重保存与加载这个划分的好处是边界清晰models/只管网络结构dataset/只管数据怎么变成张量tasks/是训练脚本config.py集中所有超参。二次开发时你改结构只动models/换数据只动dataset/不会牵一发动全身。常见做法是把config.py写成一个 dataclass字段名和论文里的符号对应比如hidden_size、num_hidden_layers、num_attention_heads、intermediate_size这样读代码时能直接和论文表格对上。2.2 技术栈与依赖技术栈很纯粹PyTorch 做张量计算和自动求导Python 标准库处理文件读写分词部分要么自己实现 WordPiece要么依赖tokenizers这类库。没有花哨的框架封装这正是它的价值——每一层都能看到。依赖大致是pip install torch numpy tqdm # 如果分词用现成库 pip install tokenizers我一般会额外装tensorboard看 loss 曲线预训练这种长任务光看终端输出很难判断模型是不是在学。注意 PyTorch 版本别太老多头注意力的scaled_dot_product_attention在较新版本里才有老版本得手写矩阵乘法性能差一截。2.3 配置对象超参集中管理BERT-base 的经典配置是 12 层、768 隐藏维度、12 个注意力头、词表 30522。项目里通常用一个配置类承载这些from dataclasses import dataclass dataclass class BertConfig: vocab_size: int 30522 # 词表大小WordPiece 词表决定 hidden_size: int 768 # 每层输出维度也是 embedding 维度 num_hidden_layers: int 12 # Transformer Block 堆叠层数 num_attention_heads: int 12# 多头注意力的头数需整除 hidden_size intermediate_size: int 3072# FFN 中间层维度通常是 hidden 的 4 倍 max_position_embeddings: int 512 # 最大序列长度 hidden_dropout_prob: float 0.1 attention_probs_dropout_prob: float 0.1参数说明hidden_size必须能被num_attention_heads整除否则拆头时维度对不上这是新手最容易翻车的地方intermediate_size取 4 倍是论文默认小数据集上可以调小省显存max_position_embeddings决定了位置编码表的大小超过这个长度的输入会被截断。我一般会把这个配置存成 JSON 一起进 checkpoint加载权重时先读配置再建模型避免结构对不上。3. 核心模块拆解一个句子如何走完 BERT 的前向3.1 Embedding 层三种信息的叠加BERT 的输入表示是 Token Embedding、Position Embedding、Segment Embedding 三者相加。Token 表示词本身Position 表示位置因为 Transformer 本身没有顺序感Segment 区分句子 A 和句子 B为 NSP 任务服务。import torch import torch.nn as nn class BertEmbedding(nn.Module): def __init__(self, config): super().__init__() # 三个独立的 embedding 表 self.token_emb nn.Embedding(config.vocab_size, config.hidden_size, padding_idx0) self.pos_emb nn.Embedding(config.max_position_embeddings, config.hidden_size) self.seg_emb nn.Embedding(2, config.hidden_size) # 只有 A/B 两种 self.layer_norm nn.LayerNorm(config.hidden_size, eps1e-12) self.dropout nn.Dropout(config.hidden_dropout_prob) def forward(self, input_ids, segment_ids): seq_len input_ids.size(1) # 位置 id 从 0 到 seq_len-1每次前向动态生成 position_ids torch.arange(seq_len, deviceinput_ids.device).unsqueeze(0) x self.token_emb(input_ids) self.pos_emb(position_ids) self.seg_emb(segment_ids) return self.dropout(self.layer_norm(x))逻辑说明三个 embedding 维度都是hidden_size直接逐元素相加相加后过 LayerNorm 再 dropout这是论文里的顺序。padding_idx0让 padding 位置的 token embedding 恒为零且不更新梯度省得 padding 干扰注意力。参数上要注意segment_ids的取值范围只能是 0 或 1构造数据时别越界否则nn.Embedding会直接报索引错误。3.2 多头注意力拆头、缩放、拼接多头注意力是 BERT 的计算核心。把hidden_size拆成num_heads份每份独立算注意力最后拼回来。缩放因子是1/sqrt(head_dim)防止点积过大导致 softmax 梯度消失。class MultiHeadAttention(nn.Module): def __init__(self, config): super().__init__() assert config.hidden_size % config.num_attention_heads 0 self.num_heads config.num_attention_heads self.head_dim config.hidden_size // config.num_attention_heads # Q/K/V 三个线性投影一次算完再拆头 self.qkv nn.Linear(config.hidden_size, config.hidden_size * 3) self.out nn.Linear(config.hidden_size, config.hidden_size) self.dropout nn.Dropout(config.attention_probs_dropout_prob) def forward(self, x, attention_mask): B, L, _ x.shape qkv self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim) q, k, v qkv.permute(2, 0, 3, 1, 4) # [3, B, heads, L, head_dim] # 缩放点积注意力 scores torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5) # mask 掉 padding 位置用极小值填充 scores scores.masked_fill(attention_mask 0, -1e9) probs self.dropout(torch.softmax(scores, dim-1)) ctx torch.matmul(probs, v) # [B, heads, L, head_dim] ctx ctx.transpose(1, 2).reshape(B, L, -1) # 拼回 hidden_size return self.out(ctx)逻辑说明qkv一次线性变换出三份比三个独立 Linear 少一次 kernel 调用permute把 heads 维度提到前面方便批量矩阵乘masked_fill用-1e9而不是-inf避免 softmax 出现 NaN这是血泪经验。参数上attention_mask的形状要和scores的[B, heads, L, L]广播兼容通常传[B, 1, 1, L]或[B, L]再 unsqueeze。改头数时只动num_attention_headshead_dim自动跟着变但记得hidden_size得整除。3.3 Transformer Block 与整体堆叠单层 Block 是“注意力 残差 LayerNorm FFN 残差 LayerNorm”的标准结构。BERT 用的是 Post-LN先残差后归一化这点和后来的 GPT 系列 Pre-LN 不同训练时对 warmup 更敏感。class TransformerBlock(nn.Module): def __init__(self, config): super().__init__() self.attn MultiHeadAttention(config) self.norm1 nn.LayerNorm(config.hidden_size, eps1e-12) self.norm2 nn.LayerNorm(config.hidden_size, eps1e-12) self.ffn nn.Sequential( nn.Linear(config.hidden_size, config.intermediate_size), nn.GELU(), # BERT 用 GELU 而非 ReLU nn.Linear(config.intermediate_size, config.hidden_size), ) self.dropout nn.Dropout(config.hidden_dropout_prob) def forward(self, x, attention_mask): x self.norm1(x self.dropout(self.attn(x, attention_mask))) x self.norm2(x self.dropout(self.ffn(x))) return x逻辑说明残差连接保证深层梯度能回传LayerNorm 稳定分布。GELU 是 BERT 的默认激活换成 ReLU 也能跑但和预训练权重不匹配。整个 BERT 就是把TransformerBlock堆num_hidden_layers层最后接一个 Pooler取[CLS]位置过一层 tanh 线性变换供分类任务用。堆叠时我一般用nn.ModuleList而不是nn.Sequential方便以后做层间分析或梯度检查点。3.4 预训练任务MLM 与 NSP 的损失构造预训练有两个目标Masked LM 随机遮词让模型预测Next Sentence Prediction 判断两句是否连续。数据构造是关键80% 换[MASK]、10% 换随机词、10% 保持原词的经典比例要写对。def create_mlm_labels(input_ids, vocab_size, mask_id, pad_id, mask_prob0.15): labels input_ids.clone() prob torch.full(input_ids.shape, mask_prob) prob[input_ids pad_id] 0.0 # padding 不参与预测 masked torch.bernoulli(prob).bool() labels[~masked] -100 # -100 是 CrossEntropy 的忽略索引 rand torch.rand(input_ids.shape) # 80% 换 [MASK]10% 换随机词10% 不变 input_ids[masked (rand 0.8)] mask_id random_words torch.randint(vocab_size, input_ids.shape) input_ids[masked (rand 0.8) (rand 0.9)] random_words[masked (rand 0.8) (rand 0.9)] return input_ids, labels逻辑说明labels只在被遮位置有真实值其余设-100让损失函数忽略-100是 PyTorchCrossEntropyLoss的默认ignore_index别改成别的值。参数mask_prob0.15是论文默认小语料上可以适当调高让模型多见监督信号。NSP 部分则是构造 50% 连续句对、50% 随机句对标签 0/1接一个二分类头。两个损失相加联合优化权重一般各取 1。4. 跑起来从零训练与加载权重两条路4.1 环境准备与数据格式先确认 PyTorch 能用到 GPUpython -c import torch; print(torch.__version__, torch.cuda.is_available())预训练数据一般是纯文本一行一句或一段。项目里的dataset/pretrain_dataset.py会负责分词、截断到max_position_embeddings、构造句对。我一般先把语料清洗成每行一句去掉空行和超长行再跑分词。词表要么用项目自带的 WordPiece 词表要么用自己的语料训一个——后者更贴合领域但词表大小别超过配置里的vocab_size否则 embedding 索引越界。4.2 预训练启动与关键参数启动脚本通常长这样python -m tasks.pretrain \ --data_path ./corpus.txt \ --vocab_path ./vocab.txt \ --batch_size 32 \ --lr 1e-4 \ --epochs 10 \ --warmup_steps 10000 \ --save_dir ./checkpoints参数说明batch_size受显存限制BERT-base 在 16G 卡上大概能到 32lr预训练一般 1e-4 量级配合 warmupwarmup_steps是学习率从 0 线性升到峰值再衰减的步数Post-LN 结构对 warmup 很敏感省了它 loss 容易炸。训练时盯两个数MLM 的 loss 应该稳步下降NSP 的 accuracy 应该从 0.5 往上走如果 NSP 一直卡在 0.5多半是负样本构造有问题。4.3 下游分类微调有了预训练权重分类任务只需在[CLS]上接一个线性层class BertForClassification(nn.Module): def __init__(self, bert, num_labels): super().__init__() self.bert bert self.classifier nn.Linear(bert.config.hidden_size, num_labels) def forward(self, input_ids, segment_ids, attention_mask): # 取 [CLS] 位置索引 0的池化输出 pooled self.bert(input_ids, segment_ids, attention_mask)[:, 0] return self.classifier(pooled)微调时学习率要比预训练小一个量级常见 2e-5 到 5e-5训练轮数 3 到 5 轮就够多了容易过拟合小数据集。加载权重时注意键名匹配预训练模型和分类模型的 BERT 部分参数名要一致不一致就用load_state_dict(strictFalse)再手动核对缺失的键。5. 避坑与排查预训练路上最容易翻车的几处5.1 loss 变成 NaN 或突然爆炸现象训练几百步后 loss 变nan或者突然从 2 点几跳到几十。原因通常是学习率过高、没做 warmup或者注意力 mask 用了-inf导致 softmax 出 NaN。解决把warmup_steps加上峰值学习率降到 5e-5 试mask 填充值统一用-1e9再加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)。我一般训练前几百步会打印梯度范数异常增大就是前兆。5.2 显存不够batch 上不去现象CUDA out of memory调小 batch 后训练又太慢。原因BERT-base 参数量约 1.1 亿加上优化器状态和激活值显存吃得很凶。解决开混合精度torch.cuda.amp激活值显存能省近一半或者用梯度累积小 batch 多步累加等效大 batch再不行就减层数或隐藏维度做个小号 BERT 先验证流程。5.3 分词结果和预期对不上现象中文语料分词后全是[UNK]或者词表里查不到常见词。原因WordPiece 词表是按特定语料训的直接拿英文词表处理中文必然大面积未知。解决用自己的语料重新训词表或者换字符级分词同时检查vocab_size配置和实际词表行数是否一致不一致会导致 embedding 越界或浪费。5.4 微调后效果不如直接训一个小模型现象拿预训练权重微调分类准确率还不如从头训的 TextCNN。原因预训练语料和下游任务领域差太远或者预训练本身没训够loss 还没降下来就停了。解决确认预训练 loss 已收敛再拿去微调领域差异大时先在目标领域语料上继续预训练一段再微调小数据集上冻结底层、只训顶层几层也能缓解过拟合。6. 二次开发改结构、换任务、验证改动的具体手法拿到这份代码后最常见的改造方向有三个换下游任务、改模型规模、接自己的数据管道。换任务最直接——分类头换成序列标注头每个 token 接一个线性层输出标签或者换成句对匹配两个[CLS]拼接再过分类层。改规模就是动config.py里的层数和维度做个小号 BERT 在单卡上快速验证想法验证通了再放大。接数据管道则是重写dataset/里的__getitem__保证返回的input_ids、segment_ids、attention_mask三个张量形状对齐。改完怎么验证没改坏我的习惯是写一个极小的过拟合测试拿 8 条样本关掉 dropout训 200 步看 loss 能不能降到接近 0。能过拟合说明前向反向链路是通的再去跑全量数据。这个测试帮我省过很多次“训练半天发现 mask 传错维度”的时间。还有一个具体技巧改注意力或 FFN 结构时先保持参数量不变做对照实验。比如把 GELU 换 ReLU其他全不动跑同样的步数比 loss 曲线这样能判断改动本身的影响而不是被参数量变化干扰。我一般会在config.py里加个variant字段不同变体存不同 checkpoint 目录方便回溯。最后说个教训。早期我做预训练时图省事checkpoint 只存模型权重不存配置和优化器状态结果中断后想续训学习率调度器状态丢了续上的曲线和原来完全对不上白跑两天。从那以后我每次存 checkpoint 都强制把 config、optimizer state、scheduler state、当前 step 一起打包续训前先assert一遍字段齐全。希望这些能帮到你少走点我踩过的弯路。本文还有配套的精品资源点击获取
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
