数据加载性能

前一章把数据做成了 token 文件，这一章解决在线侧的最后一个问题：从存储到 GPU 的搬运速度。训练每步要吃进 batch × seq × hidden 的 token 流，带宽不足时 GPU 断粮空转——"GPU 等 CPU"是最常见的隐形浪费。本章给一套完整的诊断与优化方法。

断粮的诊断：先确认是不是数据问题

症状与排查：

症状	判断
GPU 利用率锯齿状（满—空—满—空）	典型的数据饥饿
step time 方差大，CPU 侧 DataLoader 进程满载	数据管道是瓶颈
每步数据加载时间 > 计算时间	确定断粮
利用率平稳但整体低	不是数据问题，看MFU 分析

先量化：profiler（PyTorch profiler）看每步的 DataLoader 等待时间——超过 step time 的 5–10% 就值得优化。

瓶颈的三段：盘 → CPU → GPU

数据搬运的完整路径与各段的典型带宽：

存储（对象存储/并行FS/NVMe）  →  CPU 内存（mmap）  →  GPU（pinned → PCIe/NVLink）
     10–100 GB/s                    内存带宽               32–64 GB/s (PCIe)

每段的坑与解法：

一、存储段：预取与分片

训练语料放在离训练集群近的地方（本地 NVMe 或并行文件系统，见存储系统）；对象存储做底层仓库，训练前预热到本地
分片要够碎：分片大小 100–500 MB，配合多 worker 并行读；单一大文件是并行读取的天敌
顺序读：随机访问对象存储是灾难，token 流的设计（见数据管道）就是为了让训练只做顺序读

二、CPU 段：mmap 与多进程

mmap 而不是 read：token 流用内存映射，OS 页缓存自动管理，零拷贝零解析
多 worker 并行：DataLoader 的 worker 数按核数配（典型 4–16）；worker 数不足是最常见的配置错误
pin_memory：传输到 GPU 的缓冲用 pinned memory（锁页），否则 PCIe 传输要先拷进锁页中转，带宽减半；RDMA/GPUDirect 同样要求锁页（见 RDMA）
CPU 侧的处理最小化：token 流时代 CPU 只做"读 + 切 batch"，任何现做的 transform（tokenize、augment）都要挪到离线（见数据管道）

三、GPU 段：传输与重叠

non_blocking=True + pinned：tensor.to('cuda', non_blocking=True) 异步传输，与计算重叠
预取（prefetch）：第 n 步计算时，第 n+1 批已在传输——DataLoader 的 prefetch_factor、或手写双缓冲
传输量最小化：训练真正要上 GPU 的只有 token id（uint16/32）——1 个 token 2 字节，batch 8 × 8k 序列 = 128 KB/微批，PCIe 完全无压力。数据搬运的瓶颈几乎总在盘与 CPU，不在 PCIe

一套配置模板

生产训练的标准数据加载配置：

Dataset = mmap 的 token 流 + 按游标切微批（离线 tokenize，零在线处理）
num_workers = 8–16（按核数，覆盖 step 时间）
pin_memory = True
prefetch_factor = 4
persistent_workers = True      # 避免 epoch 间重启 worker
non_blocking = True            # H2D 异步

验证：profiler 确认每步的数据等待 < 5%；GPU 利用率曲线平稳（锯齿消失）。

与分布式的配合

每 rank 读自己的分片：DP 下每个进程读自己的数据分片（按 rank 切分），不做 all-broadcast
首个 rank 分发元数据：数据分片映射表由 rank 0 算好广播，各 rank 独立读自己的
CP/SP 的配合：序列切分（见上下文并行）时每卡读自己的段，数据管道按 CP 度切块
恢复对齐：数据游标（分片 + 偏移）进检查点，恢复时各 rank 精确回位（见检查点）

本章要点

数据饥饿的症状是 GPU 利用率锯齿；先 profiler 量化等待时间，超过 5% 就优化。
三段路径：存储（顺序读 + 本地预热）→ CPU（mmap + 多 worker + pinned）→ GPU（异步 + 预取）。
token 流时代 PCIe 不是瓶颈（每微批 128 KB）；瓶颈几乎总在盘与 CPU。
最常见的配置错误：worker 数不足、没 pin_memory、在线做本该离线的处理。
每 rank 读自己的分片；游标进检查点，恢复精确对齐。