124x Slower: What PyTorch DataLoader Actually Does at the Kernel Level

TL;DR

PyTorch的DataLoader可以比直接张量索引慢50-124倍，用于内存GPU工作负载。我们在RTX 4090上复制了一个真实的PyTorch问题，并跟踪了每个CUDA API调用和Linux内核事件以找到根本原因。GPU并不慢 – 它正在挨饿。DataLoader工作者在40秒内生成了200,000个CPU上下文切换和300,000个页面分配，导致GPU平均等待301ms进行每次数据传输，这本应只需几微秒。

问题

一个PyTorch用户报告说，DataLoader比简单的MLP推理工作负载中的直接张量索引慢7-22倍。即使使用num_workers=12，pin_memory=True和prefetch_factor=12，差距仍然很大。GPU利用率停留在10-20%。

我们复制了它。在我们的硬件上，差距甚至更大：

工作负载很简单：700万个样本，100个特征，2层MLP，批次大小1M。模型在几毫秒内处理一个批次。那么时间去哪了？

nvidia-smi显示的内容

没有什么有用的信息。GPU利用率在0%和30%之间闪烁。内存使用情况稳定。温度正常。GPU显然被低利用，但nvidia-smi无法告诉你为什么。

torch.profiler显示的内容

报告者尝试使用PyTorch的内置profiler，并“获得了无意义的跟踪数据”。这是一个常见的沮丧 – 应用程序级别的profiler可以显示哪些CUDA内核正在运行，但它们无法看到决定数据是否及时到达GPU的主机端调度、内存和进程生命周期事件。

内核级别跟踪显示的内容

我们在跟踪CUDA API调用（通过libcudart.so上的eBPF uprobes）和Linux内核事件（调度器上下文切换、内存页面分配、进程fork）同时运行基准测试。结果告诉了完整的故事。

完整视频演示： https://asciinema.org/a/RGwhPeXAPJdhXqxp

在视频中，我们将一个开放权重LLM（MiniMax-M2.7）通过MCP（模型上下文协议）连接到跟踪数据库：

JSON配置告诉MCP客户端在哪里找到Ingero服务器以及要加载哪个跟踪数据库：

{
"mcpServers": {
"ingero": {
"command": "./bin/ingero",
"args": ["mcp", "--db", "investigations/pytorch-dataloader-starvation.db"]
}
}
}

这使LLM能够通过7个工具直接访问跟踪数据：get_trace_stats、get_causal_chains、get_per_process_breakdown等。AI可以查询数据库，关联CUDA事件与内核调度数据，并在没有手动分析的情况下生成纯语言诊断。

4个高严重性因果链

因果链引擎检测到4个高严重性模式，所有模式都有相同的根因：

[HIGH] cudaStreamSync p99=42ms（1,638x p50=25us）- CPU 100% + 1,880 sched_switch 事件
时间线：
[SYSTEM] CPU 100%
[HOST ] 1,880个上下文切换（21s off-CPU）
[CUDA ] p99=42ms（1,638x p50=25us）
根因：DataLoader工作者争夺CPU，巨大的页面分配压力

[HIGH] cudaLaunchKernel p99=24.67ms（349x p50=70us）- CPU 100%
根因：34个sched_switch事件

[HIGH] cuMemAlloc p99=627us（4.0x p50）- CPU 100%
[HIGH] cuLaunchKernel p99=106us（4.0x p50）- CPU 100%

cudaStreamSync p99是p50的1,638倍。那不是GPU慢 – 那是GPU等待永远不会及时到达的数据。

图1：AI生成的分析，运行/investigate。模型使用Ingero的7个MCP工具查询跟踪数据库，并直接从跟踪数据中生成了带有可行建议的纯语言解释。

每个进程的细分

这就是它变得清晰的地方。主进程和其4个DataLoader工作者作为单独的实体可见：

主进程：

- cudaMemcpyAsync（主机到设备传输）：平均301ms，最大2.9秒
- cudaStreamSync：p99 = 42ms（正常25us）
- 1,567个上下文切换，平均16ms off-CPU，最大停顿5秒
- 799,018个页面分配

DataLoader工作者1：52,863个上下文切换，89,338个页面分配，最大停顿5s
DataLoader工作者2：50,638个上下文切换，83,509个页面分配，最大停顿5s
DataLoader工作者3：49,361个上下文切换，70,035个页面分配，最大停顿5s
DataLoader工作者4：38,862个上下文切换，56,354个页面分配，最大停顿5s

工作者总计：~191,000个上下文切换和~299,000个页面分配，耗时40秒。

这意味着什么

DataLoader工作者正在做三件直接索引完全避免的事情：

1. 洗牌和索引：带有shuffle=True的DataLoader生成一个索引的随机排列，然后每个工作者选择其块。这需要在整个700万个样本张量中进行随机内存访问 – 对缓存局部性非常糟糕，并触发页面错误。
2. 合并和复制：每个工作者从分散的样本中收集连续的批次张量。这意味着分配新内存（页面分配），从随机位置复制数据（缓存缺失），并将结果序列化回主进程，方法是共享内存或队列。
3. 争夺CPU：四个工作者 + 主进程在4 vCPU机器上意味着不断的抢占。每个工作者被取消调度50,000次。最坏的情况停顿是5秒 – 在此期间，GPU没有什么可处理的。

使用直接索引：X[i:i+batch_size]是已经在内存中的连续张量的零复制视图。.to(device)触发从单个连续区域的一个DMA传输。没有工作者，没有洗牌，没有合并，没有跨进程复制，没有上下文切换。GPU在几微秒内获得数据，而不是数百毫秒。

解决方案

对于GPU工作负载，整个数据集适合RAM：

1. 不要使用DataLoader。带有预洗牌索引数组的直接索引更简单，速度快100倍：

   indices = torch.randperm(num_samples)
   for i in range(0, num_samples, batch_size):
   batch = X[indices[i:i+batch_size]].to(device)
   output = model(batch)
   

2. 如果您必须使用DataLoader，将num_workers与实际CPU核心数减1相匹配。在4核机器上，num_workers=2减少争用。添加persistent_workers=True以避免fork开销。
3. 对于大于内存的数据集，DataLoader是必要的，真正的瓶颈转移到磁盘I/O。使用prefetch_factor=2（不要更高 – 更多的prefetching意味着更多的内存压力），并确保您的存储可以跟上。

更大的图景

这次调查说明了我们在GPU工作负载中看到的模式：GPU很快，主机是瓶颈，GPU指标无法看到这一点。nvidia-smi报告了低利用率，但无法解释为什么。torch.profiler捕获了CUDA内核，但错过了用户空间发生的20万个上下文切换。

唯一能看到完整图景的方法是同时跟踪两边 – 库级CUDA API调用和Linux内核调度事件 – 并按时间和进程ID关联它们。因果链“CPU 100% -> 1,880 sched_switch -> cudaMemcpyAsync 301ms -> cudaStreamSync 42ms”用一行告诉了完整的故事。没有跨栈跟踪，这将仍然是一个谜 – 就像原始报告者花了几周时间调试它一样。

图2：当被问及“核心问题是什么？”时，模型确定CPU超额订阅导致主机端调度延迟。cudaLaunchKernel从73us变慢到25.8ms（356x慢），因为CPU无法及时调度启动。

自己试试

复制基准测试：

import torch, time

from torch.utils.data import DataLoader
X = torch.randn(7_000_000, 100)

model = torch.nn.Sequential(

torch.nn.Linear(100, 512), torch.nn.ReLU(),

torch.nn.Linear(512, 512), torch.nn.ReLU(),

torch.nn.Linear(512, 10)

).cuda()
# 快速路径

start = time.time()

with torch.no_grad():

for i in range(0, len(X), 1_048_576):

model(X[i:i+1_048_576].cuda())

torch.cuda.synchronize()

print(f'直接：{time.time()-start:.3f}s')
# 慢路径

loader = DataLoader(X, batch_size=1_048_576, shuffle=True)

start = time.time()

with torch.no_grad():