【大模型】- 数据管道：从文本到训练

数据管道：从文本到训练

构建真实 LLM 使用的数据加载器

类型： 学习　｜　语言： Python　｜　🏷 前置：《分词器：BPE、WordPiece、SentencePiece》（本系列第 1 篇）

学习目标

了解 LLM 训练数据的来源和格式
实现一个将原始文本转换为训练样本的分词数据集
使用滑动窗口创建下一个 token 预测的训练对
构建高效的 DataLoader 用于批量训练
处理填充值和掩码以支持变长序列

为什么数据管道很重要

大型语言模型在海量文本数据上训练。原始文本需要经过几个转换才能成为训练批次：文本被分词成数字 ID，然后组织成固定长度的序列，最后分成输入和目标。这个管道的效率直接影响训练速度。

数据来源

现代 LLM 使用多种来源训练：

网络抓取：Common Crawl、C4、The Pile
书籍：Books3、Gutenberg
代码：GitHub repositories
学术论文：ArXiv、Semantic Scholar
对话：Reddit、论坛

对于学习，我们使用小型数据集。Hugging Face 的 datasets 库可以轻松加载它们。

数据格式

原始文本通常存储为：

纯文本文件：每行或每段一个文档
JSON/JSONL：每行一个 JSON 对象
Parquet：列式存储，大数据集效率高
CSV：简单但不适合大型文本

Tokenization Pipeline

1	Raw Text → Tokenize → Create Sliding Windows → Form Input-Target Pairs → Batch

步骤 1：加载数据

from datasets import load_dataset

dataset = load_dataset("text", data_files="data.txt")
# 或使用预制数据集
dataset = load_dataset("wikitext", "wikitext-2-raw-v1")

步骤 2：分词

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")

def tokenize_function(examples):
    return tokenizer(examples["text"])

步骤 3：创建滑动窗口

def create_windows(examples, block_size=128):
    # 连接所有文本
    concatenated = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated[list(examples.keys())[0]])
    
    # 按 block_size 分割
    result = {
        k: [t[i:i + block_size] 
            for i in range(0, total_length, block_size)]
        for k, t in concatenated.items()
    }
    return result

步骤 4：形成输入-目标对

def group_texts(examples):
    # 每个块：前 n-1 个 token 作为输入，后 n-1 个 token 作为目标
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    
    # 调整到 block_size 的倍数
    total_length = (total_length // block_size) * block_size
    
    result = {}
    for k, t in concatenated_examples.items():
        result[k] = [
            t[i:i + block_size]  # 输入
            for i in range(0, total_length, block_size)
        ]
    
    # 目标是输入右移一位
    result["labels"] = [
        t[1:i + block_size + 1]
        for i in range(0, total_length, block_size)
    ]
    
    return result

构建完整的 DataLoader

import torch
from torch.utils.data import DataLoader

# 应用分词和分组
tokenized = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"]
)

grouped = tokenized.map(
    group_texts,
    batched=True
)

# 转换为 PyTorch 格式
grouped.set_format(type="torch", columns=["input_ids", "attention_mask"])

# 创建 DataLoader
dataloader = DataLoader(grouped["train"], batch_size=32, shuffle=True)

填充和掩码

由于序列长度不同，批次内的样本需要填充到相同长度：

from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    input_ids = [item["input_ids"] for item in batch]
    labels = [item["labels"] for item in batch]
    
    # 填充到批次内最大长度
    input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=0)
    labels_padded = pad_sequence(labels, batch_first=True, padding_value=-100)
    
    # 创建注意力掩码（忽略填充值）
    attention_mask = (input_ids_padded != 0).long()
    
    return {
        "input_ids": input_ids_padded,
        "attention_mask": attention_mask,
        "labels": labels_padded
    }

dataloader = DataLoader(
    grouped["train"],
    batch_size=32,
    collate_fn=collate_fn
)

内存映射大数据集

对于大到无法放入内存的数据集：

# 使用内存映射（Hugging Face 默认）
dataset = load_dataset("text", data_files="large.txt", streaming=True)

# 或使用 WebDataset 格式
import webdataset as wds

dataset = (
    wds.WebDataset("data-{000..099}.tar")
    .decode("pil")
    .to_tuple("input_ids", "labels")
    .batched(32)
)

数据质量检查

在训练之前检查数据质量：

def analyze_dataset(dataset):
    total_tokens = sum(len(x["input_ids"]) for x in dataset)
    avg_length = total_tokens / len(dataset)
    vocab_size = len(set(token for x in dataset for token in x["input_ids"]))
    
    print(f"总样本数: {len(dataset):,}")
    print(f"总 token 数: {total_tokens:,}")
    print(f"平均序列长度: {avg_length:.1f}")
    print(f"词汇表大小: {vocab_size:,}")
    print(f"估计大小: {total_tokens * 2 / 1e9:.2f} GB (float32)")

常见陷阱

数据泄漏：确保训练/验证/测试集不重叠
tokenization 不一致：训练和推理使用相同的 tokenizer
序列边界：正确处理文档边界，避免跨文档污染
内存使用：大数据集使用流式处理或内存映射
随机性：设置随机种子以确保可重复性

性能优化

# 多进程加载
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,  # 多进程
    pin_memory=True,  # GPU 加速
    persistent_workers=True
)

# 预取
prefetch_factor=2  # 每个工作进程预取 2 个批次

总结

数据管道将原始文本转换为训练就绪的批次。关键步骤包括分词、滑动窗口、输入-目标对和高效加载。优化管道可以显著加速训练。

下一步

下一课将构建一个小型 GPT 模型，使用这个管道进行预训练。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

📝 自我检查（课程配套测验）

Q1（学前） 为什么不能简单地把所有预训练数据加载到内存？

A. Python 不支持大数组
B. 预训练语料达 TB 级，远超可用 RAM，需要流式流水线
C. 加载到内存更慢
D. 内存只用于模型权重

答案： B　解析： LLM 预训练数据通常为 1–15 TB 文本。即使有 256GB RAM 也无法容纳全量数据。流式流水线按需处理，只加载当前 batch 所需数据。

Q2（学前） 为什么预训练需要去重？

A. 节省磁盘空间
B. 重复文档使模型逐字记忆特定文本，并在重复内容上浪费训练算力
C. 加快分词
D. 减小词表

答案： B　解析： 近重复内容（样板、抓取重复）使模型记忆而非泛化。去重减少训练算力浪费，通过确保多样化训练信号提升模型质量。

Q3（学后） 为什么要将变长文档整理为固定长度训练序列？

A. 让文本更易读
B. GPU 训练需要统一张量形状，文档必须打包或填充为固定长度序列
C. 固定长度序列更准确
D. 减少总 token 数

答案： B　解析： GPU 处理形状相同的张量 batch。变长文档必须切成固定长度序列（如 2048 或 4096 token），并在文档边界设置正确的注意力掩码。

Q4（学后） 若数据流水线慢于 GPU 训练速度会怎样？

A. 训练自动降速匹配
B. GPU 空等 batch，浪费昂贵算力
C. 模型在同一 batch 上反复训练
D. 无影响——流水线异步运行

答案： B　解析： 若 dataloader 无法足够快提供 batch，GPU 在步间停滞。A100 集群每小时 $30+，流水线瓶颈直接浪费金钱。分析流水线吞吐量至关重要。

Q5（学后） 为什么在分词前要做数据质量过滤（语言检测、内容过滤）？

A. 分词器无法处理低质量文本
B. 低质量数据（垃圾、样板、有毒内容）按其在训练数据中的占比同等降低模型能力
C. 分词后无法过滤
D. 减少分词时间

答案： B　解析： 模型从所见数据学习。若 10% 训练数据是垃圾或低质量，模型会分配 10% 容量复现这些模式。早期过滤确保只有高质量信号进入模型。