跳到主要内容

数据截至 (上游 commit 3adf61e154c3)

03 · 数据管线：从文本到 GPU 里的 batch

这一章讲什么： nanoGPT 的数据侧极简到近乎粗暴——整个语料压成一个连续 token id 数组，训练时随机切片。本章讲清 data/*/prepare.py 产出的磁盘格式、get_batch 的读取机制，以及这套「穷人数据管线」省掉了什么、放弃了什么。

1. 它要解决的小问题​

语言模型训练的数据需求其实只有一句话：源源不断地产出形状为 (batch_size, block_size) 的整数张量 x，以及把 x 往左错一位的标签 y。

工业数据管线会引入 Dataset/DataLoader/worker 进程/采样器一整套抽象。nanoGPT 的问题是：这件事的最小实现长什么样？答案是：一个二进制文件 + 一个 15 行的函数。

2. 磁盘格式：train.bin / val.bin / meta.pkl​

三个 prepare 脚本产出同一套格式，分两档规模：

数据集tokenizer产物规模
shakespeare_char字符级（65 个字符手排）train/val.bin + meta.pkltrain 约 1M token（data/shakespeare_char/prepare.py:63-68 注释）
shakespearetiktoken GPT-2 BPEtrain/val.bin（无 meta）train 约 30 万 token（data/shakespeare/prepare.py:32-33 注释）
openwebtexttiktoken GPT-2 BPEtrain/val.bin（无 meta）train 约 90 亿 token / ~17GB（data/openwebtext/prepare.py:76-78 注释）

格式本身没有任何头部或索引——就是一串 uint16 小端整数，每个整数是一个 token id。

为什么 uint16 够： GPT-2 词表最大 id 是 50256 < 65536（data/openwebtext/prepare.py:62 的注释）。17GB 而不是 34GB，纯靠省字节。

meta.pkl 的角色： 只有字符级数据集需要它——里面存 vocab_size 和手排的 stoi/itos 字符↔整数映射（data/shakespeare_char/prepare.py:55-61）。训练脚本用它决定 vocab_size（train.py:138-144），采样脚本用它做编解码（见第 4 章）。用 tiktoken 的数据集不需要 meta，因为编解码规则是 tiktoken 内置的。

3. 机制一：prepare.py——文本怎么变成大数组​

以 openwebtext/prepare.py 为例，流程是：

  1. 下载：HF datasets 拉 openwebtext，约 800 万篇文档（data/openwebtext/prepare.py:23）。
  2. 切分：原数据集只有 train，手动切 0.0005 当 val（data/openwebtext/prepare.py:26-27）。
  3. tokenize：每篇文档 enc.encode_ordinary(text) 编码，末尾追加 eot（end-of-text）token 50256（data/openwebtext/prepare.py:43-48）。encode_ordinary 会忽略文本里恰好出现的特殊 token，防止注入。
  4. 拼接落盘：预先算好总长度，开一个 np.memmap 写模式文件，分 1024 个 shard 逐段写入（data/openwebtext/prepare.py:59-74）。

关键设计：文档之间没有分隔符，只有 eot token。 拼完之后整个语料就是一条 90 亿 token 的长河，模型靠 eot 学「一篇文章到此为止」。作者还在注释里留了个悬案：「I think eot should be prepended not appended... hmm」（data/openwebtext/prepare.py:46）。

shakespeare_char/prepare.py 则是另一个极端：字符排序后编号，stoi/itos 两个字典就是全部 tokenizer（data/shakespeare_char/prepare.py:30-35）。

4. 机制二：get_batch——随机切片 + 错位标签​

小问题： 一条连续 token 长河，怎么变成训练样本？

答案：不切片、不打包，直接随机裁。 get_batch（train.py:116-131）全文值得一看：

def get_batch(split):
    # We recreate np.memmap every batch to avoid a memory leak, as per
    # https://stackoverflow.com/questions/45132940/numpy-memmap-memory-usage-want-to-iterate-once/61472122#61472122
    if split == 'train':
        data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
    else:
        data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
    if device_type == 'cuda':
        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    else:
        x, y = x.to(device), y.to(device)
    return x, y

逐层拆：

  • memmap 当读接口：17GB 文件不加载进内存，操作系统按页懒加载（train.py:120）。
  • 随机起点：torch.randint 在整条长河上均匀采 batch_size 个起点（train.py:123）。没有 epoch、没有洗牌、没有「每条数据看一遍」的概念——训练 600k 步就是随机抽 600k 次。
  • 错位一个 token 造标签：y 就是 x 往右挪一位（train.py:124-125）。语言模型训练的「下一个 token 预测」就体现在这个 i+1 里。
  • 类型升格：uint16 读出来 .astype(np.int64)，因为 embedding 索引要 long（train.py:124）。
  • 异步上卡：pin_memory() + non_blocking=True 让 CPU→GPU 拷贝不阻塞主线程（train.py:127-128），配合主循环里「前向时预取下一批」（train.py:302-303）把数据加载完全藏进计算里。

坑： 每次调用都重建 memmap 对象——注释说这是为了绕开 numpy memmap 的一个内存泄漏，并附了 StackOverflow 链接（train.py:117-118）。这是「读真代码才能学到的坑」的典型：单看 API 你永远不会想到要这么做。

5. 这套管线放弃了什么​

诚实清单（相对工业数据管线）：

  • 没有文档边界对齐。 随机切片会横切文档中间、eot 会出现在序列任何位置。对预训练这基本无害，但这不是「每条样本 = 一篇完整文档」的语义。
  • 没有去重、没有质量过滤。 语料准备完全是 prepare 脚本外的责任。
  • 没有确定性采样顺序。 随机抽意味着无法保证「整个数据集被均匀覆盖 N 遍」，也没有可复现的样本顺序（除非固定种子）。
  • 没有动态 batching / 长度分桶。 所有样本定长 block_size。
  • val 集参与方式很粗。 val loss 是从 val.bin 上同样随机切片估的（train.py:219-226），不是严格的一次性遍历。

这些放弃全部服务于同一个目标：数据侧代码压到 15 行，且对 90 亿 token 依然 O(1) 内存。

6. 本章代码地图​

主题文件符号
字符级准备data/shakespeare_char/prepare.py:30-61stoi、itos、meta
OpenWebText 准备data/openwebtext/prepare.py:43-74process、memmap 写入循环
eot 追加data/openwebtext/prepare.py:44-45enc.encode_ordinary、enc.eot_token
uint16 落盘data/openwebtext/prepare.py:62-63np.uint16、np.memmap
BPE 版莎士比亚data/shakespeare/prepare.py:20-30enc.encode_ordinary
取 batchtrain.py:116-131get_batch
vocab 探测train.py:138-144meta_vocab_size