[llm-compressor] vLLM LLM-Compressor: CUDA OOM 문제 해결을 위한 메모리 핀닝 최적화

PR 링크: vllm-project/llm-compressor#2813 상태: Merged | 변경: +65 / -19

들어가며

최근 llm-compressor 라이브러리를 사용하여 대규모 모델에 AWQ(Activation-aware Weight Quantization)를 적용하는 과정에서, 특히 DDP(Distributed Data Parallel) 환경에서 CUDA OOM(Out of Memory) 문제가 빈번하게 발생하는 이슈가 보고되었습니다. 분석 결과, 기존 코드에서 모든 중간 활성화 값(Intermediate Activations)을 미리 pin_memory()로 고정(pinning)하던 과정이 시스템 메모리 부족을 유발하고 있었습니다. 본 PR은 이 메모리 핀닝 시점을 최적화하여, 필요한 시점에만 메모리를 핀닝함으로써 OOM 문제를 해결하고 효율적인 H2D(Host-to-Device) 전송을 가능하게 했습니다.

코드 분석

1. src/llmcompressor/pipelines/cache.py: 지연 핀닝 구현

기존에는 _offload_value 내부에서 무조건적으로 pin_memory()를 호출했습니다. 이를 제거하고, 대신 IntermediatesCache 클래스에 명시적인 pin_memory() 메서드를 추가하여 필요한 시점에만 호출하도록 변경했습니다.

Before:

# 기존: offload 시점에 무조건 핀닝
if torch.device(offload_device).type == "cpu" and not offloaded.is_pinned():
    offloaded = offloaded.pin_memory()

After:

# 변경: 명시적 호출을 통한 지연 핀닝
def pin_memory(self, batch_index: int, input_names: list[str] | None = None) -> None:
    intermediates = self.batch_intermediates[batch_index]
    for key, intermediate in intermediates.items():
        if input_names is None or key in input_names:
            self._pin_intermediate(intermediate)

2. src/llmcompressor/modifiers/transform/awq/base.py: 최적화된 호출 시점

AWQ의 그리드 서치(Grid Search) 과정에서 반복적으로 활성화 값을 사용하므로, 서치 직전에 해당 배치를 핀닝하도록 수정했습니다.

After:

# AWQ 그리드 서치 직전에 핀닝 수행
cache = self._parent_args_cache[parent_module]
for batch_index in range(len(cache)):
    cache.pin_memory(batch_index)

왜 이게 좋은가

1. CUDA OOM 해결 및 안정성 확보

기존 방식은 모든 데이터를 미리 핀닝하여 시스템 RAM을 과도하게 점유했습니다. 특히 대규모 모델(예: Qwen-30B)에서는 이로 인해 프로세스가 강제 종료되는 문제가 있었습니다. 이번 변경을 통해 메모리 점유율을 최적화하여 대규모 모델에서도 안정적인 학습이 가능해졌습니다.

2. 비동기 데이터 전송(Non-blocking H2D) 활용

pin_memory()를 적절히 사용함으로써, non_blocking=True 옵션을 통한 DMA(Direct Memory Access) 전송이 가능해졌습니다. 이는 GPU 연산과 데이터 전송을 오버랩(Overlap)시켜 전체적인 파이프라인 성능을 개선합니다.

3. 성능 지표

실제 벤치마크 결과, non_blocking=False를 사용했을 때 약 4시간이 소요되던 작업이, 이번 최적화를 통해 37분 수준으로 대폭 단축되었습니다. 비록 단일 GPU 환경에서는 메인 브랜치 대비 미세한 성능 차이가 있을 수 있으나, 대규모 모델에서의 안정성 확보라는 측면에서 매우 가치 있는 트레이드오프입니다.

교훈

pin_memory()는 CPU-GPU 간 데이터 전송 속도를 높이는 강력한 도구이지만, 시스템 RAM을 고정시키므로 무분별하게 사용하면 OOM을 유발합니다. "필요한 데이터만, 필요한 시점에" 핀닝하는 지연 전략이 대규모 모델 최적화의 핵심임을 다시 한번 확인했습니다.

참고 자료

• PyTorch Pin Memory Documentation
• Non-blocking Data Transfer Tutorial

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