第 3 章 数据流水线

第 3 章 数据流水线

本章目标：搞懂文本是怎么变成模型能吃的 (x, y) 张量对的。

3.1 章节地图

数据层干四件事：

┌─────────────────────────────────────────────────────────┐
│  1. 数据获取    获取文本（三级回退，鲁棒）               │
│  2. tiktoken 编码    文本 → token id 序列                │
│  3. 数据集切片   切片成 (x, y) 样本对                    │
│  4. 构造 DataLoader   批处理 + 工程优化                  │
└─────────────────────────────────────────────────────────┘
​

3.2 数据从哪来：三级回退策略

数据获取函数实现了一个鲁棒的数据获取链：

┌──────────────────────────────────────────┐
│ 1. 本地缓存 tiny_shakespeare.txt         │ ──► 命中？直接返回
└──────────────────────────────────────────┘
                  │ 未命中
                  ▼
┌──────────────────────────────────────────┐
│ 2. 框架数据集 tiny_shakespeare            │ ──► 成功？缓存并返回
└──────────────────────────────────────────┘
                  │ 失败（库缺失/下线）
                  ▼
┌──────────────────────────────────────────┐
│ 3. 直链下载原始文本                       │ ──► 缓存并返回
└──────────────────────────────────────────┘
​

核心代码骨架：

def get_dataset(cache_path: str = "data/tiny_shakespeare.txt") -> str:
    cache_path = Path(cache_path)
    cache_path.parent.mkdir(parents=True, exist_ok=True)

    # 1) 命中本地缓存
    if cache_path.exists():
        return cache_path.read_text(encoding="utf-8")

    text = None
    # 2) 尝试框架数据集
    if _HAS_DATASETS:
        try:
            ds = load_dataset("tiny_shakespeare", split="train")
            # 不同版本字段名可能为 'text'，统一取第一个字段兜底
            column = "text" if "text" in ds.column_names else ds.column_names[0]
            text = "\n".join(ds[column])
        except Exception as e:
            print(f"数据集下载失败: {e}")

    # 3) 回退：直链下载
    if text is None:
        with urllib.request.urlopen(_TINY_SHAKESPEARE_URL) as resp:
            text = resp.read().decode(encoding="utf-8")

    # 4) 写入本地缓存
    cache_path.write_text(text, encoding="utf-8")
    return text
​

工程亮点逐个讲

① 缓存优先

if cache_path.exists():
    return cache_path.read_text(encoding="utf-8")
​

避免每次训练都打网络，断网也能跑。这是「数据获取」的第一原则：网络是奢侈品，本地是必需品。

② 优雅降级（graceful degradation）

try:
    from datasets import load_dataset
    _HAS_DATASETS = True
except ImportError:
    _HAS_DATASETS = False
​

数据集库缺失时不崩溃，只把标志位置 False。后续 if 分支跳过该路径，直接走直链回退。这叫「可选依赖」处理。

对比反面教材：把 import 写在文件顶部不带 try/except，没装该库的用户一导入就报错，连回退机会都没有。

③ 字段兜底

column = "text" if "text" in ds.column_names else ds.column_names[0]
​

数据集不同版本字段名可能是 text 或 content 或别的。代码做了兜底：优先 text，否则取第一个字段。这是处理「外部数据源不稳定」的常见手法。

④ Path 替代字符串拼接

cache_path = Path(cache_path)
cache_path.parent.mkdir(parents=True, exist_ok=True)
​

pathlib.Path 比 os.path.join 更现代。mkdir(parents=True, exist_ok=True) 一行搞定「目录不存在就建、存在不报错」。

3.3 分词：tiktoken 与 BPE

为什么需要分词

模型不认识字符，只认识数字。分词器（tokenizer）把文本切成 token（子词单元），再映射成整数 id：

"To be, or not to be"
   │ tiktoken.encode
   ▼
[787, 281, 11, 453, 459, 281, 281]   ← token id 序列
​

三种分词粒度对比

GPT 系列用子词级（BPE），是当前主流。

BPE 工作原理（科普）

Byte-Pair Encoding：训练时统计语料里最频繁的相邻字节对，把它们合并成一个新 token，重复 N 次。

初始：['T','h','e',' ','c','a','t']  ← 都是单字符
       ↓ 合并 'c'+'a'='ca'（出现频繁）
       ['T','h','e',' ','ca','t']
       ↓ 合并 'ca'+'t'='cat'
       ['T','h','e',' ','cat']
​

最终词表是「单字符 + 高频合并子词」的集合。新词用已有子词拼："unbelievable" → ["un", "believable"]，"helloworld" → ["hello", "world"]——没学过的词也能拼出来。

为什么选 tiktoken + p50k_base

• tiktoken 是 OpenAI 开源的高性能 BPE 分词器（Rust 实现，比 HuggingFace 快几倍）。
• p50k_base 是 GPT-2 / text-davinci-002 用的编码，词表大小恰好 50257，与本项目 GPTConfig.vocab_size 对齐。

全局单例编码器

_ENCODER = tiktoken.get_encoding("p50k_base")

def get_encoder() -> "tiktoken.Encoding":
    return _ENCODER
​

为什么用模块级全局变量？因为 tiktoken 初始化时要加载一张不小的 BPE 合并表（几十 MB），反复 get_encoding 会重复加载。用全局单例一次加载、处处复用，省内存省时间。

⚠️ 铁律的来源：vocab_size=50257 是 p50k_base 决定的，改它会导致 token id 与 embedding 矩阵的行数对不上，模型直接崩。换词表 = 换分词器，两者必须同步。

3.4 数据集切片：自回归样本构造（本章核心）

语言模型的训练目标

预测下一个 token。给定 ['T','o',' ','b','e']，模型要预测下一个是 ,。所以要把文本切成 (x, y) 对，其中 y 是 x 左移一位：

class TextDataset(Dataset):
    def __init__(self, text: str, block_size: int):
        self.block_size = block_size
        self.tokens = _ENCODER.encode(text)
        self.n_samples = max(0, len(self.tokens) - block_size - 1)

    def __getitem__(self, idx):
        chunk = self.tokens[idx : idx + self.block_size + 1]
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:],  dtype=torch.long)
        return x, y
​

切片图解

假设 block_size=4，token 序列为 [A, B, C, D, E, F, G, ...]：

位置:  0   1   2   3   4   5   6
token: A   B   C   D   E   F   G   ...
       └────── x ──────┘
           └────── y ──────┘   ← y 整体右移一位

样本 idx=0:  x=[A,B,C,D]  y=[B,C,D,E]
样本 idx=1:  x=[B,C,D,E]  y=[C,D,E,F]
样本 idx=2:  x=[C,D,E,F]  y=[D,E,F,G]
...
​

三个关键细节

细节 1：取 block_size+1 长度的 chunk

为什么 chunk 是 block_size+1 而不是 block_size？因为要同时容纳 x（block_size 个）和 y（block_size 个，与 x 错开一位），总共需要 block_size+1 个 token。

细节 2：y = chunk[1:] 不是 chunk[:-1]

x 是「输入」，y 是「目标」，y 比 x 滞后一位：位置 i 的 x[i] 要预测 y[i] = x[i+1]（下一个 token）。所以 y 是 x 整体右移一位，用 chunk[1:]。

细节 3：样本数 = 总 token 数 − block_size − 1

每个样本要取 block_size + 1 长度的 chunk，最后一个完整 chunk 起点最多到 len - block_size - 1。max(0, ...) 是兜底：万一文本特别短，返回 0 而不是负数。

一个 block 同时训练 block_size 个预测任务

注意：一个样本有 block_size 个位置，但不是「一个样本 = 一个预测任务」。由于 GPT 的因果掩码（causal mask，见第 4 章），位置 0 只看自己预测位置 1，位置 1 看 [0,1] 预测位置 2……所有 block_size 个位置的预测任务同时在一个前向里完成，loss 是它们的平均。

x   =  [A,  B,  C,  D]
位置:   0   1   2   3
         │   │   │   │
         ▼   ▼   ▼   ▼
预测:    B   C   D   E    ← 4 个预测任务同时进行
位置 0 看 [A]     预测 B
位置 1 看 [A,B]   预测 C
​

这就是 Transformer 训练高效的根本原因——一次前向 = block_size 个训练信号。

滑动窗口 vs 不重叠切片

本项目用步长 1 的滑动窗口：相邻样本高度重叠（共享 block_size-1 个 token）。优点是数据多（33 万 token 能切 33 万样本），缺点是相邻样本相关性强。工业训练常用步长 = block_size（不重叠）+ 多 epoch，减少冗余；本项目为简单用步长 1，配合 shuffle 缓解相关性。

3.5 构造 DataLoader：批处理与工程优化

def build_dataloader(text, block_size, batch_size, num_workers=0, shuffle=True):
    dataset = TextDataset(text, block_size)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        drop_last=True,                              # ① 丢弃不完整 batch
        pin_memory=torch.cuda.is_available(),        # ② 锁页内存加速 H2D 拷贝
    )
​

工程细节 1：drop_last=True

最后一个不完整的 batch（比如 batch_size=32 但只剩 5 个样本）会被丢掉。为什么？不完整 batch 的梯度统计与完整 batch 不同，可能让训练末期出现 loss 抖动。丢掉它保持每个 batch 都是 32 个样本，梯度尺度一致。

工程细节 2：pin_memory=True

锁页内存（page-locked memory）：让数据预先放到 CPU 的「不可换页」内存区，CPU→GPU 拷贝能走 DMA 异步通道，比普通内存快。配合训练循环里的 x.to(device, non_blocking=True) 实现「数据传输」与「上一步反向传播计算」的重叠，免费的加速。

工程细节 3：num_workers 的取舍

3.6 数据流串起来

把所有环节串一遍，跟踪一个 batch 的生命周期：

1. 数据获取读取本地缓存
   → 返回 1MB 的纯文本字符串

2. 数据集初始化把文本编码为 token：
   encode(text) → [787, 281, 11, ..., 338025 个 id]
   n_samples = 338025 - 128 - 1 = 337896

3. DataLoader 用 batch_size=32 shuffle=True 包装
   内部生成打乱的索引 [12345, 892, 337000, ...]

4. 训练循环 for x, y in dataloader:
   每次 next() 触发：
   - 按索引取 32 个样本
   - 每个样本切出 (x_i, y_i) 各 shape=(128,)
   - 自动堆叠成 x.shape=(32, 128), y.shape=(32, 128)
   - pin_memory 把它们放到锁页内存

5. x.to(device, non_blocking=True)
   异步拷贝到 GPU

6. model(input_ids=x, labels=y)
   开始前向 + 算 loss
​

3.7 动手实验

1. 体验 tiktoken：在 Python 里跑 from dataset import get_encoder，用 enc.encode("Hello, world!")、enc.decode([...])、enc.encode("unbelievable")、enc.encode("你好世界") 观察子词与中文切分。
2. 构造数据集：TextDataset("Hello world this is a test", block_size=4)，打印 len(ds) 和 ds[0]，观察 x 和 y 的关系。
3. 边界测试：故意把 block_size 设成比总 token 数还大，观察 n_samples 会变成多少（答：0，会被 max(0, ...) 兜底）。
4. 思考题：如果想让两个相邻样本完全不重叠（步长 = block_size 而非 1），该怎么改 __getitem__？
   • 提示：返回 tokens[idx*block_size : idx*block_size + block_size + 1]，且 n_samples = len(tokens) // (block_size + 1)。

3.8 小结

• 数据获取用三级回退（本地缓存 → 框架数据集 → 直链下载）保证数据获取鲁棒，断网也能跑。
• tiktoken 的 p50k_base 把文本切成子词 id 序列，词表 50257 与模型 vocab_size 对齐。
• 数据集切片取 block_size+1 长度的 chunk，切成 (x, y) = (chunk[:-1], chunk[1:])，y 是 x 左移一位。
• 一个 block 同时训练 block_size 个预测任务（因果掩码保证只看过去）。
• DataLoader 的 drop_last 保持 batch 尺度一致，pin_memory + non_blocking 实现数据传输与计算重叠。

3.9 下一章

数据准备好了，去《第 4 章 模型构建》看怎么搭一个 GPT。