第11章 DataLoader 与多进程数据流水线

"Our data model looks like this (queues are indicated with curly brackets):"

—— torch/utils/data/dataloader.py:814，NOTE [ Data Loader Multiprocessing Shutdown Logic ]（该 NOTE 起于 torch/utils/data/dataloader.py:810，v2.11.0）。紧随其后的那张 ASCII 图画的是：主进程 → {index_queue} → worker 进程 → {worker_result_queue} → 主进程的 pin_memory 线程 → {data_queue} → 数据输出

11.1 一行 for 循环背后的系统

每个 PyTorch 用户都写过：

loader = DataLoader(dataset, batch_size=32, num_workers=4, pin_memory=True)

for batch in loader:
    out = model(batch.cuda(non_blocking=True))
    loss = compute_loss(out)
    loss.backward()
    optimizer.step()

for batch in loader 这一行，背后是 4 个 worker 子进程 + 1 个 pin_memory 线程 + 5 条 multiprocessing queue（每个 worker 一条 index queue，外加 1 条共享的结果队列）+ 1 条进程内 queue.Queue 协同工作的流水线。如果 dataset 是从磁盘读图（IO bound）或者要做 augment（CPU bound），这套机器把 数据准备并行化，让 GPU 不用等。

本章拆这套流水线。源码主要在 torch/utils/data/dataloader.py（1709 行）和 torch/utils/data/_utils/。

11.2 4 件套数据抽象

PyTorch 数据加载的核心抽象：

graph LR
    Sam[Sampler<br/>决定读哪些索引]
    Ds["Dataset<br/>__getitem__(idx) → 一条样本"]
    Co["collate_fn<br/>list[sample] → batch"]
    Dl[DataLoader<br/>编排所有]
    Sam --> Dl
    Ds --> Dl
    Co --> Dl
    style Dl fill:#fef3c7,stroke:#f59e0b,stroke-width:2px

逐个介绍：

11.2.1 Dataset

class Dataset:
    def __getitem__(self, index): ...
    def __len__(self): ...

Map-style Dataset：实现 __getitem__ 和 __len__。典型用例：图片数据集（每个 idx 对应一个文件）。

class ImageDataset(Dataset):
    def __init__(self, root, transform):
        self.files = sorted(os.path.join(root, f) for f in os.listdir(root))
        self.transform = transform

    def __getitem__(self, idx):
        img = Image.open(self.files[idx])
        return self.transform(img)

    def __len__(self):
        return len(self.files)

Iterable-style Dataset（IterableDataset）：实现 __iter__。用于流式数据（如 Kafka stream、巨量 webdataset）。它的特殊性是不要求 __len__（可以选择实现），长度要么不知道、要么靠 worker 自己处理。

11.2.2 Sampler

Sampler 决定"按什么顺序遍历 dataset"。最常用的几种：

• SequentialSampler：按顺序 0, 1, 2, ..., N-1
• RandomSampler：每个 epoch 随机打乱
• WeightedRandomSampler：按权重采样（用于不平衡数据集）
• BatchSampler：把 sample 索引打包成 batch
• DistributedSampler：多卡训练时按 rank 分片

DataLoader(shuffle=True) 实际就是用 RandomSampler 包装。DistributedSampler 是 DDP 训练的关键 —— 它保证 rank 0 拿到 [0, N, 2N, ...]、rank 1 拿到 [1, N+1, 2N+1, ...]，不重复不遗漏。第 17 章 DDP 章会详细讲。

11.2.3 collate_fn

collate_fn 把 list of samples 合成一个 batch。默认 default_collate 处理常见类型：

• 张量 list → torch.stack
• numpy list → torch.as_tensor + stack
• list of tuple → 按位置转置后逐项 collate，返回的是 list（不是 tuple；namedtuple 才保留原类型）
• list of dict → dict of stacked tensors
• int / float → torch.tensor（float 为 float64）
• str / bytes → list（不堆叠，作为 batch metadata）

但变长数据（如 NLP 句子）要写自定义 collate_fn 做 padding：

def my_collate(batch):
    # batch: list of (input_ids, label)
    inputs = [b[0] for b in batch]
    labels = torch.tensor([b[1] for b in batch])
    inputs_padded = torch.nn.utils.rnn.pad_sequence(inputs, batch_first=True)
    return inputs_padded, labels

collate_fn 的执行位置很关键 —— num_workers > 0 时它在 worker 进程里跑（不在主进程）。这意味着 padding / 转 tensor 等 CPU 工作能并行，主进程只接收已经合好的 batch。

11.2.4 DataLoader

DataLoader 是这 4 件套的"总编排器"：

class DataLoader(Generic[_T_co]):
    def __init__(
        self,
        dataset,
        batch_size=1,
        shuffle=None,
        sampler=None,
        num_workers=0,
        collate_fn=None,
        pin_memory=False,
        ...,
        *,                          # 以下为仅限关键字参数
        prefetch_factor=None,
        persistent_workers=False,
        ...
    ):
        ...

源码在 torch/utils/data/dataloader.py:142。它根据 num_workers 选择两套实现：_SingleProcessDataLoaderIter（同步）或 _MultiProcessingDataLoaderIter（异步）。

11.3 单进程模式：最简实现

num_workers=0 时 DataLoader 在主进程跑：

# 简化版
class _SingleProcessDataLoaderIter:
    def __next__(self):
        index = next(self._sampler_iter)        # 取下一个 batch 索引
        data = self._dataset_fetcher.fetch(index)  # 直接调 dataset.__getitem__ + collate
        if self._pin_memory:
            data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)
        return data

——拿索引 → 读数据 → 合 batch → pin → return。主进程顺序跑，下一个 batch 等当前 batch 训练完才开始读。

适用场景：

• Dataset 是内存里的张量，读取是 O(1)
• 调试时想避免多进程导致 stack trace 丢失
• 单卡 + 小数据集 + GPU 比 CPU 快得多

但绝大多数训练场景下 num_workers=0 会让 GPU 大部分时间空等。这就是多进程模式的存在意义。

11.4 多进程模式：核心架构

num_workers > 0 时 DataLoader 启动一组 worker 子进程，主进程通过 queue 与它们通信：

graph TB
    Main[主进程<br/>训练循环]
    Sam[Sampler]
    PMT[pin_memory thread<br/>主进程内]
    subgraph Workers["worker 子进程 (4 个)"]
        W0[Worker 0<br/>_worker_loop]
        W1[Worker 1<br/>_worker_loop]
        W2[Worker 2<br/>_worker_loop]
        W3[Worker 3<br/>_worker_loop]
    end
    IQ0["_index_queues[0]"]
    IQ1["_index_queues[1]"]
    IQ2["_index_queues[2]"]
    IQ3["_index_queues[3]"]
    DQ["_worker_result_queue<br/>(所有 worker 共享)"]
    OQ["_data_queue<br/>(queue.Queue, 进程内)"]
    Main -->|"取下一个 batch 索引"| Sam
    Sam -->|"round-robin 派给某个 worker"| IQ0
    Sam --> IQ1
    Sam --> IQ2
    Sam --> IQ3
    IQ0 --> W0
    IQ1 --> W1
    IQ2 --> W2
    IQ3 --> W3
    W0 -->|读取 + collate| DQ
    W1 --> DQ
    W2 --> DQ
    W3 --> DQ
    DQ --> PMT
    PMT -->|pin 后| OQ
    OQ -->|"next() 返回"| Main
    style Main fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
    style PMT fill:#dbeafe,stroke:#3b82f6
    style Workers fill:#dcfce7,stroke:#22c55e

工作流：

1. 启动时主进程创建 N 个 worker，每个跑 _worker_loop，并且为每个 worker 单独建一条 index_queue（torch/utils/data/dataloader.py:1155 的 self._index_queues = []）
2. 主进程预先往这些队列里 round-robin 派满 prefetch_factor × num_workers 个 batch 索引
3. worker 从自己那条队列 pop 索引、调 dataset.__getitem__ 读取、collate、把结果 put 到共享的 _worker_result_queue
4. 开了 pin_memory 时，主进程内的 pin 线程从 _worker_result_queue 取出、pin 完再放进 _data_queue（一个普通的 queue.Queue）；没开 pin_memory 时 _data_queue 直接就是 _worker_result_queue（torch/utils/data/dataloader.py:1228）
5. 每消费一个 batch，主进程就再派一个新索引出去，维持 pipeline

注意方向不对称：主进程→worker 是 N 条队列（这样才能把任务定向派给指定 worker，也才能给指定 worker 发关闭信号），worker→主进程只有 1 条共享队列（谁先做完谁先放）。

整套机制让 worker 与训练真正并发 —— GPU 跑 step N 时 worker 在准备 step N+1、N+2 的数据。

11.5 worker 主循环

打开 torch/utils/data/_utils/worker.py:228 的 _worker_loop（精简版）：

def _worker_loop(dataset_kind, dataset, index_queue, data_queue, ...):
    # 1. 在 C 侧装 SIGBUS/SIGSEGV 等 handler，让致命信号的死因能被主进程拿到
    signal_handling._set_worker_signal_handlers()

    # 2. 单线程模式 (avoid CPU thrashing)
    torch.set_num_threads(1)

    # 3. 设置 RNG seed
    seed = base_seed + worker_id
    random.seed(seed)
    torch.manual_seed(seed)
    if HAS_NUMPY:
        np_seed = _generate_state(base_seed, worker_id)   # 不是简单相加，见 §11.5.2
        np.random.seed(np_seed)

    # 4. 暴露 worker info 给用户回调
    global _worker_info
    _worker_info = WorkerInfo(id=worker_id, num_workers=num_workers, seed=seed, dataset=dataset)

    # 5. 创建 fetcher (Dataset 包装器, 处理 collate)
    fetcher = _DatasetKind.create_fetcher(...)

    # 6. 主循环（watchdog = ManagerWatchdog()，见 §11.29）
    while watchdog.is_alive():
        try:
            r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
        except queue.Empty:
            continue    # 超时：回到 while，让 watchdog 再查一次父进程
        if r is None:
            break    # 收到关闭信号
        elif done_event.is_set():
            continue    # 正在关闭：不再干活，只等 None
        idx, index = r
        try:
            data = fetcher.fetch(index)        # 读取 + collate
        except Exception:
            data = ExceptionWrapper(where=f"in DataLoader worker process {worker_id}")
        data_queue.put((idx, data))            # 发回主进程

几个细节：

11.5.1 torch.set_num_threads(1) 防 CPU 抢占

每个 worker 进程默认会让 PyTorch C++ 后端开多线程跑算子（如 torch.einsum），但 num_workers=4 加上每个 worker 4 个 thread = 16 thread 抢 CPU 核 → 严重 thrashing。torch.set_num_threads(1) 让每个 worker 单线程，4 个 worker 加起来正好用 4 核。

如果你的 dataset transform 用了大量 PyTorch 算子（如 vision transform），这条限制正是防止 worker 之间互相抢核。反过来说：如果你把 num_workers 设得比物理核数还多，即使有这行限制，进程之间照样会开始抢核 —— 这就是 §11.37 那条"别把 num_workers 设过头"的由来。

11.5.2 RNG seed：reproducible 但不重复

worker 0 的 seed = base_seed + 0、worker 1 = base_seed + 1，etc.。这保证：

• 不同 worker 的随机性不同（augment 不会重复）
• 每次运行一致（用同一个 base_seed 重跑能复现）
• 跨 epoch 不重复：base_seed 由主进程在每次新建 iterator 时从自己的 RNG 里现取一个，所以每个 epoch 的 worker seed 都不同。开了 persistent_workers=True 时 iterator 与 worker 都复用，seed 只在 worker 启动时设一次，后续 epoch 沿用同一条 RNG 流继续往下取（实测 get_worker_info().seed 跨 epoch 不变，随机数不重复）

一个容易记错的细节：numpy 的 seed 不是 base_seed + worker_id。torch/utils/data/_utils/worker.py:262 走的是 np_seed = _generate_state(base_seed, worker_id) —— _generate_state（torch/utils/data/_utils/worker.py:176）是从 numpy.random.SeedSequence 移植过来的哈希混合算法，目的是避免 numpy 与 Python random 拿同一个 seed —— 两者都是 MT19937，同 seed 就会产生完全相同的状态序列（引入它的提交 3b977a0d28，PR #56797，说明的就是这个动机）。torch 与 python 的 random 用的仍是简单相加。至于用户自定义的全局 RNG，PyTorch 管不到。生产代码里建议优先用 torch.rand 与 dataset 内部的 torch.Generator 实例。

11.5.3 ExceptionWrapper：跨进程异常传递

worker 抛异常时不能直接 raise（多进程间异常无法跨进程传播），用 ExceptionWrapper 把异常 + traceback 序列化成普通对象，发回主进程后由主进程 reraise。这种"序列化异常"模式让 worker 的报错在主进程能看到完整 stack trace。

11.6 pin_memory：异步 H2D 的关键

pin_memory=True 启动一个主进程内的线程（不是子进程！），把 worker 返回的 CPU tensor 搬到 pinned (page-locked) memory：

# torch/utils/data/_utils/pin_memory.py:18（简化：省略了 queue.Empty / queue.Full 的超时重试与 pin 失败时的 ExceptionWrapper）
def _pin_memory_loop(in_queue, out_queue, device_id, done_event, device):
    torch.set_num_threads(1)
    torch.accelerator.set_device_index(device_id)

    while not done_event.is_set():
        r = in_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
        idx, data = r
        if not isinstance(data, ExceptionWrapper):
            data = pin_memory(data, device)
        out_queue.put((idx, data))

pin_memory(data, device) 递归遍历 data（dict / list / tensor）调 tensor.pin_memory()。pinned memory 的好处：

• H2D 拷贝可以 async（CPU 不用等 GPU）
• CUDA driver 直接 DMA，不走 page fault
• 配合 tensor.cuda(non_blocking=True) 拿到完整异步性能

为什么用线程而不是子进程？因为 pin_memory() 调用要访问 CUDA context —— 子进程要自己 init CUDA 太重。线程在主进程里直接用主进程的 CUDA context，简单高效。

这两个开关必须配对使用：只开 pin_memory=True 但 .cuda() 时没写 non_blocking=True，拷贝仍然是同步的，等于白搭一个线程；反过来只写 non_blocking=True 而内存没 pin，CUDA 会静默退化成同步拷贝。配对之后 H2D 才有机会与训练计算 overlap —— 实际能藏掉多少，取决于你的 batch 有多大、GPU 一步算多久，要自己 profile。

11.7 persistent_workers：跨 epoch 复用 worker

默认每个 epoch 结束 DataLoader 关闭所有 worker，下个 epoch 重启。重启代价：

• 重新创建 worker 进程（spawn / forkserver 下还要重新 import 依赖、把 dataset 整个 pickle 过去）
• 每个 worker 重新设 seed、重跑 worker_init_fn、重建 fetcher
• 这些固定开销每个 epoch 都付一遍

如果 dataset 对象本身很大（如在内存里存着 ImageNet 1.28M 图片清单），spawn / forkserver 下每次重启都要把它重新 pickle 一遍，这个重启成本就尤其高。

persistent_workers=True 让 worker 跨 epoch 不退出：新 epoch 开始时主进程重建自己的 sampler 迭代器，再给每个 worker 发一个 _ResumeIteration 让它重建 fetcher 即可（torch/utils/data/dataloader.py 的 _MultiProcessingDataLoaderIter._reset）：

loader = DataLoader(dataset, num_workers=4, persistent_workers=True)
for epoch in range(100):
    for batch in loader:    # 第 2 个 epoch 起 worker 不重启
        ...

代价是 worker 持续占着内存，dataset 里的 cached 数据（如预读到内存的 ImageNet）一直占用 RAM。生产代码里几乎所有训练循环都该开 persistent_workers，除非 worker 占内存极大。

11.8 prefetch_factor：预取深度

prefetch_factor 在 num_workers > 0 时默认 2（num_workers=0 时默认 None，显式传值会直接 ValueError），意思是每个 worker 预取 2 个 batch 在 queue 里。总 prefetch = num_workers * prefetch_factor。

为什么要预取？因为 worker 跑 dataset.__getitem__ 有一定延迟（如读图、解码、augment）。如果训练消费数据的速度比 worker 快，就会有"GPU 等数据"的 stall。预取让 queue 里始终有几个 batch 待取，吸收 worker 间的速率波动。

调优建议：

• 默认 2 适合大多数场景
• IO 抖动大（如从 S3 读）→ 提升到 4-8
• 内存紧张 → 降到 1，但要确保 worker 数足够多

11.9 BatchSampler：用 batch 索引省 sampler 开销

默认 DataLoader 内部把 Sampler + batch_size 组合成 BatchSampler。BatchSampler 一次返回 batch_size 个索引：

# 默认行为
sampler = RandomSampler(dataset)            # 给单个索引
batch_sampler = BatchSampler(sampler, batch_size, drop_last)  # 给 batch 索引

for indices in batch_sampler:    # indices 是 list of 32 个索引
    samples = [dataset[i] for i in indices]
    batch = collate_fn(samples)

用户可以传自定义 batch_sampler 替代默认。常见用法：动态 batch size（如按序列长度排序后凑满 token 预算的 batch）：

class DynamicBatchSampler:
    """按 token 数构造 batch, 长序列少配, 短序列多配, 让每个 batch 总 token 数接近"""
    def __iter__(self):
        ...

这种"按 token 而非按 sample 数"在 NLP 训练里（特别是变长序列）能显著提升 GPU 利用率。

11.10 fork vs spawn：跨平台陷阱

multiprocessing 启动子进程有两种模式：

• fork（Python 3.13 及以前在 Linux 上的默认；Python 3.14 起 macOS 以外的 POSIX 平台默认改为 forkserver，见 §11.14）：子进程继承父进程内存。Dataset 对象不需要 picklable，速度快
• spawn（Windows 一直是默认；macOS 自 Python 3.8 起默认）：子进程从零启动 + pickle 传递所有参数。Dataset 必须 picklable（不能含 lambda / 闭包 / unpicklable 对象）

后果：在 Mac 调通的 DataLoader 在 Linux 上 hang，或者反过来 —— 都是真实碰到的工程坑。

# 强制指定启动方式
import torch.multiprocessing as mp
mp.set_start_method('spawn', force=True)

关于 CUDA 有一条经常被过度概括的说法，需要说准确：报错 "Cannot re-initialize CUDA in forked subprocess" 的触发条件是 主进程已经初始化过 CUDA，fork 出来的子进程又自己去初始化 CUDA（torch/csrc/utils/device_lazy_init.cpp 用 pthread_atfork 在子进程里打上 bad-fork 标记），而不是"用了 CUDA 就不能 fork"。所以：

• 主进程已经初始化过 CUDA、worker 只做纯 CPU 的读盘与 transform —— Linux 默认的 fork 完全能跑，这也是绝大多数训练脚本的实际情形
• worker 里出现任何 CUDA 调用（在 __getitem__ 里 .cuda()、用 GPU 解码库、返回 CUDA 张量）—— 才会撞上这条错误，此时要么把 CUDA 操作挪回主进程，要么改用 spawn / forkserver

DataLoader 并不会替你自动切换启动方式，它用的是 multiprocessing_context（不传就是平台默认）。理解这条差异，你看到 PicklingError 时也不会困惑 —— 那是 spawn / forkserver 在序列化 Dataset 时失败了，fork 模式下则不会有这个问题。

11.11 worker shutdown 的精细同步

那段三百来行的注释 "NOTE [ Data Loader Multiprocessing Shutdown Logic ]" 在 torch/utils/data/dataloader.py:810（_MultiProcessingDataLoaderIter 类体开头），torch/utils/data/_utils/worker.py:244、torch/utils/data/_utils/pin_memory.py:47、torch/utils/data/_utils/signal_handling.py:29 三处都指回它。简化版流程：

1. 开了 pin_memory 时先停 pin 线程：set pin_memory_thread_done_event，往 worker_result_queue 塞一个 (None, None) 把它唤醒，再 join
2. 主进程 set workers_done_event
3. 往每个 worker 的 index_queue 发 None（关闭信号）
4. worker 看到 done_event 已 set 就不再处理任务、只等 None；收到 None 才 break 主循环 → 进程退出（index_queue.get() 每 MP_STATUS_CHECK_INTERVAL 超时一次，保证不会永久阻塞）
5. 主进程带超时 join 所有 worker，对各 index_queue 调 cancel_join_thread() + close()
6. 最后从 C 侧摘掉这些 worker 的 pid（SIGCHLD 检查不再管它们），仍未退出的 worker 直接 terminate()

为什么这么复杂？因为有几种异常退出场景要处理：

• 主进程被 Ctrl+C 杀掉（worker 必须在父进程死亡时也死，否则成为继续占资源的孤儿进程）
• worker 自己崩了（main 必须感知到、报错而不是 hang）
• IterableDataset 提前 StopIteration（必须告诉 main "我没数据了，别再发任务给我"）

signal_handling._set_worker_signal_handlers()（worker.py:253）在 C 侧注册 SIGBUS / SIGSEGV（以及 SIGFPE、SIGTERM）handler —— 源码注释解释了为什么必须在 C 侧做：Python 的 signal handler 要等 C 层返回后才执行，而致命信号很可能在那之前就再次发生。ManagerWatchdog（实例化在 worker.py:314）则负责感知父进程死亡，见 §11.29。

如果你写训练代码遇到 "ctrl-c 后 python 进程 hang 不退"，几乎都是 DataLoader shutdown 流程出了问题。persistent_workers=True 引入的复杂度让这套同步逻辑更微妙。

11.12 IterableDataset 的 worker 分片

Dataset（map-style）有 __len__，DataLoader 能给每个 worker 分配不同索引。但 IterableDataset 没有按索引访问，每个 worker 默认会重复跑同一份数据流 —— 4 个 worker = 4 份重复 batch！

正确做法：dataset 内部按 worker_id 分片：

class StreamDataset(IterableDataset):
    def __init__(self, files):
        self.files = files

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            files = self.files                    # 单进程, 全部
        else:
            n = worker_info.num_workers
            i = worker_info.id
            files = self.files[i::n]              # worker i 取第 i, i+n, i+2n, ... 份
        for f in files:
            yield from read_file(f)

torch.utils.data.get_worker_info() 在主进程返回 None，在 worker 里返回 WorkerInfo(id, num_workers, seed, dataset)。这是 IterableDataset 多 worker 训练的关键 API。

第 17 章 DDP 章会再叠一层 rank 间分片（map-style 用 DistributedSampler；IterableDataset 用不了 sampler，要在 __iter__ 里同时按 rank 与 worker 切，见 §11.28），两层分片合起来让 N rank × M worker = N×M 份数据并行加载，互不重复。

11.13 DistributedSampler：多卡训练的数据切片

第 17 章 DDP 章会把这玩意儿用到底，这里先讲它的内部机制。torch/utils/data/distributed.py:17 的 DistributedSampler（下面按语义精简）：

class DistributedSampler(Sampler):
    def __init__(self, dataset, num_replicas=None, rank=None,
                 shuffle=True, seed=0, drop_last=False):
        ...

    def __iter__(self):
        # 1. 生成全局索引序列
        if self.shuffle:
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)    # 每个 epoch 不同 seed
            indices = torch.randperm(len(self.dataset), generator=g).tolist()
        else:
            indices = list(range(len(self.dataset)))

        # 2. padding 让总长度能被 num_replicas 整除
        if not self.drop_last:
            padding_size = self.total_size - len(indices)
            indices += indices[:padding_size]
        else:
            indices = indices[:self.total_size]

        # 3. 切片：rank i 拿 [i, i+N, i+2N, ...]
        indices = indices[self.rank:self.total_size:self.num_replicas]
        return iter(indices)

几个工程细节：

• 每个 epoch 必须调 sampler.set_epoch(epoch)：否则每个 epoch shuffle 顺序一样、训练 loss 曲线会异常。这是 DDP 训练新手常见 bug
• drop_last vs padding：drop_last 把不够整除的 sample 扔掉、padding 把它们重复一份补齐。padding 让 epoch 内所有 rank batch 数一致（DDP 要求），但有重复样本
• seed = 用户 seed + epoch：所有 rank 用同一个 base seed → 所有 rank 算出的全局 indices 一致 → 切片后 rank i 拿到第 i, i+N, ... 个样本，不重复不遗漏

这套机制让 N 个 rank 的训练近似等价于单 rank 训练 N×batch_size（前提：BN 用 SyncBN、学习率随总 batch 调整；padding 引入的少量重复样本除外）。是 DDP 训练正确性的基础。

11.14 forkserver：fork 与 spawn 之间的折中

§11.10 提到 fork（快但 CUDA 不安全）vs spawn（慢但安全）。还有第三种 forkserver：

mp.set_start_method('forkserver')

机制：启动一个轻量的 server 进程，每次要 fork 子进程时由 server 来 fork（而不是从主进程）。好处：

• server 是一个新启动的干净解释器（不是从主进程 fork 出来的）→ 子进程不继承主进程的 CUDA context
• 比 spawn 快（仍然是 fork，不是从零启动）
• 比 fork 安全（不继承主进程复杂状态）

代价：dataset 仍要 picklable（与 spawn 一样），不能含 lambda。

三者的成本结构（而不是具体毫秒数）是这样的：fork 只做一次 fork(2)，不重新 import 任何东西，最便宜；forkserver 也是 fork，但 fork 的源头是那个"干净的" server 进程，代价是要把参数 pickle 过去；spawn 每个子进程都要从零启动一个 Python 解释器并重新 import 你的整个依赖树（import torch 本身就不便宜），所以最贵，且启动开销随依赖规模增长。这就是为什么 persistent_workers=True 在 spawn 模式下的收益远大于 fork 模式。

选择建议：worker 不碰 CUDA 时，Linux 上用 fork 即可（Python 3.14+ 默认已是 forkserver，要 fork 得显式指定）；worker 必须碰 CUDA 时，优先 forkserver、回退 spawn。Windows 上没有 fork；macOS 上 fork 虽可用但不安全，默认就是 spawn。

11.15 DataPipes 与 DataLoader2：新一代实验

PyTorch 自 v1.8 起在主仓 torch/utils/data/datapipes/ 里加入 DataPipes（更丰富的 DataPipe 库放在 torchdata 仓库），思路是把数据加载用函数式管道表达：

from torchdata.datapipes.iter import FileLister, FileOpener, Mapper

dp = FileLister("/data/imagenet")
dp = dp.shuffle(buffer_size=10000)
dp = dp.sharding_filter()                  # 按 worker / rank 分片
dp = FileOpener(dp, mode="b")
dp = dp.map(decode_image)                   # 每条样本应用 decode_image
dp = dp.batch(batch_size=32)
dp = dp.collate()

每个 DataPipe 都是 lazy iterator。组合起来比 class MyDataset(IterableDataset) 更可读、复用度更高。

DataLoader2 是新版加载器（torchdata 中的版本；下面是它的用法），把 reading service（如 prefetching、distributed sharding）从 dataset 抽离：

loader = DataLoader2(
    datapipe=dp,
    reading_service=MultiProcessingReadingService(num_workers=4)
)

reading service 负责 worker 池、prefetching；datapipe 负责数据流水。职责分离让用户能更灵活组合（如换成 DistributedReadingService 一行切到 DDP 模式）。

实战现状要说清楚：这条路线并没有成为主流。到 v2.11，torch/utils/data/datapipes/ 仍在 PyTorch 主仓里（torch/utils/data/datapipes/datapipe.py 留着按名字发 _deprecation_warning 的钩子，但 torch/utils/data/datapipes/utils/common.py 里的 _iter_deprecated_functional_names / _map_deprecated_functional_names 两张表是空的）；DataLoader2 曾以实验形式放在主仓的 torch.utils.data.dataloader_experimental 模块里（v1.10–v1.13），v2.0 起从主仓移除，此后只在 torchdata 里，而 torchdata 后来把重心转向了 StatefulDataLoader。

结论很直接：生产代码用 DataLoader（v1）。DataPipes 值得了解的是它的思想 —— 把"读文件、分片、解码、组 batch"拆成一串可组合的 lazy iterator，这个思路在 webdataset 这类库里活得很好，只是没有以 DataPipes 的形式赢下 PyTorch 生态。

11.16 timeout：worker 健康检查

DataLoader(timeout=N) 让 next(iter(loader)) 等 N 秒拿不到数据就抛 RuntimeError，避免训练 hang 在数据加载上。机制：

• 主进程从 data_queue.get(timeout=N) 等待
• 超时后检查每个 worker 是否还活着（w.is_alive()，见 torch/utils/data/dataloader.py 的 _try_get_data）
• 如果有 worker 死了 → raise RuntimeError("DataLoader worker (pid(s) xxx) exited unexpectedly")（非 Windows 平台上 SIGCHLD handler 通常更早发现，死因会作为它的 direct cause 出现，见 §11.27）
• 如果都活着但仍超时 → raise RuntimeError("DataLoader timed out after {timeout} seconds")

默认 timeout=0 表示无限等待。生产代码建议设个保险值（如 60 秒），避免 NFS 卡死、worker silently 卡住等场景让训练 hang 几小时。

MP_STATUS_CHECK_INTERVAL 定义在 torch/utils/data/_utils/__init__.py:19，值是 5.0 秒。它被用作 worker 内部 index_queue.get(timeout=...)（torch/utils/data/_utils/worker.py:318）与 pin 线程 in_queue.get(timeout=...) 的超时 —— 每 5 秒醒一次检查"是否被叫退出"，让 shutdown 不至于卡在阻塞的 get 上。

11.17 跨进程 tensor 传输：共享内存绕开 pickle

worker 把 batch 发回主进程，不是 pickle 整个 tensor 再 send。torch/multiprocessing/reductions.py 注册了 tensor 的特殊序列化逻辑：

graph LR
    W[worker 进程<br/>collate 出 batch tensor]
    W --> SH[把 tensor 数据放到<br/>shared memory file]
    W --> META[把 metadata<br/>shape/dtype/storage_handle<br/>送到 queue]
    META --> Q[multiprocessing.Queue]
    Q --> M[主进程拿到 metadata]
    M --> RB[根据 storage_handle 重建<br/>指向同一块 shm 的 tensor]
    SH -.同一块物理内存.-> RB
    style SH fill:#fef3c7,stroke:#f59e0b
    style RB fill:#dbeafe,stroke:#3b82f6

机制：

1. worker 里的 batch 张量直接落在共享内存：default_collate 在 worker 里把样本直接 stack 进 _new_shared 分配的共享存储（torch/utils/data/_utils/collate.py 的 collate_tensor_fn）；其余张量在被 pickle 进队列时由 torch/multiprocessing/reductions.py 挪进共享内存。句柄怎么传由 set_sharing_strategy 决定：Linux 默认 'file_descriptor'，另一种是 'file_system'
2. 传输：只把 storage 的 file handle / fd 通过 queue 送过去（几十字节）
3. 主进程：用 file handle 打开同一块 shm，重建 tensor —— 零拷贝

为什么不直接 pickle？一个 ImageNet batch（256×3×224×224 的 fp32 就约 154MB）若走 pickle，要整块序列化、经管道写过去、再反序列化；共享内存只传一个句柄，传输开销基本与 batch 大小无关。

代价：每次创建的 shm 段都要在数据消费后释放。torch.multiprocessing 用引用计数自动清理，但进程被 SIGKILL（比如 OOM killer）打死时来不及跑清理，可能在 /dev/shm 下留下 torch_* 残留（主要是 file_system 策略；Linux 默认的 file_descriptor 策略在创建 shm 文件后立即 unlink，只靠 fd 保活）—— 这些是普通文件，rm 掉即可，不需要重启。另外 /dev/shm 在容器里默认往往只有 64MB，跑 DataLoader 时经常要用 --shm-size 调大，否则会撞上 "Bus error"。生产监控值得加一条 /dev/shm 占用告警。

11.18 Fetcher：dataset 与 sampler 的中间层

§11.5 主循环里的 fetcher.fetch(index) 是什么？源码 torch/utils/data/_utils/fetch.py：

class _MapDatasetFetcher(_BaseDatasetFetcher):
    def fetch(self, possibly_batched_index):
        if self.auto_collation:
            if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__:
                data = self.dataset.__getitems__(possibly_batched_index)  # 批量读取钩子
            else:
                data = [self.dataset[idx] for idx in possibly_batched_index]
        else:
            data = self.dataset[possibly_batched_index]
        return self.collate_fn(data)


class _IterableDatasetFetcher(_BaseDatasetFetcher):
    def __init__(self, dataset, auto_collation, collate_fn, drop_last):
        super().__init__(dataset, auto_collation, collate_fn, drop_last)
        self.dataset_iter = iter(dataset)
        self.ended = False

    def fetch(self, possibly_batched_index):
        if self.ended:
            raise StopIteration
        # IterableDataset 没有 __getitem__, 只能 next()
        if self.auto_collation:
            data = []
            for _ in possibly_batched_index:
                try:
                    data.append(next(self.dataset_iter))
                except StopIteration:
                    self.ended = True
                    break
            if len(data) == 0 or (
                self.drop_last and len(data) < len(possibly_batched_index)
            ):
                raise StopIteration     # worker 捕获后发 _IterableDatasetStopIteration
        else:
            data = next(self.dataset_iter)
        return self.collate_fn(data)

两种 fetcher 反映 Map vs Iterable 的本质差异：

特性	_MapDatasetFetcher	_IterableDatasetFetcher
index 来源	主进程 sampler 推送	dataset 自己迭代
是否能随机访问	是（__getitem__）	否（仅 next）
BatchSampler 作用	把单个 idx 打成 batch	仅决定 batch 大小
StopIteration	不会发生	数据耗尽时触发

理解 Fetcher 让你看到 DataLoader 一套 worker 主循环代码能同时处理两类 dataset 的工程巧思 —— 通过抽象 fetch 接口屏蔽差异。

11.19 Sample 数据类型谱与 default_collate

default_collate 处理 batch 合并的核心逻辑（torch/utils/data/_utils/collate.py:340；v2.11 里它其实是薄薄一层，真正的分派表在同文件的 default_collate_fn_map 里，下面按语义展开成 if-else 便于阅读）：

def default_collate(batch):
    elem = batch[0]
    elem_type = type(elem)
    if isinstance(elem, torch.Tensor):
        return torch.stack(batch, 0)
    elif elem_type.__module__ == "numpy" and ...:
        return torch.as_tensor(np.stack(batch))
    elif isinstance(elem, float):
        return torch.tensor(batch, dtype=torch.float64)
    elif isinstance(elem, int):
        return torch.tensor(batch)
    elif isinstance(elem, str):
        return batch                           # 字符串保持原样
    elif isinstance(elem, collections.abc.Mapping):
        return {key: default_collate([d[key] for d in batch]) for key in elem}
    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):
        return elem_type(*(default_collate(samples) for samples in zip(*batch)))
    elif isinstance(elem, collections.abc.Sequence):
        return [default_collate(samples) for samples in zip(*batch)]
    raise TypeError(default_collate_err_msg_format.format(elem_type))

设计要点：

• 递归处理嵌套结构：{"image": tensor, "label": int, "meta": {"path": str}} → {"image": stacked_tensor, "label": int_tensor, "meta": {"path": [strs]}}
• 保持 namedtuple 类型：自定义 namedtuple 输入仍输出同类型，不退化成 tuple
• str 不堆叠：字符串作为 batch metadata，保留 list

如果你的 dataset 返回了 default_collate 不认识的类型（如自定义 dataclass），就会触发 TypeError。要么写自定义 collate_fn、要么把数据转成 dict / tuple / tensor。torch.utils.data.default_collate 本身可以被 monkey-patch 但不推荐 —— 影响其他代码的行为。

11.20 _try_put_index：prefetching 的具体实现

§11.8 讲 prefetch_factor。具体怎么实现？看 dataloader.py 的 _MultiProcessingDataLoaderIter._try_put_index：

# torch/utils/data/dataloader.py:1550
def _try_put_index(self) -> None:
    max_tasks = self._prefetch_factor * self._num_workers
    if self._tasks_outstanding >= max_tasks:
        raise AssertionError("Number of outstanding tasks exceeded maximum allowed tasks")

    try:
        index = self._next_index()             # sampler 给的下一个 batch idx
    except StopIteration:
        return
    for _ in range(self._num_workers):  # find the next active worker, if any
        worker_queue_idx = next(self._worker_queue_idx_cycle)
        if self._workers_status[worker_queue_idx]:
            if self._in_order:
                break
            elif self._workers_num_tasks[worker_queue_idx] < max_tasks // sum(self._workers_status):
                # when self._in_order is False, distribute work to a worker if it has capacity
                break
    else:
        return    # 没找到可派的 worker（都已耗尽，或 in_order=False 时都满额）, 不再发任务

    self._index_queues[worker_queue_idx].put((self._send_idx, index))
    self._task_info[self._send_idx] = (worker_queue_idx,)
    self._workers_num_tasks[worker_queue_idx] += 1
    self._tasks_outstanding += 1
    self._send_idx += 1

关键逻辑：

• _tasks_outstanding 跟踪有多少任务在 worker 那里 pending。上限 = prefetch_factor × num_workers
• 轮转分配：用 itertools.cycle 造的 _worker_queue_idx_cycle 把任务依次派给各 worker
• 跳过死掉的 worker：检查 _workers_status[i]；一圈都没找到活着的就直接返回
• _in_order 分支：DataLoader(in_order=False) 时不再严格轮转，而是看 _workers_num_tasks[i] 是否还有余量，让快的 worker 多干活（代价是 batch 返回顺序不再确定）
• 每次 next() 拿到 batch 后立刻 put 一个新 index：维持 outstanding 任务数稳定

启动时主进程会预先 put 满 prefetch_factor * num_workers 个任务。后续每消费一个 batch 就 put 一个新任务，让 queue 始终满 —— 这是 sustained throughput 的关键。

这个机制的极致：prefetch_factor=4, num_workers=8 时主进程预 put 32 个任务，worker 池始终有 32 个待执行 batch 在 pipeline 里。GPU 训练慢一点 worker 就能打满；GPU 训练快了 worker 就成瓶颈 —— outstanding 数是 GPU/CPU 速率匹配的关键参数。

11.21 Pin memory 的页锁定原理

§11.6 提到 pinned memory 让 H2D 可以 async + 跳 page fault。底层是怎么实现的？

正常 malloc 出的内存是可被换页（pageable）：操作系统可能把这块内存换到 swap 区。CUDA driver 做 H2D 拷贝时：

1. driver 不能直接 DMA 到可换页内存（DMA 控制器不知道虚拟地址，OS 也可能在拷贝时换页）
2. driver 必须先把数据拷进自己持有的一块锁页暂存缓冲区（staging buffer），再从那里 DMA 到显存
3. 这条"先暂存再 DMA"的路径对 host 是 同步阻塞 的，CPU 必须等

tensor.pin_memory() 经 CachingHostAllocator（aten/src/ATen/cuda/CachingHostAllocator.cpp）默认调 cudaHostAlloc（可配置为 cudaHostRegister）申请不会被换页的内存：

• driver 知道这块内存物理地址固定 → 可以 setup DMA 后直接返回
• CPU 立刻继续做别的事，DMA 在后台跑
• 这就是为什么"pin_memory + non_blocking=True"才能 async

代价：pinned memory 不能被换出，占多了会挤压系统其余部分可用的物理内存。开太多 pin_memory 会失败、或拖累系统其他进程。DataLoader 默认只对 batch tensor pin，不对整个 dataset pin。

11.22 worker 内存增长：fork CoW 失效问题

Linux fork 用 Copy-on-Write：fork 时父子进程共享物理内存页，只在写入时复制。理论上 worker 启动几乎不占新内存。

但有一个坑：CPython 的引用计数会写入对象的 ob_refcnt。worker 即使只是"读"一个 Python 对象（如遍历一个 list），也会修改它的 refcnt → 触发 CoW → 该内存页整个被复制。

实战后果：dataset 里 cache 了大量 Python 对象（如 1.28M 个文件路径 list），worker 启动后内存看起来正常，跑一段时间后 worker 内存爆炸 —— 全部从父进程 CoW 过来了。

解决方法：

1. 把 cache 数据用 numpy / tensor 存（refcnt 操作不会蔓延到数据本身）
2. 用 spawn 启动：避免 CoW 误伤（虽然启动慢）
3. gc.freeze()（Python 3.7+）：在 fork 前把现有对象移进永久代，GC 不再扫描它们，避免 GC 自己写对象头触发 CoW；但引用计数照样会变，只能缓解、治不了本

PyTorch 文档（docs/source/data.md 的 Multi-process data loading 一节）明确建议把这类 Python 对象换成 Pandas / NumPy / PyArrow 这种不带逐元素引用计数的表示。这是把"操作系统 + Python 实现"的细节落到使用建议的典型例子。

11.23 训练 resume：sampler 状态的保存与恢复

长跑训练（几天几周）必须支持中断恢复。光保存 model + optimizer state 不够 —— dataloader 进度也要保存，否则恢复时从头开始遍历 dataset。

这里要先纠正一个流传很广的误解：PyTorch 核心库里的 Sampler 和 DataLoader 到 v2.11 都没有 state_dict() / load_state_dict()。在 torch/utils/data/sampler.py、distributed.py、dataloader.py 里 grep state_dict 是零结果。所以下面这种写法在原生 PyTorch 上会直接 AttributeError：

checkpoint["sampler"] = sampler.state_dict()   # ✗ 原生 Sampler 没有这个方法

原生能做到的只有最粗的一档：按 epoch 粒度恢复。保存 epoch 与 step，恢复时 sampler.set_epoch(epoch) 让 shuffle 顺序对上，再手动跳过前面若干个 batch：

loader.sampler.set_epoch(ckpt["epoch"])
it = iter(loader)
for _ in range(ckpt["step_in_epoch"]):
    next(it)          # 快进；数据仍然被真读了一遍，只是丢掉

这个办法正确但笨 —— 跳过的样本仍然要走完整的读盘 + transform。真正的解法在 torchdata 的 StatefulDataLoader：它是 DataLoader 的 drop-in 替代，多出 state_dict() / load_state_dict()，机制是让每个 worker 把自己的进度随数据一起报告回主进程，主进程汇总成一份可序列化的状态。要做 batch 级精确 resume，就得引入这个额外依赖。

实战判断：长跑训练值得做 dataloader resume，但先想清楚代价 —— 如果 epoch 本身只有几十分钟，用 set_epoch + 快进这一档往往就够了。

11.24 一次 step 的端到端时序

把所有机制合起来看 一次 for batch in loader: 实际发生了什么：

sequenceDiagram
    participant T as 主进程训练循环
    participant Q as output queue
    participant P as pin_memory thread
    participant W as worker 子进程
    participant DS as Dataset
    T->>+Q: next() 拿 batch N
    Q-->>T: 已有 batch N (pin 完成)
    T->>T: model(batch).backward()<br/>optim.step()
    Note over T: GPU 训练 step N 进行中
    par 同时, worker 在准备 batch N+1
        W->>+DS: dataset[idx]<br/>(磁盘读 + decode + augment)
        DS-->>-W: tensor
        W->>P: data_queue.put((idx, tensor))
        P->>P: tensor.pin_memory()
        P->>Q: output_queue.put((idx, pinned))
    and 主进程做 prefetch
        T->>W: index_queue.put(idx N+M)
    end
    T->>Q: 训练完, next() 拿 batch N+1
    Q-->>T: 已有 batch N+1

整套机制让 worker 数据准备 + pin memory + 主进程 GPU 训练构成 三阶段流水线。理想状态下三个阶段落在不同硬件上并发：worker 吃 CPU、pin 后的 H2D 走 PCIe DMA、训练吃 GPU。流水线打满时 GPU 端不会出现等数据的空窗 —— 能不能打满，取决于最慢的那一级。

这张时序图是 §11.37 工程经验的根本依据 —— 任何一个阶段慢了都会让 pipeline stall：

• worker 慢 → output queue 空 → 主进程 stall
• pin thread 慢（很少见）→ 数据卡在中间
• 主进程慢 → output queue 满 → worker 阻塞在 put

诊断瓶颈的标准方法：用 torch.profiler 看主进程的 "data fetching" 耗时占比、用 py-spy 看每个 worker 的 CPU 利用率。两边数据合起来精准定位。

11.25 DataLoader × 混合精度训练

第 20 章会讲 mixed precision（FP16/BF16）。DataLoader 这层有什么协作？

直觉做法是 worker 直接返回 fp16 tensor 节省传输带宽（pin 内存量减半）。但实战不推荐：

• transform 在 fp16 下精度损失大（特别是图像归一化）
• pin → cuda 后再 cast 到 fp16 比 worker 内 cast 更可控
• worker 跑 transform 时 fp16 的算子反而可能比 fp32 慢（多数 CPU 上 fp16 没有原生加速）

推荐做法：worker 输出 fp32 tensor，主进程 .cuda() 之后再 tensor.to(torch.float16)。这样传输量是 fp16 方案的 2 倍，但只要 PCIe 带宽还没吃满，这部分开销就能被计算 overlap 掉。

例外场景：图像训练 + 大 batch，PCIe 真的成了瓶颈。这时更划算的不是让 worker 输出 fp16，而是让它输出 uint8（每像素 1 字节，传输量只有 fp32 的 1/4），normalize 与 cast 全部挪到 GPU 上用一个 kernel 做完。NVIDIA DALI 走得更远，连 JPEG 解码也放到了 GPU 上。要不要这么做，先量一下 nvidia-smi dmon 里的 PCIe 收发带宽有没有接近卡的规格上限。

11.26 variable batch size 与 max_tokens

NLP 训练常见诉求：按 token 数构造 batch（变长序列）而非按 sample 数。让短句多凑、长句少凑、每 batch 总 token 数稳定。

实现方式：自定义 BatchSampler：

class TokenBudgetSampler(Sampler):
    def __init__(self, lengths, max_tokens=4096):
        self.lengths = lengths              # 每个 sample 的 token 数
        self.max_tokens = max_tokens

    def __iter__(self):
        # 按长度排序后, 滚动凑满 max_tokens
        sorted_idx = sorted(range(len(self.lengths)), key=lambda i: self.lengths[i])
        batch = []
        max_len_in_batch = 0
        for idx in sorted_idx:
            new_max = max(max_len_in_batch, self.lengths[idx])
            new_total = new_max * (len(batch) + 1)
            if new_total > self.max_tokens and batch:
                yield batch
                batch = [idx]
                max_len_in_batch = self.lengths[idx]
            else:
                batch.append(idx)
                max_len_in_batch = new_max
        if batch:
            yield batch

这套 sampler 的收益来自消除 padding 浪费：固定 batch_size 时一个长句会把整批 padding 到它的长度，被算掉的 padding token 全是白算的算力；按 token 预算成批之后，每批的有效 token 占比稳定得多。省下多少完全取决于你的长度分布有多长尾。fairseq、DeepSpeed 等训练库都内置了类似实现。理解 PyTorch 的 BatchSampler 接口让你看到这些上层框架是怎么实现的 —— 没有黑魔法，就是自定义 sampler + collate。

11.27 常见报错信息逐条拆解

生产中常见的 DataLoader 报错：

1. RuntimeError: DataLoader worker (pid xxx) is killed by signal: Killed.

（Linux 上它由主进程的 SIGCHLD handler 抛出，traceback 里紧接着还会出现第 2 条的 exited unexpectedly，以它为 direct cause。）worker 被 SIGKILL 杀掉，最常见的是 OOM killer。通常是 dataset / transform 内存爆。检查：

• 单个 sample 是否占内存过大（如 4K 图片解码到 numpy）
• 是否在 __getitem__ 里 cache 了不释放的对象
• num_workers 是否过多（每个 worker 有独立内存）

2. RuntimeError: DataLoader worker (pid(s) xxx) exited unexpectedly

这是 worker 以任何方式死掉后主进程最终抛出的那一条，真正的死因看 traceback 里它上面的 direct cause：is killed by signal: Segmentation fault.（段错误）、is killed by signal: Killed.（第 1 条），或 exited unexpectedly with exit code N（worker 以非零码退出）。Windows 没有 SIGCHLD，只能看到这一条。段错误与异常退出通常是 C 扩展 bug：

• OpenCV cv2.imread 偶发 segfault（特定图片格式）
• numba JIT 在 worker 里编译失败
• HDF5 跨进程读取的线程问题

排查：用 num_workers=0 跑相同代码看是否报同样错；或单独 import 触发 segfault 的库到主进程隔离测试。

3. RuntimeError: Cannot re-initialize CUDA in forked subprocess

§11.10 提过：fork 启动 + CUDA 不兼容。解法：mp.set_start_method('spawn') 或 'forkserver'。

4. _pickle.PicklingError: Can't pickle <function <lambda> at 0x...>: attribute lookup <lambda> on __main__ failed

（Python 3.10 实测措辞，其他版本略有差异；DataLoader 还会先发一条 "Got pickle error when attempting to start a worker Process" 的 UserWarning。）spawn / forkserver 启动需要 dataset / collate_fn / worker_init_fn 等参数 picklable。把 lambda 改成命名函数：

# 不行
loader = DataLoader(ds, collate_fn=lambda batch: ...)

# 改为
def my_collate(batch): ...
loader = DataLoader(ds, collate_fn=my_collate)

5. RuntimeError: received 0 items of ancdata

Linux 上多进程 + 大量 file descriptor 共享，触发 ulimit 限制。解法：ulimit -n 65536 提高上限，或 torch.multiprocessing.set_sharing_strategy('file_system') 切到 file system 共享（不用 fd）。

6. RuntimeError: generator raised StopIteration

IterableDataset 的 __iter__ 写成生成器时，里面直接 next(inner_iter) 耗尽了内层迭代器，StopIteration 会按 PEP 479 被改写成这个 RuntimeError，经 ExceptionWrapper 在主进程以 Caught RuntimeError in DataLoader worker process 0. 抛出。生成器里提前 return 是安全的（worker 正常收到 StopIteration、发 _IterableDatasetStopIteration 通知主进程）；修复是用 for x in inner_iter: yield x 或捕获 StopIteration 后 return，而不是让它从生成器里漏出去。

逐条理解这些错误能覆盖大部分常见的 DataLoader 报错。生产代码遇到上面任意一个先按这个清单排查。

11.28 N×M 分片：DDP × 多 worker 的合作

第 17 章 DDP 章会展开，这里先给"分片合作"的全景：

graph TB
    DS[全局 dataset<br/>1.28M sample]
    DS --> R0[rank 0 分片<br/>320K sample]
    DS --> R1[rank 1 分片<br/>320K sample]
    DS --> R2[rank 2 分片<br/>320K sample]
    DS --> R3[rank 3 分片<br/>320K sample]
    R0 --> W00[worker 0-0]
    R0 --> W01[worker 0-1]
    R0 --> W02[worker 0-2]
    R0 --> W03[worker 0-3]
    R1 --> W10[...]
    style DS fill:#fef3c7,stroke:#f59e0b

两层分片机制：

• 第一层 DistributedSampler：rank i 拿 [i, i+N, i+2N, ...] 索引（§11.13）
• 第二层 worker 派发：rank i 的主进程把本 rank 的 batch 按 round-robin 派给 M 个 worker（map-style 由 DataLoader 自动完成，§11.20；IterableDataset 要自己按 get_worker_info() 切，§11.12）

合起来：4 rank × 4 worker = 16 个进程并行处理全局 dataset，互不重复。

特殊情况：IterableDataset + DDP。第一层不能直接套 DistributedSampler（IterableDataset 没 __getitem__），要在 dataset 内部按 (rank, num_workers) 双重分片：

class StreamDataset(IterableDataset):
    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        rank = dist.get_rank() if dist.is_initialized() else 0
        world = dist.get_world_size() if dist.is_initialized() else 1
        n_workers = worker_info.num_workers if worker_info else 1
        worker_id = worker_info.id if worker_info else 0

        # 全局有 world * n_workers 个数据消费者
        global_id = rank * n_workers + worker_id
        global_n = world * n_workers
        files = self.files[global_id::global_n]
        ...

这是流式数据训练（webdataset、kafka 流）必写的分片模板。理解它让你跑 DDP 时数据不会重复或遗漏。

11.29 ManagerWatchdog：父子进程死亡感知

ManagerWatchdog 在 torch/utils/data/_utils/worker.py 里按平台定义了两份：Windows 版在 line 30（用 kernel32.OpenProcess + WaitForSingleObject 等父进程句柄），非 Windows 版在 line 62，运行时实例化在 line 314。下面是非 Windows 版的完整实现：

class ManagerWatchdog:
    def __init__(self) -> None:
        self.manager_pid = os.getppid()
        self.manager_dead = False

    def is_alive(self) -> bool:
        if not self.manager_dead:
            self.manager_dead = os.getppid() != self.manager_pid
        return not self.manager_dead

机制：worker 启动时记下父进程 PID（manager_pid）。每次 worker 主循环迭代检查 os.getppid()，如果不等于启动时记的 PID → 父进程已死、worker 自己也该退出。

为什么这样可以检测父死亡？Linux 下父进程死后子进程会被 init（PID 1，或最近的 subreaper）收养—— getppid() 返回值改变就是父死亡信号。比 signal handler 更鲁棒（信号可能被 mask、可能竞争）。

这里要澄清一个常见的张冠李戴：PyTorch 确实封装了 prctl(PR_SET_PDEATHSIG, ...)（torch/csrc/multiprocessing/init.cpp:34 的 _prctl_pr_set_pdeathsig），但它的调用点是 torch/multiprocessing/spawn.py:84，用于 mp.spawn 拉起的进程（比如 DDP 的多进程启动），传的信号是 SIGINT 而不是 SIGKILL。DataLoader 的 worker 走的不是这条路，它靠的就是上面这个轮询 getppid() 的 watchdog，加上 §11.11 讲的那套 shutdown 协议。

理解这个机制也就理解了它的边界：watchdog 是轮询的，只在每次主循环迭代（即每次 index_queue.get 超时或返回）时检查一次，所以空闲的 worker 从父进程死亡到自己退出最长会隔一个 MP_STATUS_CHECK_INTERVAL（5 秒），正在 fetch 的 worker 则要等手上这批做完。这已经足够避免孤儿 worker 堆积，但不是"父死子立刻死"。

11.30 DataLoader 演进时间线

PyTorch DataLoader 的迭代轨迹：

与其按版本号排时间线（这些日期很容易记串），不如看 v2.11 的 DataLoader.__init__ 签名本身 —— 每个参数都是某一轮生产反馈留下的年轮：

参数	解决的问题	本章位置
num_workers	单进程读数据跟不上 GPU	§11.4
pin_memory / pin_memory_device	H2D 拷贝无法异步（pin_memory_device 在 v2.11 已标为 deprecated，实际用当前 accelerator）	§11.6、§11.21
persistent_workers	每个 epoch 重启 worker 的固定开销	§11.7
prefetch_factor	缓冲深度需要按 IO 抖动调	§11.8
timeout	worker 卡死时训练无限 hang	§11.16
multiprocessing_context	fork / spawn / forkserver 的平台与 CUDA 差异	§11.10、§11.14
in_order	严格轮转导致快 worker 被慢 worker 拖住	§11.20

整体趋势很清楚：v1.x 阶段在磨多进程的鲁棒性（错误处理、shutdown 协议、CoW、信号处理），v2.x 阶段在磨吞吐（如 v2.6 加入的 in_order 乱序返回；stateful resume 则留在 torchdata 的 StatefulDataLoader，没进主仓）。每个参数背后都是一类真实踩过的坑，而不是设计时就想全的。

11.31 DataLoader vs DALI vs FFCV

NVIDIA DALI 与 FFCV 是 ImageNet 训练里常被提到的"替代品"。它们与 PyTorch DataLoader 的核心区别：

维度	PyTorch DataLoader	NVIDIA DALI	FFCV
transform 跑在哪	CPU（worker 进程）	GPU（CUDA kernel）	CPU（精心优化的 SIMD）
数据格式	任意（用户写 Dataset）	jpg / mp4 / TFRecord	自家 .beton 格式（预处理）
集成成本	零	中（需重写 pipeline）	中（需先转 .beton）
主要收益来源	—	把解码与 augment 从 CPU 挪到 GPU	预处理成紧凑二进制，消除解码与随机 IO
灵活性	最高	中	低（需预处理）

何时选谁：

• 多数训练：PyTorch DataLoader（足够好、生态最广）
• ImageNet / 视频训练 + GPU 富裕：DALI（GPU 解码 + augment 让 CPU 不再是瓶颈）
• 极致 throughput + 数据格式可控：FFCV（预处理过的二进制让 IO 接近 PCIe 极限）

PyTorch DataLoader 的核心优势是通用：Dataset 里能写任意 Python，代价是 transform 只能吃 CPU。DALI 与 FFCV 各自用"换执行位置"和"换数据格式"来绕开这个瓶颈，代价是灵活性 —— DALI 的 pipeline 要用它自己的算子表达，塞不进任意 Python 逻辑；FFCV 要求先把数据集离线转换一遍。

判断要不要切的标准只有一个：先确认 CPU 侧真的是瓶颈（§11.33 的 profiling 流程）。如果 GPU 利用率本来就接近满，换任何数据加载库都不会更快。

11.32 一段总结：DataLoader 是个并行编排引擎

把全章看到的拼起来：DataLoader 表面是"数据加载工具"，骨子里是个并行编排引擎：

• worker pool（进程级并发）
• pin_memory thread（线程级 + DMA 异步）
• index queue / output queue（生产消费解耦）
• sampler 状态机（next index 决定）
• ManagerWatchdog（父进程死亡感知）
• ExceptionWrapper（异常跨进程传播）
• shared memory storage（零拷贝传输）

PyTorch 把这些工程零件封装在一个 1700 来行的文件（torch/utils/data/dataloader.py）里，对用户暴露的就是 for batch in loader。理解每个零件的存在意义、能让你在这个抽象层之外做高级优化（如自己写 streaming pipeline）也不会丢核心思想。

11.33 DataLoader 性能 profiling 实战

诊断 DataLoader 是否成为瓶颈的标准方法，按精度从粗到细：

第 1 步：nvidia-smi 看 GPU 利用率

watch -n 0.5 nvidia-smi --query-gpu=utilization.gpu --format=csv

如果数字反复跳 0% / 100% / 0% / 100%（sawtooth pattern）→ GPU 在等数据。

第 2 步：看主进程是不是 CPU 瓶颈

top -p $(pgrep -f train.py | head -1)

主进程 CPU > 90% 说明它自己处理数据成了瓶颈（如做大量 .cuda() 调用）。

第 3 步：torch.profiler 量化 data loading 耗时

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=5),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log')
) as prof:
    for batch in loader:
        with torch.profiler.record_function("training_step"):
            train_step(batch)
        prof.step()

打开 TensorBoard / Perfetto 看 trace，"enumerate(DataLoader)#_MultiProcessingDataLoaderIter.__next__"（拿 batch，num_workers=0 时类名是 _SingleProcessDataLoaderIter）耗时与 "training_step" 耗时对比（下面的阈值是经验值）：

• 拿 batch 耗时 > 训练耗时 50% → DataLoader 瓶颈
• 拿 batch 耗时 < 5% → DataLoader 不是瓶颈，找别处

第 4 步：py-spy 看 worker 在干啥

py-spy dump --pid <worker_pid>

看 worker stack：

• 卡在 read() / decode_jpeg() → IO 或 decode 瓶颈
• 卡在 pickle.loads → 数据格式问题（用了大量 Python 对象）
• 卡在 numpy 函数 → CPU augment 慢

第 5 步：profile transform

from time import perf_counter

def timed_getitem(self, idx):
    t0 = perf_counter()
    sample = self.original_getitem(idx)
    print(f"getitem took {(perf_counter()-t0)*1000:.1f}ms")
    return sample

逐层定位到具体瓶颈，再针对性优化（换库、加 num_workers、预处理到磁盘等）。

这套从粗到细的流程的价值在于每一步都能排除掉一大类可能，避免一上来就盲调 num_workers 碰运气。GPU 是训练里最贵的一项资源，让它闲着等数据就是直接浪费预算。

11.34 transform 库性能基准

ImageNet 训练里一张图的标准 pipeline："JPEG 解码 → resize → crop → normalize → tensor"。不同库实现 5 步的耗时差距巨大：

库	主要差异	适用场景
PIL（torchvision 默认后端）	纯标量 C 实现，功能最全	兼容性优先
PIL-SIMD	PIL 的 SIMD fork，API 完全兼容	drop-in 替换，零代码改动
OpenCV (cv2)	解码与几何变换有 SIMD/多线程实现	比 PIL 快，功能面更广
Albumentations	基于 cv2，augment 链本身也做了优化	augment 种类多、链路长时
torchvision.transforms.v2	支持在张量上批量变换、可跑在 GPU 上	从 v1 迁移成本最低
NVIDIA DALI	解码与 augment 都在 GPU	CPU 是瓶颈、GPU 有余量时
FFCV	离线转成 .beton，运行时几乎不解码	数据集固定、长跑训练

升级路径建议（改动成本递增）：

1. torchvision.transforms → transforms.v2：接口基本兼容，还能把变换挪到张量/GPU 上做
2. PIL → PIL-SIMD：pip install pillow-simd，零代码改动
3. torchvision → Albumentations：要重写 transform 链，换来更多 augment 选择
4. CPU augment → DALI：要重写整条 pipeline，CPU 瓶颈严重时收益最大
5. 运行时 augment → FFCV 预处理：先离线转一次数据集，长跑训练里摊销掉

这张表刻意不给具体毫秒数：单图耗时高度依赖图片分辨率、JPEG 编码参数、CPU 型号与线程配置，任何抄来的数字换台机器就不成立。正确做法是在自己的数据上跑一段 §11.33 第 5 步那样的计时，再决定升级到哪一档。

11.35 远程存储训练：NFS / S3 / 对象存储的特殊问题

云上训练 dataset 常放远程存储（NFS、Lustre、S3、GCS）。DataLoader 这层的注意点：

NFS / Lustre：

• worker 起来时每个进程都打开同一文件会触发"client cache invalidation storm" → 训练前几秒卡顿
• 解决：用 os.O_DIRECT 跳过 client cache，或预热（训练前用 cat dataset/* > /dev/null 让 server 端加载）
• 文件元信息查询慢（os.path.getsize 等），把 dataset 元数据预先存成 numpy / parquet 避免逐文件 stat

S3 / GCS / OSS（对象存储）：

• 单个 GET 请求的延迟远高于本地顺序读
• num_workers 与 prefetch_factor 都要比本地盘时开大得多（具体值按实测调）才能 hide 延迟
• 用 s3fs 或 aiobotocore 的并发下载，单 batch 内多文件并行 GET
• 大数据集应该预 shard 成 webdataset / TFRecord 格式（每个文件几百 MB，少 GET 多吞吐）

三类存储的瓶颈性质不同，这决定了该怎么调：

存储	瓶颈性质	对应手段
本地 NVMe	带宽有限、延迟极低	worker 数按 CPU 核数配即可
NFS / Lustre	元数据操作慢，小文件放大	预热 + 元数据预先离线化 + 合并小文件
S3 / GCS / OSS	单请求延迟高，但可并发度很高	加大 num_workers × prefetch_factor 用并发盖延迟 + 预 shard 成大文件

关键区别在最后一行：对象存储的问题是延迟不是带宽，而延迟可以用并发掩盖 —— 这正是 num_workers × prefetch_factor 这两个旋钮的用武之地。所以远程存储不是 DataLoader 的死穴，但默认参数（num_workers=0；开了 worker 后 prefetch_factor 默认 2）几乎一定跑不动，必须按上表调。

11.36 一道真实排查题

最后看一类反复出现的现象：DDP 训练、已开 persistent_workers=True，第一个 epoch 的 GPU 利用率明显偏低，从第二个 epoch 起就上去了。怎么解释？

诊断思路：第一与后续 epoch 的差异。能想到的原因：

• persistent_workers=False：每个 epoch worker 重启 → 但已经开了 persistent，排除
• OS file cache 预热：第一 epoch 数据全从磁盘读 → 后续从 OS page cache 读。这是最常见原因
• CUDA kernel 预热：第一次见到的 kernel 形状要 JIT compile / autotune → 第一 batch 慢。但只影响开头，不解释整个 epoch

最常见的原因是 OS file cache 预热：只要 dataset 总体积明显小于机器空闲 RAM，第一个 epoch 就把它整个读进了 page cache，后续 epoch 实际是在从内存读。判断方法很简单 —— 看 free -m 的 buff/cache 列在第一个 epoch 里是不是一路涨到接近 dataset 大小。

修复方法：训练前 warmup（find /data -type f -exec cat {} > /dev/null \; 让 OS 提前缓存），或者降低 dataset 大小让缓存命中率高。这是远程存储或冷盘上训练的隐藏 quirk —— 看一眼 free -m 的 cached 列就能确认。

理解 DataLoader 的工作机制 + OS 文件系统行为，能让你解释训练曲线的细节、做出正确干预。这是一切性能调优的起点。

11.37 几条工程经验

实战 DataLoader 相关：

1. num_workers 经验值：通常设为 CPU 核心数的 1/2 到 1（如 16 核机器设 8-16）。再多就 thrashing

2. 监控 GPU 利用率：如果训练时 GPU 利用率明显偏低，先怀疑 DataLoader 跟不上。增加 num_workers 或 prefetch_factor

3. 第一次 epoch 慢：worker 初次启动 + 第一次 sample 触发 page fault。persistent_workers=True 后续 epoch 就快

4. pin_memory + non_blocking=True 必须配对：单开 pin_memory 但 cuda 调用没用 non_blocking，等同于没开

5. transform 应该在 worker 里跑：把 augment 写到 dataset 的 __getitem__ 里，而非主循环里。主循环只做 GPU 计算

6. NumPy in worker 的 seed 坑：fork 启动时 numpy 全局 RNG 在所有 worker 里相同（fork 复制了主进程的内存）。PyTorch 在 worker 启动时会 reseed numpy，但如果你在 worker 里 import 了某个用 np.random 的库，要确保它在 worker 启动之后 import

7. torch.utils.data.get_worker_info()：在 dataset 内可以查当前 worker id / num_workers，用于"按 worker 分片"（IterableDataset 必备）

8. 子进程 dataset 状态独立：worker 是 fork 出来的，每个 worker 有独立的 dataset 副本。如果你想在 dataset 里 cache 某些计算结果（如解析 json），每个 worker 各自缓存一份，不共享

11.38 一个典型优化案例：DataLoader 跟不上 GPU

社区里反复出现的最经典 DataLoader 瓶颈案例（PyTorch forum / GitHub issue 反复讨论的问题模式）：

症状：图像分类训练，GPU 利用率长期偏低
定位：
  1. nvidia-smi -l 1 看到利用率呈 sawtooth（100% → 0% → 100%），是典型的 "GPU 等数据"
  2. py-spy dump 到 worker，栈停在 PIL.Image.open / JPEG 解码里
  3. 数一下：num_workers 远小于机器核数，且每个 worker 被 set_num_threads(1) 限成单线程
判断：
  - 单个 batch 的准备时间 > 单个 step 的 GPU 计算时间，prefetch 队列被抽空
可选手段（按改动成本递增，逐个验证效果）：
  1. 调大 num_workers（上限是物理核数，见 §11.37 第 1 条）
  2. 换更快的解码/变换后端（§11.34 的升级路径）
  3. 离线预处理：把解码与 resize 提前做掉，运行时只读定长二进制

这种"先定位是哪一级最慢、再逐条验证"是 DataLoader 调优的标准流程。每一步都要用同一套指标复测，否则很容易把三个改动一起上，最后不知道是哪个起了作用。理解整章讲的"worker 池 + 三段 queue + pin_memory 线程"机制，就能知道每次该量哪一级。

11.39 跨专栏关联

• 《vLLM 推理内核深度解析》专栏第 17 章（API 服务器）：vLLM 推理时的 batched request 处理与 DataLoader 的 batch 思想一致 —— 把多个独立请求合成一个 GPU 友好的 batch
• 《Tokio 源码深度解析》专栏第 13 章（channels：mpsc / broadcast / watch / oneshot）：DataLoader 的 multiprocessing.Queue 是同步阻塞 channel，思想与 Tokio 的 mpsc::channel 相通 —— 都是 producer-consumer 模式
• 《Rust 编译器与运行时揭秘》专栏第 17 章（增量编译）：编译器的 dependency graph build 也用 worker pool + queue 处理独立任务

11.40 几条 DataLoader 设计的"通用启示"

把 DataLoader 思想抽象到任何"生产 / 消费数据流水线"系统：

第一：生产与消费独立 + queue 解耦：让两者不用互相等待。Queue 长度（prefetch_factor）调节缓冲深度

第二：worker pool 而非每次 fork：进程启动有不可忽略的固定成本，长跑任务必须 pool。persistent_workers 是这条原则的体现

第三：异常通过序列化跨进程传递：直接 throw 在多进程里没用，必须 wrap 成可序列化对象。ExceptionWrapper 是模板

第四：用 thread 处理需要主进程上下文的工作：pin_memory 用 thread 不用 process，因为它要访问 CUDA context

第五：RNG seed 每个 worker 不同 + 可复现：由主进程派发一个 base_seed、各 worker 从它派生自己的 seed，既保证不重复又保证可复现。注意派生方式要看底层 RNG 的脾气 —— PyTorch 对 torch/random 用简单相加，对 numpy 则专门做了哈希混合（§11.5.2）

第六：torch.set_num_threads(1) 防多 worker × 多 thread 抢核：任何 worker pool 都要决定每个 worker 自己开多少 thread。设错了性能崩盘

把这六条记住，写自己的多进程数据流水线（如 ETL 系统、爬虫、批处理任务）能避开几乎所有大坑。

下一章拆 TorchDynamo —— PyTorch 2.0 编译器栈的入口，看它怎么用 PEP 523 帧拦截 + 字节码分析把动态 Python 代码捕获成 FX Graph。