跳到主要内容

数据截至 (上游 commit f1e2ace65149)

03 · 数据加载与分词

这一章讲什么： 训练数据怎么从一堆文本变成 C 能高效消费的 batch。整条管线：Python 脚本把 文本 tokenize 成 .bin → DataLoader 按分片读、按 rank 切 → batch 的 inputs/targets 错位一 位。外加多选评测专用的 EvalLoader 和只会解码的 Tokenizer。

1. 它要解决的小问题​

深度学习框架里你写 DataLoader(dataset, batch_size=..., shuffle=True) 就完事。没有框架时， 「给我一个 batch」具体要做哪些事？

llm.c 的回答只有三层：定格式（.bin 长什么样）、定切法（多卡各自读哪段）、定读法 （一个 batch 怎么从文件里 fseek 出来）。

2. .bin 格式：1024 字节文件头 + token 流​

2.1 布局​

所有数据文件同一格式（写入方：write_datafile，dev/data/data_common.py:39-60；读取方： dataloader_load_shard_，llmc/dataloader.h:61-94）：

区域内容
header[0]magic = 20240520
header[1]version = 1
header[2]ntok（本文件 token 数）
header[3..255]保留（共 256 个 int32 = 1024 字节）
正文ntok 个 uint16 token id（GPT-2 BPE）

读取侧的防御很到位：magic/version 不对即报错并提示重新跑预处理；还会按 1024 + ntok*2 核对文件大小（llmc/dataloader.h:75-94）。

2.2 谁来写 .bin​

dev/data/ 下每个数据集一个脚本。以 dev/data/tinyshakespeare.py 为例：下载原文 → 按空行切 「文档」，每篇前面插一个 <|endoftext|> → tiktoken 编码 → 前 32768 个 token 作验证集，其余作 训练集。文档之间靠 EOT token 分界——训练流里根本没有「文档边界」概念，只有一条无限 token 河，这是 GPT 式预训练数据的标准抽象。

另有 fineweb.py（10B/100B 大规模预训练语料）、tinystories.py、hellaswag.py、mmlu.py （后两个写评测格式，见 §5）。

2.3 uint16 的含义与上限​

uint16 能表达 0~65535，GPT-2 词表 50257 放得下。data_common.py 里还定义了 llama-3 格式 （magic 20240801、uint32 token，dev/data/data_common.py:26-35），但 C 侧 DataLoader 只认 magic 20240520 的 uint16 格式——llama-3 的 .bin 是给 Python 侧的 train_llama3.py 用的，C 训练主线全程 GPT-2 词表（inferred：C 侧没有任何读取 uint32 token 的代码路径）。

3. DataLoader：分布式切分就是一次偏移​

3.1 核心数学​

dataloader_init（llmc/dataloader.h:142-201）里三行决定了全部分布式行为 （llmc/dataloader.h:156-158）：

loader->total_batch_size_bytes  = num_processes * B * T * sizeof(uint16_t);  // 全局 batch
loader->local_batch_offset_bytes = process_rank * B * T * sizeof(uint16_t);  // 本 rank 偏移

所有 rank 把数据看作同一条 token 河：第 i 个「全局 batch」是河里的第 i 段 num_processes*B*T 个 token，rank r 每次舀其中属于自己的 B*T。不需要广播、不需要中心调度， 每个进程独立 fseek 到 header + i*全局 + r*局部 读自己的那段（dataloader_load_batch， llmc/dataloader.h:203-219）。

3.2 B*T+1 的错位​

读 batch 时实际读 B*T+1 个 token（llmc/dataloader.h:213-214）：

for (int i = 0; i < B*T; i++) {
    loader->inputs[i]  = (int)loader->buffer[i];      // 第 i 个 token
    loader->targets[i] = (int)loader->buffer[i+1];    // 第 i+1 个 token
}

（llmc/dataloader.h:216-219）语言模型的「输入」和「标签」本就是同一条序列错开一位——所以 只需一次读取。这也解释了 §2.2 为什么说数据里没有文档结构：跨界处的预测目标就是下一个 token， 不管它是否属于下一篇文档。

3.3 分片管理与两级 shuffle​

  • 分片：-i 参数是 glob 模式（如 fineweb_train_*.bin），初始化时 glob 出全部分片并逐 个校验（llmc/dataloader.h:163-190）；任何时刻只打开一个分片文件，读完用 dataloader_advance_ 切下一个（llmc/dataloader.h:125-140），到最后一个就 reset 开新 epoch。
  • 两级 shuffle（should_shuffle=1 时）：分片顺序洗一次（shard_indices），每片内部的全局 batch 顺序再洗一次（intra_shard_indices，llmc/dataloader.h:87-96）。RNG 是 mt19937，种 子 42 + process_rank——各 rank 的洗牌序列必须一致，否则大家读的就不是同一个全局 batch 了；llm.c 的做法是让洗牌作用于「全局 batch 下标」，各 rank 用同样种子推出同样的序 列，再各取各的偏移（种子设定见 llmc/dataloader.h:172，重洗见 :110-123）。
  • 断点恢复：dataloader_resume 直接跳到 (shard_idx, sample_idx)（llmc/dataloader.h:232-237）， shuffle 状态（两个排列数组 + RNG 状态）整体存进 checkpoint（见第 4 章 §6）。

3.4 约束​

代码里写死的假设，越界即 assert：

  • 每个分片至少装得下「全体 rank 的一个全局 batch + 1」个 token（llmc/dataloader.h:186）。
  • 每个 batch 的起点必须是全局 batch 边界——变长序列、packing 之类的花样一概没有。

4. Tokenizer：只会解码​

llmc/tokenizer.h 实现的是 GPT-2 tokenizer 的一半：

  • 读 gpt2_tokenizer.bin（magic 20240328，llmc/tokenizer.h:41-84）：文件头给词表大小和 EOT id，随后是 50257 条「长度字节 + 原始字节」的记录。
  • tokenizer_decode（llmc/tokenizer.h:86-96）：token id → 字符串，就是查 token_table。
  • 没有 encode——文件头注释明说：只做无条件生成的话解码就够；要做 prompt 得加编码，而 BPE 的 正则预处理「在 C 里不好搞」（llmc/tokenizer.h:1-7）。
  • safe_printf（llmc/tokenizer.h:25-38）：单字节的怪 token（控制字符、退格）不打印——采 样输出到终端时防炸屏。

这个「半 tokenizer」精确标定了项目边界：llm.c 是预训练 + 无条件采样，不是聊天推理框架。

5. EvalLoader：把多选评测塞进 (B,T) 张量​

5.1 评测文件的格式​

HellaSwag/MMLU 这类评测是「一段上下文 + 4 个候选续写 + 正确答案下标」。文件格式（文档注释在 llmc/dataloader.h:252-265，写入方 write_evalfile，dev/data/data_common.py:62）：

header: magic 20240522 | version | num_examples | longest_example_bytes
每个 example: <START=65535> <EXAMPLE_BYTES> <EXAMPLE_INDEX> <LABEL>
              <NUM_COMPLETIONS> <context_len><context tokens>
              <comp_len><comp tokens> × 4

<EXAMPLE_BYTES> 让加载器可以不解内容直接 fseek 跳过整例（llmc/dataloader.h:316-329）。

5.2 打包与判分​

  • 打包：一条例子占 batch 的 4 行（每行 = 共享的上下文 + 一个候选续写），B=64 时一个 batch 装 16 条例子（evalloader_next_batch，llmc/dataloader.h:449-466）。上下文部分 4 行相同， 续写部分各行不同，并用 mask 标出「续写 token 的预测位置」。
  • 判分：前向得到每个位置的 loss 后，evalloader_stat_losses （llmc/dataloader.h:468-508）对每条例子的 4 行分别算 mask 区的平均 loss，选平均 loss 最低的那行——若等于 label 则记对。这正是 HellaSwag 的 acc_norm 标准做法：模型认为哪个续 写最「不意外」，哪个就是它的答案。
  • 分布式：各 rank 按例子下标均分，最后 multi_gpu_cpu_float_sum 汇总（train_gpt2.cu 的 HellaSwag 段，train_gpt2.cu:1736-1748）。

6. 关键细节与坑​

  • uint16 是硬约束。 词表超过 65535 的模型（如 Llama-3 的 128256）进不了这条 C 数据管线， 对应数据脚本写的是另一套 magic/格式（§2.3）。
  • 小分片会直接 assert。 分片 token 数不足一个全局 batch 时初始化就失败 （llmc/dataloader.h:186）；超多小分片的数据集需要先合并。
  • CPU 参考实现不洗牌也能跑。 train_gpt2.c 里 train loader 的 should_shuffle=1、val 为 0（train_gpt2.c:1093-1094）；overfit 单 batch 调试模式（-a 1）会强制关掉 train 的洗牌 （train_gpt2.cu:1602-1604）。
  • 评测假设恰好 4 个候选。 ASSUMED_NUM_COMPLETIONS 写死为 4，B 小于 4 时直接报错退出 （llmc/dataloader.h:270、:305-310）。
  • tokenizer 文件缺失只是警告。 找不到 gpt2_tokenizer.bin 时 init_ok=0，采样退化为打 印 token id，训练不受影响（llmc/tokenizer.h:43-53）。

7. 代码地图​

主题文件路径符号名
.bin 训练格式（写）dev/data/data_common.pywrite_datafile、HEADERS_INFO
.bin 训练格式（读）llmc/dataloader.hdataloader_load_shard_、HEADER_SIZE
训练 DataLoaderllmc/dataloader.hdataloader_init、dataloader_reset、dataloader_next_batch、dataloader_load_batch、dataloader_advance_、dataloader_resume
两级 shufflellmc/dataloader.h、llmc/rand.hprepare_intra_shard_indices_、random_permutation、mt19937_state
评测格式（写）dev/data/data_common.py、dev/data/hellaswag.pywrite_evalfile
EvalLoaderllmc/dataloader.hevalloader_init、evalloader_reset、evalloader_next_batch、evalloader_stat_losses
分词llmc/tokenizer.htokenizer_init、tokenizer_decode、safe_printf
数据集脚本dev/data/tinyshakespeare.py、dev/data/fineweb.py、dev/data/tinystories.pytokenize