124 veces más lento: ¿Qué hace realmente PyTorch DataLoader a nivel de kernel

TL;DR

El DataLoader de PyTorch puede ser 50-124 veces más lento que el índice de tensor directo para cargas de trabajo de GPU en memoria. Reprodujimos un problema real de PyTorch en un RTX 4090 y rastreamos cada llamada a la API de CUDA y cada evento del kernel de Linux para encontrar la causa raíz. La GPU no era lenta, estaba pasando hambre. Los trabajadores de DataLoader generaron 200,000 conmutaciones de contexto de CPU y 300,000 asignaciones de página en 40 segundos, lo que dejó a la GPU esperando un promedio de 301ms por transferencia de datos que debería tomar microsegundos.

El Problema

Un usuario de PyTorch informó que DataLoader era 7-22 veces más lento que el índice de tensor directo para una carga de trabajo de inferencia de MLP simple. Incluso con num_workers=12, pin_memory=True y prefetch_factor=12, la brecha permaneció masiva. La utilización de la GPU se mantuvo en 10-20%.

Reproducimos el problema. La brecha era aún peor en nuestro hardware:

La carga de trabajo es trivial: 7M muestras, 100 características, 2 capas de MLP, tamaño de lote 1M. El modelo procesa un lote en milisegundos. Entonces, ¿dónde va el tiempo?

Qué muestra nvidia-smi

Nada útil. La utilización de la GPU parpadea entre 0% y 30%. El uso de memoria es estable. La temperatura es buena. La GPU está claramente subutilizada, pero nvidia-smi no puede decir por qué.

Qué muestra torch.profiler

El informante intentó el perfilador integrado de PyTorch y “obtuvo keine datos de seguimiento significativos”. Esta es una frustración común: los perfiladores de nivel de aplicación pueden mostrarte qué kernels de CUDA se están ejecutando, pero no pueden ver la programación del lado del host, los eventos de memoria y los eventos de ciclo de vida del proceso que determinan si los datos llegan a la GPU a tiempo.

Qué muestra el seguimiento a nivel de kernel

Ejecutamos la prueba de rendimiento mientras rastreamos tanto las llamadas a la API de CUDA (a través de eBPF uprobes en libcudart.so) como los eventos del kernel de Linux (conmutaciones de contexto del programador, asignaciones de página de memoria, bifurcaciones de proceso) simultáneamente. Los resultados cuentan la historia completa.

Video completo: https://asciinema.org/a/RGwhPeXAPJdhXqxp

En el video, conectamos un LLM de peso abierto (MiniMax-M2.7) a la base de datos de seguimiento a través de MCP (Protocolo de Contexto de Modelo):

El archivo de configuración JSON indica al cliente de MCP dónde encontrar el servidor de Ingero y qué base de datos de seguimiento cargar:

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

Esto le da al LLM acceso directo a los datos de seguimiento a través de 7 herramientas: get_trace_stats, get_causal_chains, get_per_process_breakdown, y otras. El AI puede consultar la base de datos, correlacionar eventos de CUDA con datos de programación del kernel y producir una explicación en lenguaje plano sin análisis manual.

4 cadenas de causalidad de alta gravedad

El motor de cadena de causalidad detectó 4 patrones de alta gravedad, todos con la misma causa raíz:

[ALTA] cudaStreamSync p99=42ms (1,638x p50=25us) - CPU 100% + 1,880 conmutaciones de programador
Línea de tiempo:
[SYSTEM] CPU 100%
[HOST ] 1,880 conmutaciones de contexto (21s fuera de CPU)
[CUDA ] p99=42ms (1,638x p50=25us)
Causa raíz: Trabajadores de DataLoader luchando por CPU, presión de asignación de página masiva

[ALTA] cudaLaunchKernel p99=24.67ms (349x p50) - CPU 100%
Raíz: 34 conmutaciones de programador

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

La cudaStreamSync p99 es 1,638 veces la p50. Eso no es lentitud de la GPU, eso es la GPU esperando datos que nunca llegan a tiempo.

Figura 1: Análisis generado por AI después de ejecutar /investigate. El modelo utilizó las 7 herramientas de MCP de Ingero para consultar la base de datos de seguimiento y produjo una explicación en lenguaje plano con recomendaciones accionables derivadas directamente de los datos de seguimiento.

Desglose por proceso

Aquí es donde se vuelve claro. El proceso principal y sus 4 trabajadores de DataLoader son visibles como entidades separadas:
Proceso principal:

- cudaMemcpyAsync (transferencia de host a dispositivo): avg 301ms, max 2.9 segundos
- cudaStreamSync: p99 = 42ms (normalmente 25us)
- 1,567 conmutaciones de contexto, avg 16ms fuera de CPU, peor parada 5 segundos
- 799,018 asignaciones de página

Trabajador de DataLoader 1: 52,863 conmutaciones de contexto, 89,338 asignaciones de página, peor parada 5s
Trabajador de DataLoader 2: 50,638 conmutaciones de contexto, 83,509 asignaciones de página, peor parada 5s
Trabajador de DataLoader 3: 49,361 conmutaciones de contexto, 70,035 asignaciones de página, peor parada 5s
Trabajador de DataLoader 4: 38,862 conmutaciones de contexto, 56,354 asignaciones de página, peor parada 5s

Total en trabajadores: ~191,000 conmutaciones de contexto y ~299,000 asignaciones de página en 40 segundos.

Qué significa esto

Los trabajadores de DataLoader están haciendo tres cosas costosas que el índice directo evita por completo:

1. Barajado e índice: DataLoader con shuffle=True genera una permutación aleatoria de índices, luego cada trabajador selecciona su parte. Esto requiere acceso aleatorio a memoria a través del tensor de 7M muestras completo – terrible para la localidad de caché y desencadena fallos de página.
2. Agrupación y copia: Cada trabajador reúne muestras dispersas en un tensor de lote contiguo. Esto significa asignar nueva memoria (asignaciones de página), copiar datos de ubicaciones aleatorias (fallos de caché) y serializar el resultado de regreso al proceso principal a través de memoria compartida o una cola.
3. Competencia por CPU: Cuatro trabajadores + el proceso principal en una máquina de 4 vCPU significa preemption constante. Cada trabajador se desprograma 50,000 veces. La peor parada es de 5 segundos – durante la cual la GPU no tiene nada que procesar.

Con índice directo: X[i:i+batch_size] es una vista de tensor contiguo ya en memoria. .to(device) desencadena una transferencia de DMA desde una región contigua única. No hay trabajadores, no hay barajado, no hay agrupación, no hay copias entre procesos, no hay conmutaciones de contexto. La GPU obtiene datos en microsegundos, no en cientos de milisegundos.

La solución

Para cargas de trabajo de GPU en memoria donde el conjunto de datos completo cabe en RAM:

1. No utilice DataLoader. Índice directo con un arreglo de índice prebarajado es más simple y 100 veces más rápido:

   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. Si debe utilizar DataLoader, coincida con num_workers con sus núcleos de CPU reales menos 1. En una máquina de 4 núcleos, num_workers=2 reduce la contienda. Agregue persistent_workers=True para evitar la sobrecarga de bifurcación.
3. Para conjuntos de datos más grandes que la memoria donde DataLoader es necesario, el verdadero cuello de botella se desplaza a la E/S de disco. Utilice prefetch_factor=2 (no más alto – más prefetching significa más presión de memoria) y asegúrese de que su almacenamiento pueda seguir el ritmo.

La imagen más grande

Esta investigación ilustra un patrón que vemos constantemente en cargas de trabajo de GPU: la GPU es rápida, el host es el cuello de botella y las métricas de la GPU no pueden verlo. nvidia-smi informó una utilización baja pero no pudo explicar por qué. torch.profiler capturó kernels de CUDA pero perdió las 200,000 conmutaciones de contexto que ocurrían en el espacio de usuario.

La única forma de ver la imagen completa fue rastrear ambos lados simultáneamente – llamadas a la API de CUDA a nivel de biblioteca y eventos de programación del kernel de Linux – y correlacionarlos por tiempo y ID de proceso. La cadena de causalidad “CPU 100% -> 1,880 conmutaciones de programador -> cudaMemcpyAsync 301ms -> cudaStreamSync 42ms” cuenta la historia completa en una línea. Sin seguimiento entre pilas, esto habría seguido siendo un misterio – como lo fue para el informante original que pasó semanas depurándolo.

Figura 2: Cuando se le preguntó “¿Cuál es el problema principal?”, el modelo identifica la sobre-suscripción de CPU que causa retrasos en la programación del lado del host. cudaLaunchKernel pasó de 73us a 25.8ms (356 veces más lento) porque la CPU no podía programar el lanzamiento a tiempo.

Pruebe usted mismo

Reproduzca la prueba de rendimiento:

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()

# Ruta rápida
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'Directo: {time.time()-start:.3f}s')

# Ruta lenta
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')

Rastree con Ingero para ver qué está sucediendo debajo del capó:

git clone https://github.com/ingero-io/ingero.git
cd ingero && make build
sudo ./bin/ingero trace --duration 60s # en un terminal
python3 benchmark.py # en otro terminal
./bin/ingero explain --since 60s # después de que se complete la prueba de rendimiento

O omita la reproducción y explore nuestros datos de seguimiento directamente. La base de datos de investigación (764KB) está en el repositorio:

# Ver cadenas de causalidad de la investigación
./bin/ingero explain --db investigations/pytorch-dataloader-starvation.db --since 5m

# Desglose por proceso (ver trabajadores de DataLoader vs proceso principal)
./bin/ingero explain --db investigations/pytorch-dataloader-starvation.db --per-process --since 5m

# Conecte su asistente de AI para investigación interactiva
./bin/ingero mcp --db investigations/pytorch-dataloader-starvation.db

Investigue con AI (recomendado). La forma más rápida de analizar el seguimiento es conectar cualquier AI compatible con MCP directamente a la base de datos. No se requiere análisis manual.

Cree un archivo de configuración:

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

Luego conecte su modelo:

# Con Ollama + MiniMax (lo que usamos en el video)
ollmcp -m minimax-m2.7:cloud -j /tmp/ingero-mcp-dataloader.json

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

# Con cualquier cliente compatible con MCP
# Agregue la configuración anterior a la configuración de MCP de su AI

Escriba /investigate para desencadenar el análisis guiado, o haga cualquier pregunta: “¿Qué causó la hambruna de la GPU?” El AI tiene acceso a 7 herramientas que consultan la base de datos de seguimiento directamente.

GitHub: github.com/ingero-io/ingero
Problema original: pytorch/pytorch#154318
Video: https://asciinema.org/a/RGwhPeXAPJdhXqxp

Investigación realizada en TensorDock RTX 4090 (24GB), Ubuntu 22.04, PyTorch 2.10.0+cu128.