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

1次阅读

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

在自然语言处理任务中，BERT 等大型预训练模型的微调过程经常面临显存不足的问题。尤其在处理长文本序列时，输入长度超过 512 个 token 的情况会导致显存需求急剧上升。原生 BERT 实现中，Adam 优化器需要保存模型参数的动量和方差估计，这会占用相当于模型参数两倍的显存空间。例如，在使用 BERT-base（110M 参数）时，仅优化器状态就需要额外消耗约 1.7GB 显存。

梯度累积的核心思想是通过多次前向传播累积梯度，再进行一次参数更新。数学表达为：

$$\theta_{t+1} = \theta_t – \eta\cdot\frac{1}{N}\sum_{i=1}^N g_i$$

其中 $N$ 为累积步数，$g_i$ 为第 $i$ 步的梯度。这种方法使得等效 batch size 扩大 $N$ 倍，而实际显存占用仅增加前向传播的中间结果。

自动混合精度训练通过以下方式节省显存：

将部分计算转为 FP16 格式
维护 FP32 主权重副本
动态损失缩放防止梯度下溢

关键配置参数包括：

opt_level: O1（默认）或 O2
loss_scale: “dynamic”（推荐）

基于论文《BERT: Pre-training of Deep Transformers for Language Understanding》的建议，不同网络层采用差异化的学习率：

$$lr_i = lr_{base} \cdot \alpha^{L-i}$$

其中 $L$ 为总层数，$\alpha$ 为衰减系数（通常取 0.95）。

import torch
from apex import amp

# 初始化模型和优化器
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

# 混合精度初始化
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

# 分层学习率设置
param_groups = [{'params': model.bert.embeddings.parameters(), 'lr': 1e-5},
    {'params': model.bert.encoder.layer[:6].parameters(), 'lr': 2e-5},
    {'params': model.bert.encoder.layer[6:].parameters(), 'lr': 3e-5},
    {'params': model.classifier.parameters(), 'lr': 4e-5}
]
optimizer = torch.optim.AdamW(param_groups)

# 训练循环
accum_steps = 4
scaler = amp.GradScaler()

for epoch in range(3):
    for step, batch in enumerate(train_loader):
        inputs = {k:v.cuda() for k,v in batch.items()}

        with amp.autocast():
            outputs = model(**inputs)
            loss = outputs.loss / accum_steps

        scaler.scale(loss).backward()

        if (step+1) % accum_steps == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

在 IMDb 影评数据集上的测试结果：