把 TB 级数据喂给大模型时最容易卡死在 IO 瓶颈上

很多人在跑模型的时候，显卡利用率低得可怜，盯着监控一看，其实是 CPU 在死磕数据加载，这就是典型的 IO 瓶颈。如果你的数据读取速度跟不上 GPU 的计算速度，买再贵的 H100 也是在浪费电。要解决这个问题，不能只靠简单的 read_csv，得在数据管线（Pipeline）上下功夫。

针对大规模数据集的实操，我建议从这几个维度去优化加载流程：

• 二进制格式替代文本： 别再用 CSV 或 JSON 了，直接转成 Parquet 或 TFRecord。这类列式存储或二进制格式支持压缩且读取极快，能直接减少磁盘 IO 次数。
• 异步预取（Prefetching）： 利用多线程或多进程在 GPU 计算当前 Batch 时，提前把下一个 Batch 加载到内存里。
• 内存映射（Memory Mapping）： 对于内存装不下的超大文件，用 mmap 这种方式让 OS 管理页缓存，避免一次性把文件全部读入内存导致 OOM。

如果你在使用 PyTorch，建议直接用 DataLoader 配合 num_workers 参数，但记得这个值不能盲目设高，否则会导致内存碎片化。一个比较稳妥的配置参考：

from torch.utils.data import DataLoader

# 根据 CPU 核心数设置 num_workers，pin_memory 必须开启以加速数据从内存拷贝到显存
train_loader = DataLoader(
    dataset=train_dataset, 
    batch_size=64, 
    shuffle=True, 
    num_workers=8, 
    pin_memory=True, 
    prefetch_factor=2
)

对于更极端的场景，比如处理数亿行的数据集，可以考虑用 Apache Arrow 来做内存管理，它能实现零拷贝（Zero-copy）读取，直接在内存中操作数据而不经过序列化和反序列化。

从零搭建一个高性能的数据加载流，核心就在于“消除等待”。先把数据格式洗成二进制，再把预取机制跑起来，通常能让训练速度提升 30% 以上。