22 · 数据管线与存储

# ⭐ 实验一：拿掉数据管线
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 一定会有等待

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（分布式下尤其重要）
)

# ❌ 同步拷贝，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 序列首选

   1000 个 worker 同时向存储发起小文件请求
   → 元数据服务器成为瓶颈
   → 存储系统被打爆，所有节点一起变慢 ⭐

   ⭐ 三个必须处理对的问题：

   ① 【不重不漏】：每个 rank 拿到不同的数据
      → DistributedSampler（第 10 章）

   ② 【每轮重新洗牌】：sampler.set_epoch(epoch)
      → 漏了等于没 shuffle ⭐

   ③ 【各 rank 步数一致】：
      → drop_last=True，或补齐
      → 否则某个 rank 提前结束 → AllReduce 挂死（第 21 章）💀

   问题：数据太大无法全部载入，只能流式读
        但流式读的顺序是固定的 → 没有随机性

   ✅ 解法：【缓冲区洗牌】（shuffle buffer）
     维护一个 N 条样本的缓冲区，每次随机取一条，再补一条进来

   ⚠️ 缓冲区太小 → 随机性不够（相邻样本仍然相关）
   ⭐ 经验：缓冲区 ≥ 10000 条，或者【先在文件层面 shuffle 分片顺序】

   ⚠️ 变长文本 + padding 的浪费：

   样本长度: [120, 2048, 340, 89, 1500]
   padding 到 2048 → 有效 token 只占 40%
   → 【60% 的算力在算 padding】💀

   ✅ 序列打包（packing）：
     把多条短样本【拼成】一条满长度的序列

   [样本A][样本B][样本C][样本D]  ← 拼满 2048
      ↑ 用 attention mask 或 position_ids 隔开，防止跨样本注意 ⭐

   → 有效 token 占比从 40% 提到 95%+
   → 等于【训练速度翻倍】⭐

   本地 NVMe      ~5 GB/s     ⭐ 最快，容量有限
   本地 SSD       ~1 GB/s
   网络存储(NFS)  ~1-10 GB/s  取决于网络和并发
   对象存储(S3)   ~几百 MB/s  但可以【高并发】

   ⭐ 常见架构：
   S3（全量数据）→ 预取到本地 NVMe（当前 epoch 用的分片）→ 训练