一次训练 step 不是只有 forward/backward。 当 GPU 还在算当前 batch 时,CPU worker 应该已经把下一个 batch 准备好。 如果下一个 batch 没准备好,主进程会卡在 next(dataloader), GPU 利用率掉下去,nsys 里会看到 kernel 之间出现长空洞。 所以优化数据 pipeline 的第一步不是改参数,而是把每个环节拆开计时。

三类典型瓶颈

PyTorch 支持 map-style 和 iterable-style 两类 dataset。 大规模训练的数据常来自 tar shard、远端对象存储、日志流或已经打包好的 token block, 这些场景更适合 IterableDataset。 但 IterableDataset 有一个关键坑: 多 worker 时 dataset 对象会被复制到每个 worker,如果不主动切分 worker 的数据范围, 很容易重复读同一份样本。

推荐起步配置

from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=global_micro_batch,
    num_workers=8,                 # 从 4/8 起步,看 CPU 和 data_time 再调
    prefetch_factor=2,             # 总预取约 2 * num_workers 个 batch
    pin_memory=True,              # CUDA 训练通常开启
    persistent_workers=True,       # 长训练建议开启
    drop_last=True,
)

for batch in loader:
    batch = move_to_cuda(batch, non_blocking=True)
    loss = train_step(batch)

/dev/shm 为什么会炸

小文件随机读是训练数据系统的大敌。 WebDataset 的核心做法是把样本打进 tar shards: 同一个样本的不同字段共享 basename,例如 000001.jpg、000001.json。 训练时顺序读取 tar,在样本级做 shuffle、decode、map、batch。 它本质上是 PyTorch IterableDataset 风格的数据管线, 特别适合对象存储、分布式训练和大规模 image/text/audio 数据。

WebDataset 训练管线示例

import webdataset as wds
from torch.utils.data import DataLoader

urls = "s3://bucket/laion/train-{000000..001023}.tar"

dataset = (
    wds.WebDataset(urls, shardshuffle=1000)
    .shuffle(10000)                 # sample-level shuffle buffer
    .decode("pil")
    .to_tuple("jpg", "json")
    .map(preprocess)
)

loader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,
    prefetch_factor=2,
    pin_memory=True,
    persistent_workers=True,
)

WebDataset 调参关注点

Mosaic Streaming 的目标是让大规模远端数据像本地 dataset 一样被训练任务消费。 StreamingDataset 是一个 PyTorch IterableDataset,核心参数是: remote 指远端数据目录, local 指本地缓存目录, predownload 控制每个 worker 提前下载多少样本, cache_limit 控制本地缓存上限。 它还强调 deterministic shuffle 和 mid-epoch resume,适合弹性训练和云存储数据集。

StreamingDataset 示例

from streaming import StreamingDataset
from torch.utils.data import DataLoader

dataset = StreamingDataset(
    remote="s3://bucket/mds/train",
    local="/local_nvme/cache/train",
    split="train",
    shuffle=True,
    batch_size=64,             # 与 DataLoader per-device batch 对齐
    predownload=1024,           # 每 worker 提前缓存的样本目标
    cache_limit="500gb",
)

loader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,
)

Mosaic Streaming 字段表

当你自己写 IterableDataset 时, 最关键的工程问题是:在多机、多 rank、多 worker 下,每个 reader 应该读不同的数据切片。 PyTorch 文档明确提醒:多进程加载时同一个 dataset 对象会在 worker 中复制, 因此需要在 __iter__() 里用 get_worker_info() 或在 worker_init_fn 里改写每个 worker 的范围,避免重复数据。

import itertools
import os
import torch
import torch.distributed as dist
from torch.utils.data import IterableDataset, get_worker_info


class ShardedJsonlDataset(IterableDataset):
    def __init__(self, files, seed=1234):
        self.files = list(files)
        self.seed = seed

    def _rank_info(self):
        if dist.is_available() and dist.is_initialized():
            return dist.get_rank(), dist.get_world_size()
        return 0, 1

    def __iter__(self):
        rank, world = self._rank_info()
        worker = get_worker_info()
        worker_id = worker.id if worker else 0
        num_workers = worker.num_workers if worker else 1

        # 全局 reader id = rank 内 worker id 组合,确保每个 reader 读不同文件
        reader_id = rank * num_workers + worker_id
        num_readers = world * num_workers

        files = self.files[reader_id::num_readers]
        for path in files:
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    sample = parse_jsonl_line(line)
                    yield tokenize_and_pack(sample)

自定义 IterableDataset 检查表

今天的交付物是一张调参表:不同 num_workers、prefetch_factor、 pin_memory 组合下的 data_time、h2d_time、 GPU util 和内存/共享内存占用。 目标不是把 worker 一路加大,而是找到吞吐不再增长、资源开始爆的拐点。

# dataloader_probe.py
import argparse
import time
import torch
from torch.utils.data import DataLoader


def move_to_cuda(batch):
    if torch.is_tensor(batch):
        return batch.cuda(non_blocking=True)
    if isinstance(batch, dict):
        return {k: move_to_cuda(v) for k, v in batch.items()}
    if isinstance(batch, (list, tuple)):
        return type(batch)(move_to_cuda(v) for v in batch)
    return batch


def probe(dataset, batch_size, workers, prefetch, pin_memory, steps):
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=workers,
        prefetch_factor=prefetch if workers > 0 else None,
        pin_memory=pin_memory,
        persistent_workers=workers > 0,
        drop_last=True,
    )

    it = iter(loader)
    data_times, h2d_times = [], []
    for _ in range(steps):
        t0 = time.perf_counter()
        batch = next(it)
        t1 = time.perf_counter()
        batch = move_to_cuda(batch)
        torch.cuda.synchronize()
        t2 = time.perf_counter()
        data_times.append(t1 - t0)
        h2d_times.append(t2 - t1)

    print({
        "workers": workers,
        "prefetch": prefetch,
        "pin_memory": pin_memory,
        "data_ms": 1000 * sum(data_times) / len(data_times),
        "h2d_ms": 1000 * sum(h2d_times) / len(h2d_times),
    })

实验矩阵

系统观测命令

# GPU 是否空等
nvidia-smi dmon -s pucm

# 磁盘 await/util 是否高
iostat -xz 1

# worker CPU 是否打满
pidstat -ru -p ALL 1

# shared memory 是否不足
df -h /dev/shm
ipcs -m

# 容器启动时常见修复
docker run --shm-size=64g ...

记录模板

1. data_time 和 GPU util 之间是什么关系?
2. Map-style Dataset 和 IterableDataset 在分布式训练中的核心差异是什么?
3. 为什么多 worker IterableDataset 必须显式按 worker/rank 切分?
4. prefetch_factor 的总预取量如何计算?它为什么会增加内存压力?
5. pin_memory=True 配合 non_blocking=True 解决的是哪段瓶颈?
6. WebDataset 的 shard shuffle 和 sample shuffle 分别解决什么问题?
7. Mosaic Streaming 的 remote、local、predownload、cache_limit 分别控制什么?