124배 느림: PyTorch DataLoader가 커널 수준에서 실제로 하는 일

TL;DR

PyTorch의 DataLoader는 메모리 내 GPU 워크로드에서 직접 텐서 인덱싱보다 50-124배 느릴 수 있습니다. 우리는 RTX 4090에서 실제 PyTorch 이슈를 재현하고 모든 CUDA API 호출과 Linux 커널 이벤트를 추적하여 근본 원인을 찾았습니다. GPU가 느린 것이 아니라 굶주리고 있었습니다. DataLoader 워커들이 40초 동안 200,000회의 CPU 컨텍스트 스위치와 300,000회의 페이지 할당을 생성하여, 마이크로초 단위로 걸려야 할 데이터 전송마다 평균 301ms 동안 GPU가 대기하게 했습니다.

문제점

한 PyTorch 사용자가 보고한 바에 따르면, 간단한 MLP 추론 워크로드에서 DataLoader가 직접 텐서 인덱싱보다 7-22배 느렸습니다. 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 호출(libcudart.so에 대한 eBPF uprobes를 통해)과 Linux 커널 이벤트(스케줄러 컨텍스트 스위치, 메모리 페이지 할당, 프로세스 포크)를 동시에 추적하면서 벤치마크를 실행했습니다. 결과는 완전한 이야기를 말해줍니다.

전체 비디오 설명: https://asciinema.org/a/RGwhPeXAPJdhXqxp

비디오에서 우리는 오픈 웨이트 LLM(MiniMax-M2.7)을 MCP(Model Context Protocol)를 통해 트레이스 데이터베이스에 연결했습니다:

JSON 구성은 MCP 클라이언트에게 Ingero 서버를 찾을 위치와 로드할 트레이스 데이터베이스를 알려줍니다:

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

이를 통해 LLM은 get_trace_stats, get_causal_chains, get_per_process_breakdown 등 7개의 도구를 통해 트레이스 데이터에 직접 접근할 수 있습니다. AI는 데이터베이스를 쿼리하고, CUDA 이벤트를 커널 스케줄링 데이터와 연관시키며, 수동 분석 없이도 일반 언어로 진단을 생성할 수 있습니다.

4개의 HIGH 심각도 인과 관계 체인

인과 관계 체인 엔진은 모두 동일한 근본 원인을 가진 4개의 높은 심각도 패턴을 감지했습니다:

[HIGH] cudaStreamSync p99=42ms (1,638x p50=25us) - CPU 100% + 1,880 sched_switch 이벤트
타임라인:
[SYSTEM] CPU 100%
[HOST ] 1,880 컨텍스트 스위치 (21초 off-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: /investigate 실행 후 AI가 생성한 분석. 모델은 Ingero의 7개 MCP 도구를 사용하여 트레이스 데이터베이스를 쿼리하고, 트레이스 데이터에서 직접 도출된 실행 가능한 권장 사항과 함께 일반 언어 설명을 생성했습니다.

프로세스별 분석

여기서 명확해집니다. 메인 프로세스와 그 4개의 DataLoader 워커들은 별개의 개체로 보입니다:

메인 프로세스:

- cudaMemcpyAsync (호스트-장치 전송): 평균 301ms, 최대 2.9초
- cudaStreamSync: p99 = 42ms (보통 25us)
- 1,567 컨텍스트 스위치, 평균 16ms off-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

워커 전체 합계: 40초 동안 약 191,000회 컨텍스트 스위치 및 약 299,000회 페이지 할당.

이것이 의미하는 바

DataLoader 워커들은 직접 인덱싱이 완전히 피하는 세 가지 비용이 큰 작업을 수행하고 있습니다:

1. 셔플링 및 인덱싱: shuffle=True인 DataLoader는 인덱스의 무작위 순열을 생성한 다음 각 워커가 자신의 청크를 선택합니다. 이는 전체 7백만 샘플 텐서에 걸친 무작위 메모리 접근을 필요로 하며, 캐시 지역성에 끔찍하고 페이지 폴트를 유발합니다.
2. 병합 및 복사: 각 워커는 흩어진 샘플들을 연속된 배치 텐서로 모읍니다. 이는 새로운 메모리 할당(페이지 할당), 무작위 위치에서 데이터 복사(캐시 미스), 그리고 결과를 공유 메모리나 큐를 통해 메인 프로세스로 다시 직렬화하는 것을 의미합니다.
3. CPU 경쟁: 4개의 워커 + 4-vCPU 머신의 메인 프로세스는 지속적인 선점을 의미합니다. 각 워커는 50,000번 스케줄 해제됩니다. 최악의 경우 정지는 5초 동안 지속되며, 그 동안 GPU는 처리할 것이 없습니다.

직접 인덱싱의 경우: X[i:i+batch_size]는 이미 메모리에 있는 연속된 텐서의 제로-카피 뷰입니다. .to(device)는 단일 연속 영역에서 하나의 DMA 전송을 트리거합니다. 워커 없음, 셔플링 없음, 병합 없음, 프로세스 간 복사 없음, 컨텍스트 스위치 없음. GPU는 수백 밀리초가 아닌 마이크로초 단위로 데이터를 받습니다.

해결책

전체 데이터셋이 RAM에 맞는 메모리 내 GPU 워크로드의 경우:

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를 추가하세요.
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”라는 인과 관계 체인은 한 줄로 완전한 이야기를 말해줍니다. 크로스-스택 트레이싱 없이는 이것은 수수께끼로 남았을 것입니다. 원래 보고자가 몇 주 동안 디버깅했던 것처럼 말입니다.

그림 2: “핵심 문제는 무엇인가?”라고 물었을 때, 모델은 호스트 측 스케줄링 지연을 유발하는 CPU 과다 할당을 식별합니다. cudaLaunchKernel은 CPU가 제시간에 런치를 스케줄할 수 없어서 73us에서 25.8ms(356배 느림)로 증가했습니다.

직접 시도해 보세요

벤치마크 재현:
<div style