BERT预训练数据集构建实战：从数据清洗到分布式训练的完整解决方案

1次阅读

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

当前 BERT 预训练数据准备面临三大挑战：多源异构数据（如 Common Crawl/Wikipedia/ 专业文献）的格式差异显著，低质量文本（广告 / 乱码 / 重复内容）占比常超 30%，以及多语言混合场景下的语种识别准确率问题。据 HuggingFace 社区统计，85% 的预训练失败案例源于数据清洗阶段未处理的隐式噪声。

PySpark 实现核心清洗逻辑（同时展示 RDD/DataFrame API）：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from bs4 import BeautifulSoup
import langdetect

# RDD 版本（适合非结构化数据处理）def clean_html_rdd(text: str) -> str:
    try:
        soup = BeautifulSoup(text, 'html.parser')
        return soup.get_text(separator='').encode('utf-8').decode('utf-8','ignore').strip()
    except Exception as e:
        print(f"HTML parsing failed: {e}")
        return ""

# DataFrame 版本（带类型注解和异常处理）@udf("string")
def clean_html_df(text: str) -> str:
    try:
        if not text:
            return ""return BeautifulSoup(text,'html.parser').get_text()
    except Exception as e:
        print(f"UDF processing error: {e}")
        return ""spark = SparkSession.builder.appName("BERT-Preprocess").getOrCreate()

# 原始数据加载（演示 Parquet 格式处理）df = spark.read.parquet("s3://raw_data/*.parquet")

# 执行清洗（包含 UTF- 8 强制转换）cleaned_df = df.select(clean_html_df(df["raw_text"]).alias("clean_text"),
    df["source"]
).filter("length(clean_text) > 50")  # 过滤短文本

语言检测优化技巧：
– 使用 fasttext 替代 langdetect 提升检测速度 300%
– 对中文 / 日文等相似语种添加自定义规则

Tokenizer 类型	内存占用（10GB 文本）	处理速度	特殊 Token 支持
SentencePiece	14.2GB	2.1MB/s	需手动添加
WordPiece	8.7GB	3.4MB/s	自动处理

内存优化关键发现：
– 启用 optimized=True 参数可减少 WordPiece 内存占用 40%
– 对于中文数据，BPE 比 Unigram 算法节省 20% 词表空间

# Horovod+PyTorch 关键配置
FROM nvcr.io/nvidia/pytorch:21.03-py3

RUN pip install horovod[pytorch]==0.23.0 \
    && apt-get update \
    && apt-get install -y openssh-server \
    && mkdir -p /var/run/sshd

# 设置 NCCL 参数
ENV NCCL_DEBUG=INFO \
    NCCL_SOCKET_IFNAME=^lo,docker0 \
    HOROVOD_GPU_OPERATIONS=NCCL

# 启动脚本示例
CMD ["sh", "-c", "horovodrun -np 8 -H worker0:4,worker1:4 python train.py"]

import hashlib
from pyspark.sql import functions as F

def check_contamination(train_df, test_df, col_name="text"):
    """计算训练集与测试集的 MD5 重叠率"""
    train_hashes = train_df.select(F.md5(F.col(col_name)).alias("hash")
    ).distinct().cache()

    test_hashes = test_df.select(F.md5(F.col(col_name)).alias("hash")
    ).distinct().cache()

    overlap = train_hashes.join(test_hashes, "hash").count()
    contamination_rate = overlap / test_hashes.count()

    print(f"Contamination rate: {contamination_rate:.4%}")
    return contamination_rate < 0.001  # 阈值设为 0.1%

Attention Mask 压缩：
使用 scipy.sparse 存储连续 0 值区块
对 padding 超过 50% 的 batch 启用动态 mask

梯度检查点技术：

from torch.utils.checkpoint import checkpoint

class CheckpointedBERT(torch.nn.Module):
    def forward(self, inputs):
        return checkpoint(super().forward, inputs)