BERT 预训练实战PyTorch 复现 MLM 与 NSP 双任务Loss 降至 1.2BERTBidirectional Encoder Representations from Transformers作为自然语言处理领域的里程碑模型其核心价值在于通过Masked Language ModelMLM和Next Sentence PredictionNSP两个预训练任务学习深层次的语言表示。本文将带您从零开始使用PyTorch框架完整实现BERT的预训练过程并分享将训练Loss降至1.2的实战调优经验。1. 环境准备与数据预处理在开始BERT预训练之前我们需要搭建合适的开发环境并准备训练数据。以下是关键步骤# 环境依赖安装 pip install torch1.12.1 transformers4.28.1 datasets2.11.0对于训练数据建议使用Wikipedia或BookCorpus等开源语料。数据预处理包括以下几个关键环节文本清洗去除特殊字符、HTML标签等非文本内容分词处理使用WordPiece分词器构建词汇表样本生成创建MLM和NSP任务所需的训练样本from transformers import BertTokenizer tokenizer BertTokenizer.from_pretrained(bert-base-uncased) text The quick brown fox jumps over the lazy dog. tokenized_text tokenizer.tokenize(text) # 输出: [the, quick, brown, fox, jumps, over, the, lazy, dog, .]2. BERT模型架构实现BERT的核心是由多层Transformer编码器堆叠而成。下面我们实现关键的模型组件2.1 嵌入层实现BERT的输入嵌入由三部分组成Token Embeddings词向量表示Segment Embeddings区分句子A和BPosition Embeddings编码位置信息import torch import torch.nn as nn class BertEmbeddings(nn.Module): def __init__(self, config): super().__init__() self.token_embeddings nn.Embedding(config.vocab_size, config.hidden_size) self.position_embeddings nn.Embedding(config.max_position_embeddings, config.hidden_size) self.segment_embeddings nn.Embedding(config.type_vocab_size, config.hidden_size) self.LayerNorm nn.LayerNorm(config.hidden_size) self.dropout nn.Dropout(config.hidden_dropout_prob) def forward(self, input_ids, token_type_idsNone): seq_length input_ids.size(1) position_ids torch.arange(seq_length, dtypetorch.long, deviceinput_ids.device) position_ids position_ids.unsqueeze(0).expand_as(input_ids) if token_type_ids is None: token_type_ids torch.zeros_like(input_ids) token_embeddings self.token_embeddings(input_ids) position_embeddings self.position_embeddings(position_ids) segment_embeddings self.segment_embeddings(token_type_ids) embeddings token_embeddings position_embeddings segment_embeddings embeddings self.LayerNorm(embeddings) embeddings self.dropout(embeddings) return embeddings2.2 Transformer编码器层每个Transformer编码器层包含多头自注意力机制和前馈神经网络class BertLayer(nn.Module): def __init__(self, config): super().__init__() self.attention BertAttention(config) self.intermediate BertIntermediate(config) self.output BertOutput(config) def forward(self, hidden_states, attention_maskNone): attention_output self.attention(hidden_states, attention_mask) intermediate_output self.intermediate(attention_output) layer_output self.output(intermediate_output, attention_output) return layer_output3. 预训练任务实现3.1 Masked Language Model (MLM)MLM任务随机遮盖输入序列中的部分token让模型预测被遮盖的原始tokenclass BertMLMHead(nn.Module): def __init__(self, config): super().__init__() self.dense nn.Linear(config.hidden_size, config.hidden_size) self.transform_act_fn nn.GELU() self.LayerNorm nn.LayerNorm(config.hidden_size) self.decoder nn.Linear(config.hidden_size, config.vocab_size, biasFalse) self.bias nn.Parameter(torch.zeros(config.vocab_size)) def forward(self, hidden_states): hidden_states self.dense(hidden_states) hidden_states self.transform_act_fn(hidden_states) hidden_states self.LayerNorm(hidden_states) hidden_states self.decoder(hidden_states) self.bias return hidden_states提示在实际实现中建议采用80-10-10的遮盖策略80%替换为[MASK]10%随机替换10%保持不变。3.2 Next Sentence Prediction (NSP)NSP任务判断两个句子是否连续class BertNSPHead(nn.Module): def __init__(self, config): super().__init__() self.seq_relationship nn.Linear(config.hidden_size, 2) def forward(self, pooled_output): return self.seq_relationship(pooled_output)4. 训练优化策略要实现Loss降至1.2的目标需要精心设计训练策略4.1 学习率调度采用带warmup的线性衰减学习率from torch.optim import AdamW optimizer AdamW(model.parameters(), lr5e-5, eps1e-8) scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_steps10000, num_training_stepstotal_steps )4.2 梯度累积与混合精度训练scaler torch.cuda.amp.GradScaler() for step, batch in enumerate(train_dataloader): with torch.cuda.amp.autocast(): outputs model(**batch) loss outputs.loss loss loss / gradient_accumulation_steps scaler.scale(loss).backward() if (step 1) % gradient_accumulation_steps 0: scaler.step(optimizer) scaler.update() scheduler.step() optimizer.zero_grad()4.3 关键超参数配置参数推荐值说明batch_size32-256根据GPU显存调整max_seq_length128-512影响计算复杂度learning_rate5e-5基础学习率warmup_steps10,000warmup步数num_train_epochs3-5训练轮次adam_epsilon1e-8Adam优化器参数5. 模型评估与调试训练过程中需要监控多个指标Loss曲线观察训练和验证Loss的变化趋势准确率MLM和NSP任务的预测准确率梯度范数防止梯度爆炸或消失def evaluate(model, eval_dataloader): model.eval() total_loss 0 total_samples 0 with torch.no_grad(): for batch in eval_dataloader: outputs model(**batch) loss outputs.loss total_loss loss.item() * batch[input_ids].size(0) total_samples batch[input_ids].size(0) return total_loss / total_samples6. 实战调优经验分享在将Loss降至1.2的过程中我们总结了以下关键经验动态遮盖策略每次epoch重新随机选择遮盖位置避免模型记忆特定模式学习率预热前10,000步逐步提高学习率稳定训练初期梯度裁剪设置梯度阈值为1.0防止梯度爆炸层归一化调整在注意力机制后添加额外的层归一化注意力头剪枝训练后期可尝试剪枝部分注意力头提升效率# 动态遮盖实现示例 def create_masked_lm_labels(input_ids, tokenizer, mlm_probability0.15): labels input_ids.clone() probability_matrix torch.full(labels.shape, mlm_probability) masked_indices torch.bernoulli(probability_matrix).bool() # 80%替换为[MASK], 10%随机替换, 10%保持不变 indices_replaced torch.bernoulli(torch.full(labels.shape, 0.8)).bool() masked_indices input_ids[indices_replaced] tokenizer.convert_tokens_to_ids(tokenizer.mask_token) indices_random torch.bernoulli(torch.full(labels.shape, 0.5)).bool() masked_indices ~indices_replaced random_words torch.randint(len(tokenizer), labels.shape, dtypetorch.long) input_ids[indices_random] random_words[indices_random] return labels通过以上方法和技巧的组合应用我们成功将BERT预训练的Loss稳定在1.2左右模型在下游任务上的微调表现也达到了预期效果。