[TensorRT] PyTorch-TensorRT, 엔진 빌드 후 메모리 누수 해결: `malloc_trim` 도입으로 최적화
PR 링크: pytorch/TensorRT#4792 상태: Merged | 변경: +88 / -17
들어가며
PyTorch 모델을 TensorRT 엔진으로 변환하는 과정은 상당한 컴퓨팅 자원을 요구합니다. 특히, 모델의 크기가 커질수록 컴파일 과정에서 사용되는 CPU 메모리 양도 기하급수적으로 늘어날 수 있습니다. 이 과정에서 발생하는 메모리 누수는 시스템 불안정으로 이어지거나, 더 큰 모델을 다룰 수 없게 만드는 병목 현상을 초래할 수 있습니다.
이번 PR(#4797)은 PyTorch-TensorRT에서 TensorRT 엔진을 빌드한 후에도 해제되지 않고 프로세스에 남아있는 CPU 메모리 문제를 해결하는 데 초점을 맞추고 있습니다. 특히 리눅스 환경에서 glibc의 malloc_trim 함수를 활용하여 빌드 과정에서 사용되었던 메모리를 운영체제에 적극적으로 반환함으로써, 컴파일 시 최대 메모리 사용량을 줄이고 컴파일 후에도 시스템에 더 적은 메모리를 점유하도록 개선했습니다.
이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 변경이 왜 성능 향상으로 이어지는지, 그리고 이 최적화가 주는 일반적인 교훈은 무엇인지 살펴보겠습니다.
코드 분석
이번 PR의 핵심 변경 사항은 py/torch_tensorrt/dynamo/utils.py 파일의 메모리 관리 로직과 관련 문서 파일들의 수정입니다.
1. py/torch_tensorrt/dynamo/utils.py - 메모리 해제 로직 개선
가장 중요한 변경은 release_host_and_device_memory 함수와 새로 도입된 trim_host_heap 함수입니다. 기존에는 TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM 환경 변수가 1로 설정된 경우에만 libc.so.6의 malloc_trim을 호출하여 메모리를 해제하려 시도했습니다. 하지만 이 PR에서는 이 로직이 변경되었습니다.
Before:
--- a/py/torch_tensorrt/dynamo/utils.py
+++ b/py/torch_tensorrt/dynamo/utils.py
@@ -1038,16 +1039,38 @@ def release_host_and_device_memory() -> None:
torch.cuda.ipc_collect()
torch.cuda.synchronize()
- if (
- platform.system() == "Linux"
- and os.environ.get("TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM", "0") == "1"
- ):
- try:
- libc = ctypes.CDLL("libc.so.6")
- if libc.malloc_trim(0) != 1:
- logger.warning("Failed to release CPU memory.")
- except Exception:
- logger.warning("Failed to release CPU memory.")
+ trim_host_heap()
+```
**After:**
```diff
@@ -1038,16 +1039,38 @@ def release_host_and_device_memory() -> None:
torch.cuda.ipc_collect()
torch.cuda.synchronize()
- if (
- platform.system() == "Linux"
- and os.environ.get("TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM", "0") == "1"
- ):
- try:
- libc = ctypes.CDLL("libc.so.6")
- if libc.malloc_trim(0) != 1:
- logger.warning("Failed to release CPU memory.")
- except Exception:
- logger.warning("Failed to release CPU memory.")
+ trim_host_heap()
+```
이 변경은 `release_host_and_device_memory` 함수 내에서 직접 `malloc_trim` 호출 로직을 제거하고, 대신 `trim_host_heap()` 함수를 호출하도록 변경한 것입니다. `trim_host_heap` 함수는 다음과 같이 구현되었습니다.
```diff
@@ -4,6 +4,7 @@
from __future__ import annotations
import ctypes
+import functools
import gc
import importlib.util
import logging
@@ -1038,16 +1039,38 @@ def release_host_and_device_memory() -> None:
torch.cuda.ipc_collect()
torch.cuda.synchronize()
- if (
- platform.system() == "Linux"
- and os.environ.get("TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM", "0") == "1"
- ):
- try:
- libc = ctypes.CDLL("libc.so.6")
- if libc.malloc_trim(0) != 1:
- logger.warning("Failed to release CPU memory.")
- except Exception:
- logger.warning("Failed to release CPU memory.")
+ trim_host_heap()
+```
```diff
@@ -1038,16 +1039,38 @@ def release_host_and_device_memory() -> None:
torch.cuda.ipc_collect()
torch.cuda.synchronize()
- if (
- platform.system() == "Linux"
- and os.environ.get("TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM", "0") == "1"
- ):
- try:
- libc = ctypes.CDLL("libc.so.6")
- if libc.malloc_trim(0) != 1:
- logger.warning("Failed to release CPU memory.")
- except Exception:
- logger.warning("Failed to release CPU memory.")
+ trim_host_heap()
+```
```diff
+@functools.lru_cache(maxsize=1)
+def _glibc_malloc_trim() -> Optional[Callable[[int], int]]:
+ """glibc's ``malloc_trim``, or ``None`` where the C library doesn't provide it."""
+ if platform.system() != "Linux":
+ return None
+ try:
+ malloc_trim = ctypes.CDLL("libc.so.6").malloc_trim
+ except (OSError, AttributeError):
+ return None
+ malloc_trim.argtypes = [ctypes.c_size_t]
+ malloc_trim.restype = ctypes.c_int
+ return malloc_trim
+
+
def trim_host_heap() -> None:
+ """Return freed heap memory to the operating system.
+
+ Building a TensorRT engine frees on the order of gigabytes of host allocations that
+ glibc keeps in its arenas instead of returning to the OS, so the process holds them
+ for the rest of its life. On unified memory systems that memory is also unavailable
+ to the GPU. ``malloc_trim`` hands it back. Set
+ ``TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=0`` to skip this.
+ """
+ if os.environ.get("TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM", "1") == "0":
+ return
+ malloc_trim = _glibc_malloc_trim()
+ if malloc_trim is not None:
+ # The return value only says whether there was anything to release.
+ malloc_trim(0)
trim_host_heap 함수는 다음과 같은 중요한 변경 사항을 포함합니다:
- 기본 동작 변경: 이전에는
TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM환경 변수가 명시적으로1로 설정되어야malloc_trim이 호출되었습니다. 이제는 이 환경 변수가0으로 설정되지 않는 한 (즉, 기본적으로)malloc_trim이 호출됩니다. 리뷰어의 코멘트(Probably no need to change the API, we can change the default just so other peoples flows dont break)에서 알 수 있듯이, 이는 기존 사용자들의 흐름을 깨지 않으면서도 기본적으로 메모리 누수 문제를 해결하려는 의도입니다. _glibc_malloc_trim헬퍼 함수:malloc_trim함수를 로드하고 캐싱하는 로직이_glibc_malloc_trim이라는 별도의 함수로 분리되었습니다. 이는functools.lru_cache를 사용하여 라이브러리 로딩 비용을 줄이고,malloc_trim함수가 없을 경우(예: glibc가 아닌 C 라이브러리 사용 시)None을 반환하도록 하여 안정성을 높입니다.- 명확한 설명:
trim_host_heap함수의 docstring은 왜malloc_trim이 필요한지, 그리고 이것이 통합 메모리 시스템(unified memory systems)에서 왜 중요한지에 대해 명확하게 설명하고 있습니다.
2. 문서 변경 (docsrc/debugging/troubleshooting.rst, docsrc/tutorials/resource_memory/resource_management.rst)
메모리 관리 로직 변경에 따라 관련 문서도 업데이트되었습니다.
Before (docsrc/debugging/troubleshooting.rst):
--- a/docsrc/debugging/troubleshooting.rst
+++ b/docsrc/debugging/troubleshooting.rst
@@ -112,7 +112,8 @@ Memory Errors During Compilation
TRT compilation can use up to 5× the model size in CPU memory.
- * Set ``TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=1`` to reduce to ~3× model size.
+ * Make sure ``TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM`` is not set to ``0``; trimming
+ the builder's freed memory is on by default on Linux.
* Disable ``offload_module_to_cpu`` (``False``) to drop another 1× copy.
----
After (docsrc/tutorials/resource_memory/resource_management.rst):
--- a/docsrc/tutorials/resource_memory/resource_management.rst
+++ b/docsrc/tutorials/resource_memory/resource_management.rst
@@ -28,16 +28,16 @@ This can exceed system limits when compiling large models.
**Ways to lower CPU memory usage:**
-1. **Enable memory trimming**
+1. **Keep memory trimming enabled**
- Set the following environment variable:
+ On Linux, Torch-TensorRT returns the memory the TensorRT builder frees to the
+ operating system after each engine build (glibc ``malloc_trim``). Without it the
+ process keeps on the order of gigabytes of freed builder memory for the rest of its
+ life. Trimming is on by default; to turn it off, set:
.. code-block:: bash
- export TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=1
-
- This reduces approximately **2x** of redundant model copies, limiting
- total CPU memory usage to up to **3x** the model size.
+ export TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=0
문서에서는 TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=1을 설정하라는 내용에서, 이제는 이 변수를 0으로 설정하지 않으면 기본적으로 메모리 트림이 활성화된다는 점을 강조합니다. 또한, malloc_trim이 없을 경우 발생하는 경고 메시지를 제거하고, 이 기능이 기본적으로 활성화되는 이유를 명확히 설명합니다.
3. 테스트 케이스 추가 (tests/py/dynamo/runtime/test_compile_memory.py)
새로운 테스트 파일 test_compile_memory.py가 추가되어 trim_host_heap 함수의 동작을 검증합니다.
test_trims_by_default: 환경 변수가 설정되지 않았을 때malloc_trim이 기본적으로 호출되는지 확인합니다.test_env_var_disables_trimming:TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=0환경 변수가 설정되었을 때malloc_trim이 호출되지 않는지 확인합니다.test_missing_malloc_trim_is_a_no_op:_glibc_malloc_trim이None을 반환할 때 아무런 동작도 하지 않는지 확인합니다.test_finds_glibc_malloc_trim: 리눅스 환경에서_glibc_malloc_trim이 제대로malloc_trim함수를 찾는지 확인합니다.
이 테스트들은 trim_host_heap 함수의 변경된 동작과 안정성을 보장하는 데 중요한 역할을 합니다.
왜 이게 좋은가?
이 PR의 가장 큰 장점은 메모리 효율성 향상입니다. PR 설명에 따르면, RTX 5090에서 3 GiB FP16 모델을 컴파일할 때 다음과 같은 개선 효과가 측정되었습니다:
- Peak compile memory (host + GPU) drops by 1.5 GiB: 컴파일 중 최대 호스트 + GPU 메모리 사용량이 1.5 GiB 감소했습니다.
- The process keeps 1.3 GiB less after compiling: 컴파일 완료 후 프로세스가 점유하는 메모리가 1.3 GiB 감소했습니다.
이는 특히 통합 메모리 시스템(Unified Memory Systems)을 사용하는 Jetson, DGX Spark와 같은 환경에서 매우 중요합니다. 이러한 시스템에서는 호스트 메모리와 GPU 메모리가 동일한 풀에서 할당되므로, 호스트 메모리 사용량 증가는 곧 GPU가 사용할 수 있는 메모리 감소로 직결됩니다. malloc_trim을 통해 불필요하게 점유되던 메모리를 OS에 반환함으로써, GPU가 더 많은 메모리를 활용할 수 있게 되어 잠재적으로 더 큰 모델을 컴파일하거나 더 효율적인 엔진을 생성할 수 있게 됩니다.
또한, 이 PR은 기존 사용자 경험을 해치지 않으면서 기본 설정을 개선했다는 점에서 좋은 평가를 받을 수 있습니다. 리뷰 코멘트에서도 언급되었듯이, 기본값 변경은 많은 사용자가 별도의 설정 없이도 메모리 최적화의 이점을 누릴 수 있게 합니다. TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM=0으로 옵트아웃하는 방식은 사용자가 명시적으로 메모리 트림을 비활성화해야 함을 의미하며, 이는 메모리 누수 방지를 기본값으로 설정하는 강력한 의지를 보여줍니다.
일반적인 교훈
- 메모리 관리의 중요성: 딥러닝 모델 컴파일 및 추론 과정에서 메모리 관리는 성능과 안정성에 직결되는 핵심 요소입니다. 특히 라이브러리 내부에서 발생하는 메모리 할당 및 해제 패턴을 이해하고 최적화하는 것이 중요합니다.
- 운영체제 API 활용: glibc의
malloc_trim과 같은 운영체제 수준의 메모리 관리 도구를 적절히 활용하면, 애플리케이션 레벨에서 직접 관리하기 어려운 메모리 누수 문제를 효과적으로 해결할 수 있습니다. - 기본값의 힘: 좋은 기본값 설정은 사용자의 편의성을 높이고, 많은 경우 성능을 자동으로 개선합니다. 사용자가 명시적으로 비활성화하지 않는 한 최적화된 동작을 제공하는 것이 바람직합니다.
- 테스트의 중요성: 새로운 기능이나 변경 사항이 도입될 때는 반드시 관련 테스트 케이스를 추가하여 기능의 정확성과 안정성을 보장해야 합니다. 이 PR에서는
test_compile_memory.py를 통해malloc_trim로직의 올바른 동작을 검증했습니다.
결론
이번 PR은 PyTorch-TensorRT의 컴파일 과정에서 발생하는 CPU 메모리 누수 문제를 해결하여 시스템 효율성을 크게 향상시킨 중요한 개선입니다. malloc_trim 함수의 도입과 기본값 변경을 통해 사용자는 별도의 설정 없이도 더 적은 메모리로 모델을 컴파일할 수 있게 되었습니다. 이는 특히 메모리 제약이 있는 환경에서 딥러닝 모델을 최적화하고 배포하는 데 큰 도움이 될 것입니다.
References
- glibc malloc_trim documentation (Note: Direct link to
malloc_trimusage might be within broader memory allocator documentation) - torch.export
- torch_tensorrt.dynamo.compile
참고 자료
- https://www.gnu.org/software/libc/manual/html_node/Hooks-for-a-Memory-Allocator.html#Hooks-for-a-Memory-Allocator
- https://pytorch.org/docs/stable/torch.export.html
- https://docs.nvidia.com/deeplearning/tensorrt/pytorch-api/pytorch_tensorrt.dynamo.html#torch_tensorrt.dynamo.compile
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
PR Analysis 의 다른글
- 이전글 [TensorRT] PyTorch-TensorRT의 메모리 효율성 개선: 점진적 GPU 메모리 해제를 통한 컴파일 최적화
- 현재글 : [TensorRT] PyTorch-TensorRT, 엔진 빌드 후 메모리 누수 해결: `malloc_trim` 도입으로 최적화
- 다음글 [flashinfer] Blackwell 아키텍처를 위한 FlashInfer의 FP8/BF16 혼합 정밀도 Attention 최적화 분석
댓글