Megatron-LM 详细学习笔记第三章：数据准备与分词管线

快撑死的鱼

862人浏览 · 2025-11-04 19:48:35

快撑死的鱼 · 2025-11-04 19:48:35 发布

第三章：数据准备与分词管线

本章快速使用清单

本章目标：
- 掌握大规模语料的典型清洗、去重与质量过滤流程，理解数据质量对模型能力的决定性影响。
- 学会训练和评估自定义的分词器（Tokenizer），并权衡词表大小带来的利弊。
- 精通 Megatron 特有的 .bin/.idx 二进制数据集格式，并能够独立完成从原始文本到该格式的打包过程。
必备命令/脚本：
- spm_train --input=<corpus.txt> --model_prefix=<mymodel> --vocab_size=50000 ...：使用 SentencePiece 训练一个 BPE 分词器。
- tools/preprocess_data.py --input <in.jsonl> --output-prefix <prefix> --vocab <v.json> ...：Megatron-LM 官方脚本，将 JSONL 格式的文本文件打包成 .bin/.idx 格式。
- pip install datasketch：安装用于大规模数据去重的 MinHash LSH 库。
验收标准：
- 成功将一份小型原始文本语料（例如，几百 MB 的 JSONL 文件）处理成 Megatron 所需的 .bin 和 .idx 文件对。
- 能够清晰阐述所训练分词器的词表大小、特殊 token 等关键配置，并解释其选择理由。
- 产出一份数据统计报告，其中包含总 tokens 数量、各数据源的占比和 token 数量等核心信息。
常见坑点：
1. 文本编码错误：处理多来源语料时，未统一为 UTF-8 编码，导致分词或处理脚本在中途崩溃。
2. 去重不彻底：仅进行精确行级去重，忽略了段落、文档级别的近似重复，导致模型在训练集上“过拟合”于重复样本，泛化能力差。
3. 特殊 Token 处理不当：EOD (End of Document) token 在打包时被遗漏或错误使用，破坏了 Megatron 的文档边界识别，影响学习效果。
4. 内存管理不善：在处理 TB 级数据时，使用一次性将所有数据读入内存的脚本，导致进程被系统杀死。必须使用流式处理（streaming）或分块处理（chunking）方法。
5. 验证集污染：在清洗和去重过程中，无意中将验证集或测试集中的样本泄露到了训练集中，导致评估结果虚高，无法真实反映模型性能。

3.1 语料清洗与去重

“Garbage in, garbage out.” 这句格言在大型语言模型领域被体现得淋漓尽致。一个高质量、多样化且经过精心清洗的数据集，是训练出强大、可靠模型的基础。这一阶段的目标是剔除噪声、消除冗余、过滤有害内容，并规范化文本格式。

1. 语料清洗 (Corpus Cleaning)

清洗是一个多阶段、经验性的过程，其具体步骤取决于数据来源的“脏乱”程度。

清洗任务	解决的问题	常用方法与工具	难度与挑战
HTML/XML 标签去除	网页抓取数据中包含大量 `<div>`, `<a>`, `<script>` 等无用标签。	使用 `BeautifulSoup`, `lxml` 等库解析并提取纯文本内容。正则表达式有时也可用，但鲁棒性较差。	需要处理不规范的 HTML；避免错误地移除代码片段中的 `<` 或 `>` 符号。
样板文件移除 (Boilerplate Removal)	去除网页中的导航栏、页眉、页脚、广告、版权声明等重复性、低信息量内容。	`jusText`, `trafilatura` 等库专门用于从 HTML 中提取正文。基于 DOM 树的密度和文本/标签比进行启发式判断。	难以定义普适的“样板”规则，可能会误伤正文。
格式与编码统一	文本中包含奇怪的控制字符（如 `\x00`）、多种编码格式（GBK, UTF-8）。	将所有文本强制转换为 UTF-8 编码，并移除 Unicode 中不可打印的字符。	确定原始编码可能需要探测库（如 `chardet`），处理混合编码文件很棘手。
语言判定	训练单语种模型时，需要剔除其他语言的文本。	`fastText` (language identification model), `pycld2`, `langdetect` 等库。	短文本判定不准；代码中夹杂的英文注释可能会被误判。
个人可识别信息 (PII) 过滤	移除或脱敏语料中的姓名、电话号码、身份证号、邮箱地址等敏感信息。	基于正则表达式的规则匹配；使用 NER (命名实体识别) 模型来识别并替换 PII。	正则表达式容易漏报或误报；NER 模型本身也非 100% 准确，且有计算开销。
毒性与不当内容过滤 (Toxicity & NSFW Filtering)	移除仇恨言论、暴力、色情等不适合模型学习的内容。	使用预训练的文本分类模型（如 Google 的 Perspective API 或开源的 Detoxify 模型）对文本打分，并根据阈值过滤。维护关键词/黑名单。	文化差异导致“毒性”定义模糊；对抗性样本可以轻易绕过分类器；过度过滤可能引入偏见。

2. 数据去重 (Deduplication)

去重是保证模型泛化能力、防止其“背诵”训练样本的关键步骤。大规模语料中的重复内容远超想象，从句子、段落到整个文档都可能存在。

精确去重 (Exact Deduplication)
- 方法: 计算每个文档或段落的哈希值（如 SHA256），并存储在一个 set 或哈希表中。如果遇到已存在的哈希值，则丢弃该样本。
- 优点: 简单、快速、精确。
- 缺点: 只能处理完全一致的重复。对于仅有微小差别（如一个标点、一个空格）的文本无能为力。
近似去重 (Fuzzy Deduplication)
- 核心思想: 这是大规模去重的难点和重点。目标是识别那些内容高度相似但不完全相同的文档。
- 主流技术：MinHash LSH (Locality-Sensitive Hashing)
  1. 分片 (Shingling): 将文档打碎成一系列重叠的 n-grams (通常 n=5 到 10)。例如，"the cat sat" 的 3-grams 是 ("the", "cat", "sat")。
  2. 哈希 (Hashing): 将每个 shingle 哈希成一个整数。
  3. 最小哈希 (MinHash): 使用 K 个不同的哈希函数，对文档的所有 shingle 哈希值进行计算，并分别取 K 个最小值，形成一个 K 维的“签名向量”（signature）。核心原理：两个集合的 Jaccard 相似度约等于它们 MinHash 签名的相似度。
  4. 局部敏感哈希 (LSH): 将 K 维签名向量分成 b 个 band，每个 band 有 r 行 (K = b * r)。对每个 band 进行哈希，将其映射到一个桶（bucket）中。只要有两个文档在至少一个 band 上的哈希值相同，它们就被视为“候选对”（candidate pair）。
  5. 验证: 对所有候选对，精确计算其 Jaccard 相似度或其他人为定义的相似度，超过阈值（如 0.8）的则认为重复，并只保留其一。
- 优点: 能够在数十亿文档级别上高效地找到近似重复，计算和内存开销可控。
- 工具: datasketch 库提供了高效的 MinHash 和 LSH Forest 实现。

3.2 分词器选择与训练

计算机不理解文本，只理解数字。分词器（Tokenizer）是连接这两者的桥梁，它负责将原始文本字符串转换成一个整数序列（Token IDs），并能再转换回来。

1. 分词算法：BPE 与 SentencePiece

字节对编码 (Byte Pair Encoding, BPE): GPT 系列模型广泛采用的子词（subword）分词算法。
- 训练过程:
  1. 初始化: 词表（vocabulary）由所有单个字符组成。
  2. 迭代合并: 不断统计语料中出现频率最高的相邻“对”（pair），并将它们合并成一个新的、更长的词表单元。
  3. 终止: 当词表达到预设大小（vocab_size）或不再有可合并的对时停止。
- 优点: 能有效处理未登录词（Out-of-Vocabulary, OOV），任何词都可以被拆解成已知的子词单元。它在常见词（如 “transformer”）和罕见词（如 “Megatron”）之间取得了很好的平衡。
SentencePiece: Google 开源的、集成了多种分词算法（包括 BPE 和 Unigram）的工具库。
- 核心优势:
  1. 直接操作 Unicode: 它将文本视为一个 Unicode 字符序列，从而摆脱了对特定语言预处理（如按空格分词）的依赖。
  2. 可逆性: 保证 tokenize(detokenize(text)) 与原始 text（在规范化后）一致。
  3. 空格处理: 将空格视为一个普通字符 _，使其能够被包含在子词单元中，从而无损地重建原始文本。

2. 词表大小 (Vocabulary Size) 的权衡

选择词表大小是训练分词器时最重要的超参数之一，它直接影响模型性能和效率。

特性	小词表 (e.g., 32k)	大词表 (e.g., 128k)
序列长度	更长。一个词可能被拆分成多个 tokens (e.g., `Megatron` -> `Mega`, `tron`)。	更短。常见词和许多专有名词本身就是一个 token。
模型计算量	更高。序列长度是 Transformer 计算复杂度的主要因素 (`O(N^2)`)。	更低。序列更短，计算更快。
Embedding/Softmax 层大小	更小。`vocab_size * hidden_size`，参数量少，内存占用小。	更大。这部分参数量会显著增加，成为模型规模的一大组成部分。
处理 OOV 能力	强。几乎所有词都能由子词组合而成。	相对弱。虽然仍是子词模型，但罕见词被拆分的可能性降低。
跨语言能力	对于多语言模型，小词表可能导致语言间共享更多基础子词，但会使序列变得极长。	大词表可以为每种语言分配更多的专属 tokens，但会稀释共享。

实践建议：现代大模型（如 Llama 2）倾向于使用中等偏大的词表（如 32k-65k），并确保词表中包含用于控制格式（如代码缩进）和特殊任务的控制 tokens。

实践示例：使用 SentencePiece 训练一个分词器

# a. 准备纯文本语料 (假设为 corpus.txt)
# b. 运行训练命令
spm_train \
    --input=corpus.txt \
    --model_prefix=my_tokenizer \
    --vocab_size=50257 \
    --character_coverage=0.9995 \
    --model_type=bpe \
    --pad_id=0 \
    --unk_id=1 \
    --bos_id=2 \
    --eos_id=3 \
    --user_defined_symbols='<sep>,<cls>'

--model_prefix: 输出模型文件的前缀（会生成 .model 和 .vocab 文件）。
--vocab_size: 目标词表大小。
--character_coverage: 确保至少覆盖 99.95% 的字符，防止罕见字符被丢弃。
--model_type=bpe: 指定使用 BPE 算法。
--pad_id, --unk_id, --bos_id, --eos_id: 定义特殊 tokens 的 ID。

3.3 Megatron 数据集打包：`.bin/.idx` 索引

Megatron-LM 为了实现最高效的数据加载，设计了一种独特的二进制数据格式。它由一个数据文件 (.bin) 和一个索引文件 (.idx) 组成。这个过程由 tools/preprocess_data.py 脚本完成。

核心概念

输入格式: 该脚本期望的输入是一个 JSONL 文件，其中每一行都是一个 JSON 对象，且包含一个 text 字段。
```
{"text": "This is the first document."}
{"text": "This is the second, slightly longer document."}
```
文档 (Document): JSONL 文件中的每一行被视为一个独立的文档。
打包与拼接 (Packing & Concatenation):
- 脚本会逐一读取每个文档，使用你提供的分词器将其转换为 Token IDs。
- 在每个文档的 Token ID 序列末尾，添加一个特殊的 EOD (End-of-Document) token。
- 然后，将所有文档的 Token ID 序列（包含 EOD）拼接成一个巨大的一维数组。
切片成样本 (Slicing into Samples/Sequences):
- 这个巨大的一维数组会被依次切片成固定长度的块，这个长度就是模型的序列长度 (--seq-length)。每一个块就是一个训练样本。
- 关键点: 切片时不考虑文档边界。一个训练样本可能完整包含多个短文档，也可能只包含一个长文档的一部分，甚至可能跨越两个文档的边界（例如，前半部分是 doc_A 的结尾，后半部分是 doc_B 的开头）。模型需要通过 EOD token 自行学习文档的边界。

.bin 和 .idx 文件详解

.bin 文件 (数据文件):
- 这是一个纯二进制文件，内容就是那个巨大的一维 Token ID 数组。
- 其数据类型通常是 uint16 (如果 vocab_size < 65536) 或 int64，由 NumPy 数组直接写入。
- 它不包含任何元信息，只是原始数据。
.idx 文件 (索引文件):
- 这是实现高效随机访问的关键。它是一个索引文件，使得 dataloader 可以直接跳转到任意一个样本的起始位置，而无需扫描整个 .bin 文件。
- 文件结构:
  1. Magic Number: 文件开头的几个字节，用于标识文件类型（在 Megatron 中是 MMAP）。
  2. Version Number: 版本号，目前通常是 1。
  3. Data Type Code: 一个字节，表示 .bin 文件中数据的数据类型（如 uint16）。
  4. Document Index Header: 包含一些关于文档组织的信息，用于数据混合。
  5. Sample Pointer Array: 这是一个 uint64 类型的数组，是 .idx 文件的核心。
    - pointers[i] 的值是第 i 个样本在 .bin 文件中的起始字节偏移量。
    - pointers[0] 永远是 0。
    - 数组的最后一个元素是整个 .bin 文件的总大小（以字节为单位）。
    - 因此，第 i 个样本的数据位于 .bin 文件的 [pointers[i], pointers[i+1]) 字节范围内。

mmap 数据加载模式 (--data-impl mmap)

Megatron 强烈推荐使用 mmap (memory-map) 模式加载数据。它利用操作系统虚拟内存的特性，将磁盘上的 .bin 文件映射到进程的地址空间，但并不会立即将整个文件读入物理内存。只有当代码实际访问到某一块数据时，操作系统才会通过缺页中断机制，将对应的数据页从磁盘加载到内存中。

优点:
- 启动速度极快: 无需等待 TB 级的 .bin 文件被读完。
- 内存占用低: 物理内存只加载实际需要的数据，多个进程可以共享同一份物理内存页。
- 高效随机访问: 配合 .idx 文件，可以几乎无开销地跳转到任意样本。

3.4 数据混合与配比

现实世界中的预训练通常使用来自多个领域（如网页、书籍、代码、论文）的语料。如何智慧地混合这些数据，对模型的最终能力有显著影响。

Epoch 的重新定义: 对于海量数据集，传统意义上的 “epoch”（完整遍历一次数据集）不再适用。在大模型训练中，一个 “epoch” 通常被重新定义为处理一定数量的 tokens（例如，3000 亿 tokens）。
数据权重 (Data Weighting):
- 不同的数据源质量和价值不同。通常会给高质量、信息密度高的数据源（如高质量书籍、筛选过的代码）分配更高的权重，在混合时对其进行过采样 (over-sampling)。
- Megatron 的 preprocess_data.py 支持 --data-impl mmap 和 --split 参数，它在生成数据集时会为不同的数据源（来自不同输入文件的）记录文档索引，从而允许在训练时按权重进行采样。
温度采样 (Temperature Sampling):
- 如果严格按照权重比例来采样数据，可能会导致模型在训练初期长时间只看到某个高权重的数据源，不利于学习。
- 温度采样是一种更平滑的策略。假设数据源 i 的权重为 w_i，采样温度为 T。那么采样到数据源 i 的概率 P_i 正比于 w_i^(1/T)。
  - 当 T -> inf 时，所有数据源的采样概率趋于均匀。
  - 当 T -> 0 时，采样概率急剧倾向于权重最高的数据源。
  - 通常选择一个 T > 1.0 的值，以增加多样性，同时仍然偏向于高质量数据。

数据混合配置示例

数据源	原始 Tokens (Billion)	权重	采样后 Tokens (Billion)	备注
CommonCrawl (filtered)	1200	0.6	720	覆盖广泛，但质量偏低。
BooksCorpus2 + Wikipedia	250	1.0	250	高质量文本。
The Stack (Code)	400	0.8	320	提升代码和逻辑能力。
ArXiv (Papers)	80	1.2	96	提升学术和推理能力。
总计	1930	-	1386	假设我们的训练目标是 1.4T tokens。

3.5 在线/离线评估集构建与覆盖率检查

一个高质量的评估集（验证集/测试集）是衡量模型进步和做出正确研究决策的标尺。

构建原则:
1. 严格分离 (Strict Segregation): 评估集必须与训练集完全没有交集。这是最重要的一条原则。
2. 代表性 (Representativeness): 评估集应覆盖模型预期应用的各个领域和任务类型。
3. 高质量 (High Quality): 评估集本身需要经过精心的清洗和筛选，不应包含噪声或错误。
4. 适当规模 (Appropriate Size): 规模要足够大，以保证评估结果的统计显著性；但又不能太大，以免评估过程耗时过长。通常在几万到几十万样本之间。
数据污染检查 (Contamination Check):
- 问题: 训练数据（尤其是来自网络爬取的数据）可能无意中包含了评估集的内容，导致模型在评估时“作弊”，分数虚高。
- 检查方法:
  1. 从评估集中提取 n-grams (e.g., 8-grams or 13-grams)。
  2. 使用高效的字符串搜索工具（如 grep）或更高级的数据结构（如 Bloom Filter），在整个训练数据集中检查这些 n-grams 是否存在。
  3. 如果发现大量重叠，说明存在严重的数据污染，需要从训练集中剔除这些受污染的文档。
在线/离线评估:
- 在线评估 (Online Evaluation): 在训练过程中，每隔 N 个 iterations，在验证集上计算一次困惑度（Perplexity）等指标。这是监控训练进程、判断模型是否收敛或过拟合的主要手段。
- 离线评估 (Offline Evaluation): 在训练结束后，使用最终的模型检查点，在一系列下游任务的测试集（如 SuperGLUE, MMLU）上进行全面的性能评测。这是衡量模型通用能力的最终标准。

实践产出：`prepare_dataset.sh` + 统计报告

一个完整的、可复用的数据准备脚本，应该将上述流程串联起来。

prepare_dataset.sh (骨架)

#!/bin/bash
set -e # Exit immediately if a command exits with a non-zero status.

# --- Configuration ---
RAW_DATA_DIR="/path/to/raw_jsonl_files"
PROCESSED_DIR="/path/to/processed_data"
TOKENIZER_PATH="/path/to/my_tokenizer.model"
MEGATRON_LM_PATH="/path/to/Megatron-LM"
VOCAB_FILE="/path/to/vocab.json" # If needed by Megatron script
MERGE_FILE="/path/to/merges.txt" # If needed by Megatron script
OUTPUT_PREFIX="my_pretraining_dataset"

# --- 1. Data Cleaning and Deduplication ---
echo "Step 1: Cleaning and Deduplication..."
# This step is highly custom. Assume we have a Python script for this.
python custom_cleaner.py --input_dir ${RAW_DATA_DIR} --output_dir ${PROCESSED_DIR}/cleaned

# --- 2. (Optional) Train Tokenizer ---
# This is usually a one-time step.
# echo "Step 2: Training tokenizer..."
# spm_train --input=${PROCESSED_DIR}/all_cleaned.txt --model_prefix=my_tokenizer ...

# --- 3. Concatenate all cleaned files for Megatron ---
echo "Step 3: Concatenating data..."
find ${PROCESSED_DIR}/cleaned -name "*.jsonl" | xargs cat > ${PROCESSED_DIR}/all_in_one.jsonl

# --- 4. Run Megatron Preprocessing ---
echo "Step 4: Running Megatron's preprocess_data.py..."
python ${MEGATRON_LM_PATH}/tools/preprocess_data.py \
    --input ${PROCESSED_DIR}/all_in_one.jsonl \
    --output-prefix ${PROCESSED_DIR}/${OUTPUT_PREFIX} \
    --vocab-file ${VOCAB_FILE} \
    --merge-file ${MERGE_FILE} \
    --tokenizer-type GPT2BPETokenizer \
    --workers $(nproc) \
    --append-eod \
    --dataset-impl mmap

# --- 5. Generate Statistics Report ---
echo "Step 5: Generating statistics..."
# A custom script to analyze the generated .bin/.idx and report stats
python generate_stats.py --data_prefix ${PROCESSED_DIR}/${OUTPUT_PREFIX} > ${PROCESSED_DIR}/stats_report.txt

echo "Data preparation complete!"

统计报告 (stats_report.txt)

--- Dataset Statistics: my_pretraining_dataset ---
Date: 2025-11-03

Tokenizer Info:
- Type: SentencePiece BPE
- Vocab Size: 50257

Dataset Info:
- Total Documents: 15,234,890
- Total Tokens (approx.): 25,123,456,789
- Sequence Length: 2048
- Total Samples: 12,267,312

Data Source Distribution (if multiple files were processed):
- common_crawl.jsonl: 18.2B tokens (72.5%)
- books.jsonl: 4.1B tokens (16.3%)
- code.jsonl: 2.8B tokens (11.2%)

第三章总结

本章详细阐述了将原始、杂乱的文本转化为大模型可以直接“食用”的高质量“口粮”的全过程。我们从繁琐但至关重要的清洗与去重工作开始，强调了其对模型泛化能力的基础性作用。接着，我们深入探讨了分词器的选择与训练，特别是 BPE 算法和词表大小的权衡。核心部分在于对 Megatron 特有的 .bin/.idx 数据格式进行了深入解剖，解释了其设计原理和 mmap 加载机制的高效性。最后，我们讨论了在多源数据场景下的混合策略以及构建可靠评估集的重要性。完成本章的学习后，你应该具备独立构建一个完整、高效、可靠的大模型数据准备管线的能力。本章的学习后，你应该具备独立构建一个完整、高效、可靠的大模型数据准备管线的能力。