[TensorRT] PyTorch-TensorRT의 메모리 효율성 개선: 점진적 GPU 메모리 해제를 통한 컴파일 최적화

PR 링크: pytorch/TensorRT#4793 상태: Merged | 변경: +79 / -6

들어가며

PyTorch 모델을 TensorRT로 컴파일하는 과정에서 가장 큰 병목 중 하나는 메모리 사용량입니다. 특히 대규모 모델을 다룰 때, module.to("cpu")를 호출하여 모델을 CPU로 옮기는 과정에서 PyTorch는 기존 GPU 메모리를 즉시 해제하지 않고 캐싱합니다. 이로 인해 컴파일러는 모델의 CPU 복사본과 GPU 복사본을 동시에 메모리에 유지해야 하며, 이는 특히 통합 메모리(Unified Memory) 시스템에서 심각한 메모리 부족(OOM)을 유발합니다. 이번 PR은 이 문제를 해결하기 위해 모듈 이동 중 점진적으로 CUDA 캐시를 비우는 방식을 도입했습니다.

코드 분석

py/torch_tensorrt/dynamo/utils.py

기존의 deallocate_module 함수는 단순히 module.to(CPU_DEVICE)를 호출한 뒤 마지막에 torch.cuda.empty_cache()를 수행했습니다. 이 방식은 이동이 완료될 때까지 모든 가중치가 GPU와 CPU 양쪽에 존재하게 만듭니다.

Before:

def deallocate_module(module: torch.fx.GraphModule) -> None:
    module.to(CPU_DEVICE)
    torch.cuda.empty_cache()
    gc.collect()

After:

def deallocate_module(module: torch.nn.Module, release_every_bytes: int = 1 << 30) -> None:
    moved_since_release = 0
    def to_cpu(t: torch.Tensor) -> torch.Tensor:
        nonlocal moved_since_release
        if t.device.type == "cuda":
            if moved_since_release >= release_every_bytes:
                torch.cuda.empty_cache()
                moved_since_release = 0
            moved_since_release += t.nbytes if t.layout == torch.strided else 0
        return t.to(CPU_DEVICE)

    module._apply(to_cpu)
    gc.collect()
    torch.cuda.empty_cache()

핵심 변경 사항은 module._apply를 사용하여 텐서 단위로 이동을 수행하고, 설정된 임계값(release_every_bytes)만큼 데이터가 이동할 때마다 torch.cuda.empty_cache()를 호출하는 것입니다. 이를 통해 GPU 메모리 점유율이 점진적으로 감소하게 됩니다.

왜 이게 좋은가

이 최적화의 핵심은 "메모리 피크(Peak)의 평탄화"입니다. 3GiB 모델을 기준으로 테스트했을 때, 호스트와 GPU를 합친 메모리 피크 사용량이 1.0W에서 0.33W로 대폭 감소했습니다.

1. 통합 메모리 효율성: Jetson이나 DGX와 같은 시스템에서는 CPU와 GPU가 동일한 DRAM을 공유합니다. 이 최적화는 두 복사본이 동시에 메모리를 점유하는 시간을 최소화하여 OOM 발생 가능성을 획기적으로 낮춥니다.
2. 점진적 해제: release_every_bytes를 통해 메모리 해제 빈도를 조절함으로써, 너무 잦은 empty_cache 호출로 인한 오버헤드와 메모리 부족 사이의 균형을 맞췄습니다.

일반적인 교훈은 "대규모 객체를 이동하거나 변환할 때, 전체 작업이 끝날 때까지 기다리지 말고 중간중간 리소스를 회수하는 전략을 취하라"는 것입니다. 이는 특히 리소스가 제한된 환경에서 대형 모델을 다루는 엔지니어들에게 필수적인 패턴입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.to
• https://pytorch.org/docs/stable/generated/torch.cuda.empty_cache.html

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