[onnxruntime] ONNX Runtime CUDA 데이터 로딩 최적화: Pinned Buffer와 병렬 I/O를 통한 성능 개선

PR 링크: microsoft/onnxruntime#32437 상태: Merged | 변경: +1156 / -44

들어가며

대규모 AI 모델을 서비스할 때, 모델 가중치를 디스크에서 GPU로 로드하는 과정은 전체 InferenceSession 생성 시간의 상당 부분을 차지합니다. 기존 ONNX Runtime의 CUDA Execution Provider(EP)는 mmap을 통해 파일을 매핑한 뒤, 이를 페이지 가능한(pageable) CPU 메모리로 읽고, 다시 CUDA 드라이버가 관리하는 내부 스테이징 메모리를 거쳐 GPU로 전송하는 복잡한 경로를 거쳤습니다. 이 과정에서 발생하는 불필요한 메모리 복사와 동기화 오버헤드는 대형 모델 로딩 시 병목 현상의 주원인이었습니다.

본 PR은 이러한 문제를 해결하기 위해 Pinned Host Buffer를 직접 활용하고, 병렬 CPU 읽기를 도입하여 데이터 전송 경로를 획기적으로 단축했습니다.

코드 분석

1. CUDA Provider 옵션 확장 (cuda_provider_options.h)

사용자가 병렬 읽기 스레드 수를 제어할 수 있도록 새로운 옵션을 추가했습니다.

// Before: 옵션 없음
// After:
struct OrtCUDAProviderOptionsV2 {
  // ...
  size_t external_data_loader_reading_threads = 4; // 0: 비활성화, 1: 동기, 2~64: 병렬
};

2. 커스텀 외부 데이터 로더 구현 (cuda_external_data_loader.cc)

기존의 mmap 기반 경로를 제거하고, 두 개의 64 MiB Pinned Buffer를 순환(ping-pong)하며 사용하는 로직을 구현했습니다. 한 버퍼가 GPU로 cudaMemcpyAsync되는 동안, 다른 버퍼는 CPU에서 디스크 데이터를 읽어 채우는 파이프라이닝 구조입니다.

// 핵심 로직: 병렬 읽기 및 비동기 복사
// CPU는 읽기 작업을 수행하고, CUDA 스트림은 H2D 복사를 수행함
for (const auto& chunk : chunks) {
  // 4개의 스레드가 disjoint range를 병렬로 읽음
  pool.enqueue([&]() { read_file_into_buffer(chunk); });
}
// 모든 읽기가 완료되면 GPU로 비동기 전송
cudaMemcpyAsync(dst, pinned_buffer, size, cudaMemcpyHostToDevice, stream);

왜 이게 좋은가

성능 수치

• Cold Cache 환경: 31.03초 → 26.62초 (14.2% 개선)
• Warm Cache 환경: 27.79초 → 24.42초 (12.1% 개선)
• 병렬 스레드 최적화: 4개의 읽기 스레드 사용 시, 기존 대비 29.8%의 로딩 시간 단축을 기록했습니다.

최적화의 교훈

1. 데이터 경로 단순화: mmap -> pageable -> pinned -> GPU 경로를 disk -> pinned -> GPU로 단축하여 커널 레벨의 페이지 폴트와 드라이버 내부 복사를 최소화했습니다.
2. 파이프라이닝: CPU의 I/O 작업과 GPU의 DMA 전송을 비동기적으로 겹침(overlap)으로써 하드웨어 자원 활용도를 극대화했습니다.
3. 안전한 폴백(Fallback): Pinned 메모리 할당 실패 시 기존의 pageable 경로로 자동 전환되도록 설계하여 시스템 안정성을 확보했습니다.

리뷰어 피드백 반영

리뷰 과정에서 std::async 사용 시 발생할 수 있는 리소스 누수 문제와 32비트 환경에서의 정수 오버플로우 문제를 지적받았습니다. 특히, synchronize_streams()를 모든 예외 경로에 배치하여 비동기 작업 중 오류 발생 시에도 GPU 상태가 꼬이지 않도록 견고하게 수정되었습니다.

참고 자료

• CUDA Runtime API (cudaMemcpyAsync)
• ONNX Runtime CUDA Provider Options

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.