124 раз повільніше: Що PyTorch DataLoader фактично робить на рівні ядра

TL;DR

DataLoader PyTorch може бути на 50-124 рази повільніше за прямий індексування тензорів для вбудованих в пам’яті GPU робіт. Ми відтворили реальну проблему PyTorch на RTX 4090 і відслідкували кожний виклик CUDA API та подію Linux ядра, щоб знайти причину. GPU не був повільним – він голодував. Робочі процеси DataLoader генерували 200 000 перемикань контексту CPU та 300 000 виділень сторінок за 40 секунд, залишаючи GPU в очікуванні в середньому 301 мілісекунду на кожен перехід даних, який повинен займати мікросекунди.

Проблема

Користувач PyTorch повідомив, що DataLoader був на 7-22 рази повільніше за прямий індекс тензорів для простого MLP-інференсу. Навіть з num_workers=12, pin_memory=True та prefetch_factor=12 розрив залишався величезним. Використання GPU становило 10-20%.

Ми відтворили це. Розрив був ще гіршим на нашому обладнанні:

Робота тривіальна: 7 млн зразків, 100 ознак, 2-шаровий MLP, розмір партії 1 млн. Модель обробляє партію за мілісекунди. Тому де йде час?

Що показує nvidia-smi

Нічого корисного. Використання GPU миготить між 0% і 30%. Використання пам’яті стабільне. Температура нормальна. GPU явно недовантажений, але nvidia-smi не може сказати чому.

Що показує torch.profiler

Повідомник спробував вбудований профайлер PyTorch і “отримав жодних значимих даних трасування”. Це звичайна проблема – профайлери на рівні застосунку можуть показати, які ядра CUDA виконуються, але вони не бачать планування на стороні хоста, події з пам’яттю та життєвий цикл процесів, які визначають, чи прибуває дані вчасно на GPU.

Що показує трасування на рівні ядра

Ми запустили бенчмарк, відслідковуючи одночасно виклики CUDA API (за допомогою eBPF uprobes на libcudart.so) та події Linux ядра (перемикання контексту, виділення сторінок пам’яті, форк процесів). Результати розповідають повну історію.

Повне відео: 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 перемикань контексту
Хронологія:
[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 перемикання контексту

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

cudaStreamSync p99 становить 1,638 рази p50. Це не повільність GPU – це GPU, який чекає на дані, які ніколи не прибувають вчасно.

Фігура 1: Аналіз, згенерований AI після запуску /investigate. Модель використала 7 інструментів MCP Ingero для запиту бази даних трасування та створила діагноз у вигляді простої мови з рекомендаціями, отриманими безпосередньо з даних трасування.

Розбивка за процесами

Це стає зрозумілим. Основний процес і його 4 робочих процеси DataLoader видно як окремі сутності:

Основний процес:

- cudaMemcpyAsync (перехід з хоста на пристрій): середнє 301 мілісекунда, максимальне 2,9 секунди
- cudaStreamSync: p99 = 42 мілісекунди (зазвичай 25 мікросекунд)
- 1,567 перемикань контексту, середнє 16 мілісекунд поза 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 роблять три дорогих речі, яких прямий індексування уникнути:

1. Перемішування та індексування: DataLoader з shuffle=True генерує випадкову пермутацію індексів, потім кожен робочий процес вибирає свій фрагмент. Це вимагає випадковий доступ до пам’яті по всьому 7-мільйонному тензору зразків – погано для локальності кешу та викликає помилки сторінок.
2. Агрегування та копіювання: Кожен робочий процес збирає розрізнені зразки в контигуальний тензор партії. Це означає виділення нової пам’яті (виділення сторінок), копіювання даних з випадкових місць (промахи кешу) та серіалізацію результату назад у основний процес через спільну пам’ять або чергу.
3. Конкуренція за CPU: Чотири робочих процеси плюс основний процес на 4-ядерному процесорі означають постійну перемикання контексту. Кожен робочий процес відключається 50 000 раз. Найгірша затримка становить 5 секунд – протягом якої GPU нічого не обробляє.

З прямим індексуванням: X[i:i+розмір_партії] – це нуль-копійний вигляд контигуального тензору, який вже знаходиться в пам’яті. .to(пристрій) викликає один DMA-перехід з однієї контигуальної області. Жодних робочих процесів, жодного перемішування, жодного агрегування, жодної копії між процесами, жодних перемикань контексту. GPU отримує дані за мікросекунди, а не сотні мілісекунд.

Виправлення

Для вбудованих в пам’яті GPU робіт, де весь набір даних поміщається в оперативній пам’яті:

1. Не використовувати DataLoader. Прямий індексування з попередньо перемішаним масивом індексів простіший і на 100 разів швидший:

   індекси = torch.randperm(кількість_зразків)
   for i in range(0, кількість_зразків, розмір_партії):
   партія = X[індекси[i:i+розмір_партії]].to(пристрій)
   вивід = модель(партія)
   

2. Якщо ви повинні використовувати DataLoader, збігайте num_workers з фактичними ядрами CPU minus 1. На 4-ядерному процесорі num_workers=2 зменшує конкуренцію. Додайте persistent_workers=True, щоб уникнути накладних витрат на форк.
3. Для наборів даних більших за пам’ять де DataLoader необхідний, справжня瓶шея зсується до вводу/виводу з диска. Використовуйте prefetch_factor=2 (не вище – більше.prefetching означає більше тиску на пам’ять) та переконайтесь, що ваша система зберігання може впоратися.

Більша картина

Це розслідування ілюструє закономерність, яку ми бачимо постійно у роботах з GPU: GPU швидке, хост -瓶шея, а метрики GPU не можуть пояснити чому. nvidia-smi повідомила низьке використання, але не могла пояснити чому. torch.profiler захопила ядра CUDA, але пропустила 200 000 перемикань контексту, які відбувалися в просторі користувача.

Єдиний спосіб побачити повну картину полягав у тому, щоб відслідкувати одночасно обидва боку – виклики CUDA API на рівні бібліотеки та події планування Linux ядра – і корелювати їх за часом та ідентифікатором процесу. Ланцюг причин “CPU 100% -> 1,880 перемикань контексту -> cudaMemcpyAsync 301 мілісекунда -> cudaStreamSync 42 мілісекунди” розповідає повну історію в одному рядку. Без трасування по всьому стеку це залишилося б загадкою – як і для оригінального повідомника, який провів тижні, відлажуючи це.

Фігура 2: Коли запитано “що є основною проблемою?”, модель ідентифікує перевантаження CPU, яке викликає затримки планування на стороні хоста. cudaLaunchKernel перейшло з 73 мікросекунд до 25,8 мілісекунд (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 # в одному терміналі
python3 benchmark.py # в іншому терміналі
./bin/ingero explain --since 60s # після завершення бенчмарку

Або пропустіть відтворення і дослідіть нашу трасову базу даних безпосередньо. База даних розслідування (764 КБ) знаходиться в репозиторії:

# Перегляньте ланцюги причин з розслідування
./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 (рекомендовано). Найшвидший спосіб проаналізувати трасову базу даних – підключити будь-який клієнт AI, сумісний з MCP, безпосередньо до бази даних. Жодного ручного аналізу не потрібно.

Створіть файл конфігурації:

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
# Додайте конфігурацію вище до налаштувань MCP вашої AI