BERT预训练实战：从零构建高效微调流水线

1次阅读

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

在 NLP 业务落地时，BERT 预训练面临三大挑战：

硬件资源消耗：基础 BERT-large 模型需要 16GB 显存，业务数据微调时显存占用可能翻倍
长文本处理效率：当序列长度超过 512 时，Self-Attention/ 自注意力的计算复杂度呈平方级增长
领域适应成本：医疗 / 金融等垂直领域需要 Domain-Adaptive Pretraining/ 领域自适应预训练，但 Full-Pretraining/ 完整预训练成本过高
Full-Pretraining 优势：模型通用性强，适合多任务场景
Domain-Adaptive Pretraining 优势：领域任务表现提升 5 -15%，但需要领域语料和重新预训练

使用 transformers.Trainer 作为基础框架
集成 datasets 库实现高效数据加载
通过 accelerate 库支持多机多卡训练

Gradient Accumulation/ 梯度累积：
原理：分多步计算梯度后再更新参数
效果：batch_size= 8 时，累积 4 步等效于 batch_size=32
AMP/ 自动混合精度：
FP16 存储：减少 50% 显存占用
FP32 计算：保持数值稳定性
需配合 torch.cuda.amp.GradScaler 使用
Gradient Checkpointing/ 梯度检查点：
时间换空间：重新计算中间激活值
实现：model.gradient_checkpointing_enable()

from pytorch_lightning import LightningModule

class BERTPretrainer(LightningModule):
    def __init__(self, model_name="bert-base-uncased"):
        super().__init__()
        self.model = AutoModelForMaskedLM.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

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

from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    input_ids = [torch.tensor(x["input_ids"]) for x in batch]
    attention_mask = [torch.ones(len(x)) for x in input_ids]

    return {"input_ids": pad_sequence(input_ids, batch_first=True),
        "attention_mask": pad_sequence(attention_mask, batch_first=True)
    }

环境变量设置：

export CUDA_VISIBLE_DEVICES=0,1,2,3
export WORLD_SIZE=4

PyTorch Lightning 参数：

trainer = Trainer(
    devices=4,
    strategy="ddp",
    precision=16
)

梯度裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

学习率 warmup：

scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=1000,
    num_training_steps=total_steps
)

Shuffle 与 Dataloader Worker：
Worker 数 >0 时要设置worker_init_fn

示例：

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    numpy.random.seed(worker_seed)

NaN 值检测：

if torch.isnan(loss).any():
    optimizer.zero_grad()
    continue

解决方案：
降低学习率（尝试 1e- 5 到 5e-5）
减小 batch size
启用 AMP 的keep_batchnorm_fp32=True

测试环境：AWS p3.8xlarge（V100 32GB * 8）

优化方案	显存占用	训练速度
基线方案	22.1GB	1.0x
+AMP	14.3GB	1.7x
+ 梯度累积	9.8GB	1.2x
全优化方案	7.5GB	2.3x

通过本文方案，我们成功将单卡可训练的序列长度从 512 提升到 1024。实际业务中建议：
1. 小规模数据：直接使用 Domain-Adaptive Pretraining
2. 通用场景：采用 HuggingFace 提供的预训练模型 + 本文微调方案
3. 超长文本：考虑 Reformer/Longformer 等变体模型

正文完