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%.

Мы воспроизвели это. Разрыв был еще хуже на нашем оборудовании:

Нагрузка тривиальна: 7M образцов, 100 признаков, 2-слойный MLP, размер пакета 1M. Модель обрабатывает пакет за миллисекунды. Итак, куда делось время?

Что показывает 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 = 42ms (обычно 25us)
- 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 генерирует случайную перестановку индексов, затем каждый работник выбирает свою часть. Это требует случайного доступа к памяти по всему 7M-образцу тензора – это плохо для кэш-локальности и вызывает ошибки страниц.
2. Сбор и копирование: Каждый работник собирает разбросанные образцы в контигуальный тензор пакета. Это означает выделение новой памяти (выделения страниц), копирование данных из случайных мест (промахи кэша) и сериализацию результата обратно в основной процесс через общую память или очередь.
3. Конкуренция за CPU: Четыре работника + основной процесс на 4-ядерной машине означают постоянную прерывание. Каждый работник отменяется 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 необходим, реальная проблема смещается в сторону ввода-вывода с диска. Используйте prefetch_factor=2 (не выше – больше предварительного чтения означает больше давления на память) и убедитесь, что ваше хранилище может справиться.

Большая картина

Это расследование иллюстрирует закономерность, которую мы видим постоянно в нагрузках 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 (рекомендуется). Быстрейший способ проанализировать трассировку – подключить любой совместимый с 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 клиентом
# Добавьте конфиг выше в настройки MCP вашего AI

Введите /investigate, чтобы запустить руководимый анализ, или задайте любой вопрос: “Что вызвало голодание GPU?” AI имеет доступ к 7 инструментам, которые запрашивают базу данных трассировки напрямую.

GitHub: github.com/ingero-io/ingero
Оригинальная проблема: pytorch/pytorch#154318
Видео-демо: https://asciinema.org/a/RGwhPeXAPJdhXqxp

Расследование проведено на TensorDock RTX 4090 (24 ГБ), Ubuntu 22.04, PyTorch 2.10.0+cu128.