vLLM 메모리 누수 디버깅: 힙(Heap) 너머 UCX와 mmap의 미스터리
(mistral.ai)요약:
- 문제 상황: vLLM의 Prefill/Decode 분리(disaggregated) 서빙 환경에서 분당 400MB의 시스템 메모리(RSS) 누수가 발생했으나, 일반적인 Python 프로파일러로는 감지되지 않음.
- 원인 분석: Heaptrack과 pmap으로 힙이 아닌 익명 메모리 매핑(mmap)에서 누수를 확인하고, BPFtrace와 자동화된 GDB 스크립트를 통해 원인을 추적함.
-
범인 규명: 고성능 통신 라이브러리인 UCX가 최적화를 위해
mmap/munmap호출을 가로채고 있었으며, 해제된 메모리를 즉시 반환하지 않고 무한정 큐에 쌓아두는 것이 원인임. -
해결책: 환경 변수
UCX_MEM_MMAP_HOOK_MODE=none을 설정하여 UCX의 메모리 후킹 기능을 비활성화함으로써 문제를 해결함.
상세요약:
1. 미스터리한 메모리 누수
Mistral AI 팀은 vLLM을 사용한 Prefill/Decode 분리 서빙(NIXL 기반) 환경에서 분당 400MB씩 시스템 메모리가 선형적으로 증가하는 현상을 발견했습니다.
- 증상: Python 힙 메모리는 안정적이었으나, 운영체제 수준의 RSS(Resident Set Size)가 계속 증가하다가 결국 OOM(Out of Memory)으로 이어짐.
-
초기 시도 실패:
Memray,Guppy 3같은 Python 도구는 정상으로 표시되었고, 표준GDB는 프로세스를 크래시 시켰으며,Valgrind는 너무 느려 사용할 수 없었습니다.
2. 커널 레벨로의 심층 분석
문제의 원인이 애플리케이션 레벨(Python/C++)이 아닌 더 낮은 레벨에 있음을 직감하고 시스템 도구를 활용했습니다.
-
Heaptrack: 힙 할당(malloc/free)은 안정적이지만 RSS가 증가함을 시각적으로 확인. 이는 누수가
glibc의 힙 관리 밖인 익명 메모리 매핑(anonymous memory mappings) 에서 발생함을 시사했습니다. -
pmap:
/proc/<pid>/maps를 모니터링하여 특정 익명 매핑 영역이 계속 커지고 주소가 변경됨을 확인했습니다. 이는mremap또는munmap후mmap사이클이 반복됨을 의미했습니다. -
BPFtrace:
LD_PRELOAD로도 잡히지 않는(glibc를 우회하는) 시스템 콜을 추적하기 위해 BPFtrace를 사용했습니다. 그 결과mmap호출이 직접적인syscall을 통해 발생하고 있음을 확인했습니다.
3. 범인 검거: 자동화된 GDB 스크립팅
BPFtrace로 문제의 시스템 콜 주소를 확인한 후, GDB를 사용하여 해당 주소(SYS_mmap)에서만 멈추도록 스크립트를 작성했습니다.
사용된 GDB 스크립트 예시:
# mmap 시스템 콜(번호 9)에 조건부 브레이크포인트 설정
break syscall if $rdi == 9
commands
silent
# 시스템 콜 리턴 지점에 임시 브레이크포인트 설정
tbreak *0x00007ffff7d9525d
commands
silent
# 스택 트레이스와 반환된 주소 출력
bt
printf "Syscall returned: rax = 0x%012lx\n", $rax
continue
end
continue
end
이 스택 트레이스를 통해 UCX(Unified Communication X) 라이브러리가 Python의 mmap/munmap 호출을 중간에서 가로채고(intercept) 있다는 결정적 단서를 잡았습니다.
4. 원인: UCX의 과도한 최적화
UCX는 InfiniBand 전송 성능을 높이기 위해 메모리 할당/해제를 후킹합니다.
-
매커니즘:
munmap이 호출될 때 메모리를 운영체제에 바로 반환하지 않고, 나중에 재사용하거나 정리하기 위해 '무효화 큐(invalidation queue)'에 넣어둡니다. -
버그: 기본 설정(
UCX_RCACHE_MAX_UNRELEASED=inf)으로 인해 이 큐가 무한히 커질 수 있었고, vLLM의 특정 사용 패턴에서는 정리 로직(ucp_worker_progress)이 제대로 동작하지 않아 메모리가 계속 쌓이기만 했습니다.
5. 해결 방법
vLLM의 경우 거대한 KVCache 메모리 영역 하나만 등록하면 되므로, UCX의 복잡한 메모리 후킹 기능이 굳이 필요하지 않았습니다.
-
즉각적인 해결: 환경 변수
UCX_MEM_MMAP_HOOK_MODE=none을 설정하여 UCX의 메모리 후킹을 완전히 끄는 것으로 누수를 막았습니다. -
대안:
UCX_RCACHE_MAX_UNRELEASED=1024와 같이 큐의 크기를 제한하여 강제로 정리가 일어나게 할 수도 있습니다. - 조치: vLLM 커뮤니티를 위해 해당 수정 사항이 병합되었으며, 향후 NIXL 릴리스에서도 기본 동작이 개선될 예정입니다.