BERT预训练模型流程图解：从原理到工程实现的关键路径

1次阅读

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

（图示说明：Tokenization → WordPiece 分字 → Positional Encoding → 12 层 Transformer 编码器 → MLM/NSP 损失计算）

三合一嵌入 ：Token Embeddings + Segment Embeddings + Position Embeddings
关键公式：$E = E_{tok} + E_{seg} + E_{pos}$

# shape: (batch_size, seq_len, hidden_size)
embeddings = word_embeddings(input_ids) \
            + position_embeddings(position_ids) \
            + token_type_embeddings(segment_ids)

QKV 投影 ：$Q = XW^Q$, $K = XW^K$, $V = XW^V$
缩放点积 ：$Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$

# shape: (batch, heads, seq_len, head_dim)
q = self.q_proj(x).view(bsz, seq_len, self.num_heads, -1).transpose(1, 2)
k = self.k_proj(x).view(bsz, seq_len, self.num_heads, -1).transpose(1, 2)
v = self.v_proj(x).view(bsz, seq_len, self.num_heads, -1).transpose(1, 2)

改造 DataLoader 的 collate_fn：

def pad_collate(batch):
    max_len = max(len(x['input_ids']) for x in batch)
    return {'input_ids': pad_sequence([x['input_ids'] for x in batch], batch_first=True),
        'attention_mask': torch.stack([torch.cat([torch.ones(len(x['input_ids'])),
            torch.zeros(max_len - len(x['input_ids']))
        ]) for x in batch])
    }

NaN 值诊断步骤 ：
检查梯度裁剪阈值（建议 2.0）
验证 LayerNorm 的 epsilon 值（BERT 默认 1e-12）
监控各层激活值范围

with torch.autograd.detect_anomaly():
    outputs = model(inputs)
    loss = outputs.loss
    loss.backward()

当处理多个文档拼接输入时，需要：
– 在文档边界处添加特殊分隔符
– 确保跨文档的 attention 权重归零

优势：
显存占用减少 70%
更适合移动端部署
劣势：
微调表现下降约 2 - 3 个点
需要更长的训练周期

（测试环境：NVIDIA V100 32GB，batch_size=32）
| 优化方法 | 显存节省 | 速度提升 |
|——————|———-|———-|
| 梯度累积（4 步）| 42% | 1.2x |
| 动态序列分块 | 37% | 1.5x |
| FP16 混合精度 | 55% | 2.1x |

正文完