📑 本页目录（点开跳转）
22 · 数据管线与存储
⏱ 24 分钟 ｜ 🎁 最容易被忽略的瓶颈
🎯 一句话
你花了几个月优化模型和并行策略，结果 GPU 有 30% 的时间在等数据。 第 4 章那个"把 dataloader 换成随机张量"的实验之所以排第一位， 就是因为这个问题极其常见，而且极其容易被误诊成"模型慢"。
🚨 一、五分钟确诊
# 🧩 骨架：`vocab` 来自你自己的代码，这一段只看写法
# ⭐ 实验一：拿掉数据管线
import time
import torch
fake = torch.randint(0, vocab, (B, S), device='cuda')
for _ in range(50):
train_step(fake) # 固定张量，完全没有数据加载
# → 如果 MFU 显著上升，瓶颈就在数据管线
# ⭐ 实验二：只跑 dataloader，不训练
t0 = time.time()
for i, batch in enumerate(loader):
if i >= 100: break
print(f"数据管线吞吐 {100 * B / (time.time() - t0):.0f} samples/s")
# → 和训练需要的速度对比
对照
⭐ 判据：数据管线的吞吐必须 ≥ 训练吞吐的 1.2 倍
否则 GPU 一定会有等待
💡 profiler 里的信号（第 4 章）： GPU 时间轴上有规律的空隙，且空隙出现在每个 step 的开头 —— 就是数据没跟上。
🔧 二、DataLoader 的六个参数
# 🧩 骨架：`DataLoader` 来自你自己的代码，这一段只看写法
loader = DataLoader(
dataset,
batch_size=32,
num_workers=8, # ⭐ 通常设为 CPU 核数 / GPU 数
pin_memory=True, # ⭐ 锁页内存，让 H2D 拷贝能异步
prefetch_factor=4, # 每个 worker 预取几个 batch
persistent_workers=True,# ⭐ 别每个 epoch 重建 worker（重建很贵）
drop_last=True, # 避免最后一个不完整 batch（分布式下尤其重要）
)
| 参数 | 常见错误 |
|---|---|
num_workers=0 |
主进程加载 → 必然阻塞 ⭐ |
num_workers 太大 |
内存爆、CPU 争抢，反而更慢 |
pin_memory=False |
H2D 拷贝无法异步重叠 |
persistent_workers=False |
每个 epoch 重建 worker，小数据集上开销巨大 ⭐ |
⭐ 让 H2D 拷贝也重叠起来
# 🧩 骨架：`batch` 来自你自己的代码，这一段只看写法
# ❌ 同步拷贝，GPU 干等
import torch
batch = batch.cuda()
# ✅ 异步 + 预取下一个 batch
batch = batch.cuda(non_blocking=True) # ⭐ 需要 pin_memory=True
# ⭐ 更进一步：手动双缓冲预取
class Prefetcher:
def __init__(self, loader):
self.it, self.stream = iter(loader), torch.cuda.Stream()
self._preload()
def _preload(self):
try: self.next = next(self.it)
except StopIteration: self.next = None; return
with torch.cuda.stream(self.stream):
self.next = {k: v.cuda(non_blocking=True) for k, v in self.next.items()}
def __next__(self):
torch.cuda.current_stream().wait_stream(self.stream)
batch = self.next
if batch is None: raise StopIteration
self._preload() # ⭐ 立刻开始加载下一个
return batch
📦 三、存储格式：小文件是灾难
关键信息
- ❌ 一百万个小文件（每个 4KB）
- → 每次读都是一次文件系统操作
- → 在网络文件系统上，元数据操作比读数据本身还慢 💀
- ✅ 打包成大文件：
- WebDataset（tar 包） ⭐ 流式、支持分片、云存储友好
- Parquet / Arrow ⭐ 列式，适合表格和文本
- TFRecord / MDS 顺序读友好
- 内存映射（np.memmap） ⭐ 预处理好的 token 序列首选
⭐ 对 LLM 预训练：最好的做法是提前把文本 tokenize 成一个巨大的 uint16/uint32 数组， 用
np.memmap读。 训练时只是内存映射 + 切片，几乎零开销。tokens = np.memmap("train.bin", dtype=np.uint16, mode="r") # 取一个样本就是一次切片，快到可以忽略 x = torch.from_numpy(tokens[i:i+seq_len].astype(np.int64))
为什么小文件在分布式训练里特别糟：
关键信息
🔀 四、分布式下的数据划分
因果链
⭐ 流式数据集的洗牌难题
流程图
✂️ 五、序列打包：一个常被忽略的大优化
信息关系
因果链
⚠️ 打包必须正确隔离： 如果不隔离，样本 B 会 attend 到样本 A 的内容 —— 模型会学到错误的关联。 FlashAttention 支持
cu_seqlens（变长序列）， 这是正确且高效的做法（第 8 章提过：别用 padding mask）。
💽 六、存储层次
信息关系
💡 一个实用技巧：训练前先把下一个 epoch 的数据分片预取到本地盘。 用后台线程做，和训练完全重叠 —— 等于把网络存储的延迟藏起来了。
🔗 和站内其他章的关系
| 相关的地方 | 这里的位置 |
|---|---|
| 第 4 章 dataloader 实验 | 本章是它的展开 ⭐ |
| 第 10 章 set_epoch | 分布式采样 |
| 第 21 章 checkpoint 存数据位置 | 恢复时的正确性 |
| 第 8 章 cu_seqlens | 序列打包的正确实现 |
| 《模型上线之后》04 训练推理一致性 七种不一致 | 本章只保证喂得快，不保证喂对；口径、时间窗口、缺失值处理在那边 ⭐ |
| 《推荐算法》07 特征工程 特征穿越 | 打包和洗牌都救不了你 —— 样本里混进了未来信息，管线越快错得越快 ⚠️ |
| 《Kaggle竞赛方法论》10 数据增强与输入管线 大规模图像读取方案 | 图像侧的同一个问题：小文件 → LMDB / TFRecord / 预处理缓存，和本章的打包同一个思路 |
✅ 检查点
- 怎么在五分钟内确诊数据管线瓶颈？判据是什么？
- profiler 里数据没跟上的信号是什么？
- DataLoader 的四个常见错误配置？
pin_memory和non_blocking的关系是什么？- 为什么小文件在分布式训练里特别糟？
- LLM 预训练数据的最佳存储方式是什么？
- 分布式数据划分必须处理对的三个问题？
- 流式数据集怎么洗牌？缓冲区太小会怎样？
- 序列打包能带来多大收益？为什么必须正确隔离？
👀 答案
- ①把 dataloader 换成固定随机张量看 MFU 是否上升 ②只跑 dataloader 不训练测吞吐。判据：数据管线吞吐必须 ≥ 训练吞吐的 1.2 倍。
- GPU 时间轴上有规律的空隙，且空隙出现在每个 step 的开头。
num_workers=0（主进程加载必然阻塞）、num_workers太大（内存爆+CPU 争抢）、pin_memory=False（H2D 无法异步）、persistent_workers=False（每个 epoch 重建 worker）。non_blocking=True需要pin_memory=True才生效——只有锁页内存才能做异步 DMA 拷贝。- 因为1000 个 worker 同时发起小文件请求会打爆元数据服务器，所有节点一起变慢。网络文件系统上元数据操作比读数据本身还慢。
- 提前 tokenize 成一个巨大的 uint16/uint32 数组，用
np.memmap读。训练时只是内存映射+切片，几乎零开销。 - ①不重不漏（DistributedSampler）②每轮重新洗牌（set_epoch）③各 rank 步数一致（drop_last=True，否则 AllReduce 挂死）。
- 缓冲区洗牌：维护 N 条样本的缓冲区，每次随机取一条再补一条。太小则随机性不够，相邻样本仍然相关。经验：≥10000 条，或先在文件层面 shuffle 分片顺序。
- 有效 token 占比从 40% 提到 95%+，等于训练速度翻倍。必须隔离是因为不隔离的话样本 B 会 attend 到样本 A，模型学到错误关联。正确做法是用 FlashAttention 的
cu_seqlens。
🛑 可以停在这里
⚡ 走神救援
⭐ 最容易被忽略的瓶颈，也最容易被误诊成「模型慢」。
⭐⭐ 五分钟确诊两步：把数据加载换成固定的随机张量，看利用率是不是上去了；再单独跑数据管线测吞吐。⭐ 判据：数据管线的吞吐必须明显高于训练吞吐；profiler 上的信号是每个 step 开头都有一段有规律的空隙。
加载器参数里最容易漏的两个：⭐ 异步拷贝只有配上锁页内存才真的生效；⭐ 不开持久化 worker，每个 epoch 都要重建一次。
存储格式：⚠️⭐ 小文件是灾难——上千个 worker 同时发小文件请求会打爆元数据服务；✅ 打包成整块。⭐ 语言模型预训练的最佳做法是提前切好词、存成一个巨大的整数数组用内存映射读——训练时就只是切片，几乎零开销。
分布式三件事：不重不漏、每轮重设种子、⭐ 各个 rank 的步数必须一致——⚠️ 否则同步操作会直接挂死。 流式数据要用足够大的缓冲区洗牌，⚠️ 缓冲区太小的话相邻样本仍然相关。
⭐⭐ 序列打包是最常被忽略的大优化：变长文本补齐到固定长度时，有效 token 可能只占四成——也就是说大半算力在算填充；打包之后接近全满，等于训练速度翻倍。⚠️ 但必须正确隔离边界，否则后一个样本会看到前一个样本的内容，模型学到错误关联。
⭐ 存储层次的实用技巧：后台预取下一轮的分片到本地盘，把网络存储的延迟藏起来。
下一节 👉 23-可观测性与成本.md