レポート

124倍遅い:PyTorch DataLoaderがカーネルレベルで実際に何をしているか

mm
Unite.AI を Google の優先ソースに追加
A conceptual widescreen illustration of an hourglass containing glowing digital data streams and circuit patterns, with the bottom half featuring a large data block and a GPU hardware component sitting idle to the side, representing data processing delays and GPU starvation.

この記事は、eBPF uprobesを使用したカーネルレベルのGPUトレース調査の結果に基づいています。トレースデータベースは、独立した検証のためにIngeroのオープンソースリポジトリに公開されています。

TL;DR

PyTorchのDataLoaderは、直接のテンソルインデックス化よりも50〜124倍遅いことがあります。私たちは、RTX 4090で実際のPyTorchの問題を再現し、CUDA API呼び出しとLinuxカーネルイベントをトレースして根本原因を発見しました。GPUは遅くなかった、飢えていたのです。DataLoaderワーカーは40秒で20万回のCPUコンテキストスイッチと30万回のページ割り当てを生成し、GPUは平均301msの遅延を経験しました。

問題

PyTorchユーザーは、DataLoaderが直接のテンソルインデックス化よりも7〜22倍遅いと報告しました。num_workers=12、pin_memory=True、prefetch_factor=12であったにもかかわらず、ギャップは巨大なままでした。GPU利用率は10〜20%でした。

私たちはそれを再現しました。ハードウェアのギャップは私たちのほうがひどかった:

方法 時間 直接との比較
直接のテンソルインデックス化 0.39s 1倍
DataLoader (shuffle=True) 48.49s 124倍遅い
DataLoader (最適化、4ワーカー、pin_memory) 43.29s 111倍遅い

ワークロードはごく普通です:700万サンプル、100フィーチャー、2層のMLP、バッチサイズ1M。モデルはバッチを数ミリ秒で処理します。そこで時間はどこに行きますか?

nvidia-smiの表示

何も役に立たない。GPU利用率は0%と30%の間で変動します。メモリ使用量は安定しています。温度は問題ありません。GPUは明らかに活用されていないですが、nvidia-smiはその理由を示すことができません。

torch.profilerの表示

報告者はPyTorchの組み込みプロファイラーを試しましたが、「有意義なトレースデータを取得できませんでした。」これは、CUDAカーネルを実行していることを示すことができますが、ホスト側のスケジューリング、メモリ、プロセスライフサイクルイベントを示すことができないため、データがGPUに到着するタイミングを決定することはできません。

カーネルレベルトレースの表示

私たちは、CUDA API呼び出しとLinuxカーネルイベントの両方を同時にトレースしました。結果は全体の物語を語っています。

完全なビデオウォークスルー: https://asciinema.org/a/RGwhPeXAPJdhXqxp

ビデオでは、オープンウェイトLLM(MiniMax-M2.7)をMCP(モデルコンテキストプロトコル)を介してトレースデータベースに接続しました:

ollmcp -m minimax-m2.7:cloud -j /tmp/ingero-mcp-dataloader.json

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つのHIGH-severity因果連鎖

因果連鎖エンジンは、すべて同じ根本原因を持つ4つのHIGH-severityパターンを検出しました:

[HIGH] cudaStreamSync p99=42ms (1,638x p50=25us) - CPU 100% + 1,880 sched_switch events
タイムライン:
[SYSTEM] CPU 100%
[HOST ] 1,880コンテキストスイッチ (21秒オフCPU)
[CUDA ] p99=42ms (1,638x p50=25us)
根本原因:DataLoaderワーカーがCPUを争い、ページ割り当て圧力が大きな原因です
[HIGH] cudaLaunchKernel p99=24.67ms (349x p50=70us) - CPU 100%
根本原因:34 sched_switch events

[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:/investigateを実行した後のAI生成分析。モデルは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ページ割り当て、最悪のストール5秒
DataLoaderワーカー2:50,638コンテキストスイッチ、83,509ページ割り当て、最悪のストール5秒
DataLoaderワーカー3:49,361コンテキストスイッチ、70,035ページ割り当て、最悪のストール5秒
DataLoaderワーカー4:38,862コンテキストスイッチ、56,354ページ割り当て、最悪のストール5秒

ワーカー全体の合計:約191,000コンテキストスイッチと299,000ページ割り当て、40秒間。

これは何を意味するのか

DataLoaderワーカーは、直接のインデックス化では全く回避される3つの高価なことを行っています:

  1. シャッフルとインデックス化: DataLoader with shuffle=Trueは、ランダムなインデックスの順列を生成し、各ワーカーはそのチャンクを選択します。これには、7Mサンプルのテンソル全体に対するランダムなメモリアクセスが必要です。キャッシュの局所性が悪く、ページフォールトを引き起こします。
  2. 結合とコピー: 各ワーカーは散在したサンプルを連続したバッチテンソルにまとめます。これには、新しいメモリの割り当て(ページ割り当て)、ランダムな場所からのデータのコピー(キャッシュミス)、および結果のシリアライズが必要です。
  3. CPUの競合: 4つのワーカーとメインプロセスが4つのvCPUマシン上で実行され、常にプリエンプションが発生します。各ワーカーは50,000回デスケジュールされます。最悪のストールは5秒です。その間、GPUは何も処理しません。

直接のインデックス化の場合: X[i:i+batch_size]は、すでにメモリ上にある連続したテンソルのゼロコピービューです。.to(device)は、単一の連続した領域からの1つの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コア数に合わせます。4コアマシンの場合、num_workers=2は競合を減らします。persistent_workers=Trueを追加してフォークオーバーヘッドを回避します。
  3. メモリよりも大きなデータセットの場合、 DataLoaderは必要です。実際のボトルネックはディスクI/Oにシフトします。prefetch_factor=2(より高い値ではなく、より多くのプリフェッチはより多くのメモリ圧力を意味します)を使用し、ストレージが追いつくことを確認します。

大きな絵

この調査は、GPUワークロードで私たちが見るパターンを示しています:GPUは速いですが、ホストはボトルネックであり、GPUメトリクスはそれを見えません。nvidia-smiは低利用率を報告しましたが、理由を示すことができませんでした。torch.profilerはCUDAカーネルをキャプチャしましたが、ユーザースペースで発生する200,000のコンテキストスイッチを逃しました。

全体像を見せる唯一の方法は、CUDA API呼び出しとLinuxカーネルスケジューリングイベントの両方を同時にトレースし、時間とプロセスIDでそれらを関連付けることでした。因果連鎖「CPU 100% → 1,880 sched_switch → cudaMemcpyAsync 301ms → cudaStreamSync 42ms」は、1行で全体の物語を語っています。クロススタックトレースなしでは、これは謎のままだったでしょう。それは、週間を費やしてデバッグした元の報告者にとってもそうでした。

図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}秒')

# 遅いパス 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}秒')

Ingeroでトレースして、下層で何が起こっているかを見てみましょう:

git clone https://github.com/ingero-io/ingero.git
cd ingero && make build
sudo ./bin/ingero trace --duration 60s # 1つのターミナル
python3 benchmark.py # 別のターミナル
./bin/ingero explain --since 60s # ベンチマーク終了後

または、ベンチマークの再現をスキップし、トレースデータを直接調査します。調査データベース(764KB)はリポジトリにあります:

# 因果連鎖を調査する
./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で実行されました。

デビッド・メイルは、Ingeroの共同著者およびメンテナーであり、CUDAレベルのGPU観測可能性用のオープンソースeBPFエージェントです。彼は、プロダクションAIワークロードのカーネルレベルトレースに特化しています。