124 倍慢：PyTorch DataLoader 在内核级别实际做了什么

TL;DR

PyTorch 的 DataLoader 可以比直接 tensor 索引慢 50-124 倍，用于内存 GPU 工作负载。我们在 RTX 4090 上复现了一个真实的 PyTorch 问题，并跟踪了每个 CUDA API 调用和 Linux 内核事件，以找到根本原因。GPU 并不慢 – 它正在等待数据。

问题

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

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

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

nvidia-smi 显示的内容

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

torch.profiler 显示的内容

报告者尝试使用 PyTorch 的内置 profiler，并“没有获得任何有意义的跟踪数据”。这是一个常见的挫折 – 应用级别的 profiler 可以显示 CUDA 内核的运行情况，但不能看到主机侧的调度、内存和进程生命周期事件，这些事件决定了数据是否及时到达 GPU。

内核级别跟踪显示的内容

我们在运行基准测试的同时跟踪了 CUDA API 调用（通过 eBPF uprobes on libcudart.so）和 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 离 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 离 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. 洗牌和索引： DataLoader 带有 shuffle=True 生成一个随机的索引排列，然后每个工作线程选择其 chunk。这需要对整个 7M 样本 tensor 进行随机内存访问 – 对缓存局部性非常糟糕，并触发页错误。
2. 合并和复制： 每个工作线程收集分散的样本到一个连续的批次 tensor。这意味着分配新内存（页分配），从随机位置复制数据（缓存失误），并将结果序列化回主进程通过共享内存或队列。
3. 竞争 CPU： 四个工作线程 + 主进程在一个 4 核 CPU 机器上意味着不断的抢占。每个工作线程被抢占 50,000 次。最坏情况是 5 秒 – 在此期间，GPU 没有任何东西可处理。

使用直接索引： X[i:i+batch_size] 是一个已经在内存中的连续 tensor 的零复制视图。.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 核心数匹配，减少争用。添加 persistent_workers=True 以避免 fork 开销。
3. 对于大于内存的数据集， DataLoader 是必要的，真正的瓶颈转移到磁盘 I/O。使用 prefetch_factor=2（不是更高 – 更多的 prefetching 意味着更多的内存压力），并确保存储可以跟上。

更大的图景

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

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

图 2：当被问及“核心问题是什么？”时，模型确定 CPU 过度订阅导致主机侧调度延迟。cudaLaunchKernel 从 73us 变为 25.8ms（356 倍更慢），因为 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():
for batch in loader:
model(batch.cuda())
torch.cuda.synchronize()
print(f'DataLoader：{time.time()-start:.3f}s')

使用 Ingero 跟踪以查看底层发生了什么：

git clone https://github.com/ingero-io/ingero.git
cd ingero && make build
sudo ./bin/ingero trace --duration 60s # 在一个终端
python3 benchmark.py # 在另一个终端
./bin/ingero explain --since 60s # 基准测试完成后

或者跳过复现并直接探索我们的跟踪数据：

# 查看调查中的因果链
./bin/ingero explain --db investigations/pytorch-dataloader-starvation.db --since 5m

# 每个进程细分（查看 DataLoader 工作线程与主进程）
./bin/ingero explain --db investigations/pytorch-dataloader-starvation.db --per-process --since 5m

# 连接你的 AI 助手进行交互式调查
./bin/ingero mcp --db investigations/pytorch-dataloader-starvation.db

使用 AI 进行调查（推荐）。 分析跟踪的最快方法是将任何 MCP 兼容的 AI 直接连接到数据库。无需手动分析。

创建一个配置文件：

cat > /tmp/ingero-mcp-dataloader.json << 'EOF'
{
"mcpServers": {
"ingero": {
"command": "./bin/ingero",
"args": ["mcp", "--db", "investigations/pytorch-dataloader-starvation.db"]
}
}
}
EOF

然后连接你的模型：

# 使用 Ollama + MiniMax（我们在视频中使用的）
ollmcp -m minimax-m2.7:cloud -j /tmp/ingero-mcp-dataloader.json

# 使用 Claude Code
claude --mcp-config /tmp/ingero-mcp-dataloader.json

# 使用任何 MCP 兼容客户端
# 将上面的配置添加到你的 AI 的 MCP 设置中

输入 /investigate 以触发引导式分析，或提出任何问题：“是什么导致了 GPU 饥饿？”AI 有权访问 7 个工具，这些工具直接查询跟踪数据库。

GitHub： github.com/ingero-io/ingero
原始问题： pytorch/pytorch#154318
视频演示： https://asciinema.org/a/RGwhPeXAPJdhXqxp

调查在 TensorDock RTX 4090（24GB）、Ubuntu 22.04、PyTorch 2.10.0+cu128 上进行。