BERT预训练实战指南：从零构建高效语言模型

1次阅读

共计 2178 个字符，预计需要花费 6 分钟才能阅读完成。

BERT 作为 NLP 领域的里程碑模型，其预训练过程面临三大核心挑战：

数据规模要求高：原始 BERT-base 训练需要 16GB 文本数据（如 Wikipedia+BookCorpus），数据清洗和格式转换耗时巨大
计算资源消耗大：单卡训练需数周时间，即使使用 TPU/ 多 GPU 集群也面临显存不足问题
收敛稳定性差：学习率策略不当易导致训练崩溃，动态掩码和位置编码的耦合可能引发梯度异常

基于 TensorFlow 1.x 静态计算图
固定序列长度（如 512）导致短文本计算浪费
全局批处理（global batch）依赖 AllReduce 同步

动态 padding：按 batch 内最长序列实时填充，显存利用率提升 30%
梯度累积：模拟大 batch 训练（如gradient_accumulation_steps=4）
智能缓存：自动复用已处理的 tokenized 结果

from transformers import BertTokenizer
import tensorflow as tf  # 用于 TFRecord 生成

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def create_tfrecord(text_pair):
    inputs = tokenizer(text_pair[0], text_pair[1], 
        max_length=512,
        truncation=True,
        padding='max_length',
        return_tensors='tf'
    )
    features = {
        'input_ids': tf.train.Feature(int64_list=tf.train.Int64List(value=inputs['input_ids'].numpy()[0])
        ),
        'attention_mask': tf.train.Feature(int64_list=tf.train.Int64List(value=inputs['attention_mask'].numpy()[0])
        )
    }
    return tf.train.Example(features=tf.train.Features(feature=features))

import pytorch_lightning as pl
from transformers import BertForPreTraining

class BertPretrainer(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = BertForPreTraining.from_pretrained('bert-base-uncased')

    def forward(self, input_ids, attention_mask):
        return self.model(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

    def training_step(self, batch, batch_idx):
        outputs = self(**batch)
        loss = outputs.loss
        self.log('train_loss', loss)
        return loss

# deepspeed_config.json
{
  "train_batch_size": 4096,
  "gradient_accumulation_steps": 8,
  "optimizer": {
    "type": "AdamW",
    "params": {"lr": 6e-5}
  },
  "fp16": {"enabled": true},
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu"}
  }
}