PyTorch 2.14에서 torch.cuda.use_mem_pool과 torch.compile을 함께 볼 때 확인할 점 | DAKER 커뮤니티
PyTorch 2.14부터 torch.cuda.use_mem_pool을 torch.compile로 올린 구간 안에서도 사용할 수 있습니다. 이 변화는 커스텀 메모리 풀을 쓰는 환경에서 특히 중요합니다. 그동안은 풀을 바깥에서만 열어 두면 컴파일 경로의 임시 버퍼가 기본 할당자로 빠지는지 따로 의심해 봐야 했는데, 이제는 그 경계를 더 분명하게 점검할 수 있습니다.
지금 확인해 볼 만한 이유도 분명합니다. 컨텍스트 안에서 발생하는 할당이 폴백이나 외부 커널까지 포함해 같은 CUDA 메모리 풀로 모인다면, 서빙이나 배치 파이프라인에서 메모리 거동을 더 예측하기 쉬워집니다. 반대로 풀 수명과 컴파일 아티팩트 수명이 어긋나면 캡처나 재실행 시점에 할당자 상태가 달라질 수 있어, 초기에 비교해 두는 것이 좋습니다.
PyTorch 2.14부터 torch.cuda.use_mem_pool을 torch.compile 구간 안에서도 쓸 수 있고, 컨텍스트 안에서 나는 할당은 폴백·외부 커널을 포함해 지정한 CUDA 메모리 풀로 모입니다.
무엇이 달라졌는가
핵심은 torch.cuda.use_mem_pool의 적용 범위를 컴파일된 경로 안까지 가져갈 수 있다는 점입니다. 이때 컨텍스트 안에서 발생하는 할당은 지정한 풀에 붙습니다. 따라서 풀을 컴파일된 함수 바깥에서만 열어 둔 경우와, 함수 안에서 직접 연 경우를 비교하면 컴파일 경로의 임시 버퍼가 어디에 잡히는지 더 분명하게 볼 수 있습니다.
이 변화는 어제 다룬 torch.switch, CUDA 그래프 while_loop, NVGEMM과는 다른 축의 이야기입니다. 여기서는 컴파일 구간과 메모리 풀의 결합만 따로 보는 것이 좋습니다.
재현해 볼 실험 조건
작은 모델 하나를 torch.compile한 뒤, 같은 스텝을 풀 컨텍스트 밖에서 한 번, 안에서 한 번 돌려 보면 됩니다. 비교할 항목은 피크 예약 메모리, 세그먼트 수, 스텝 시간입니다. 여기에 메모리 스냅샷을 함께 보면 컴파일 구간의 할당이 실제로 풀에 붙는지도 확인할 수 있습니다.
- CUDA 메모리 풀을 만들고, 컴파일된 함수 밖에서만
use_mem_pool을 연 채 스텝을 측정합니다. - 같은 풀을 컴파일된 함수 안에서 연 뒤, 폴백·외부 커널이 끼는 경로를 한 번 더 측정합니다.
- 피크 메모리, 세그먼트, 스텝 시간을 비교하고, 컴파일 구간 할당이 풀에 붙는지 메모리 스냅샷으로 확인합니다.
실무에서 먼저 볼 포인트
커스텀 풀이나 공유 버퍼를 쓰는 서빙·배치 파이프라인이라면, 컴파일 과정에서 생기는 임시 텐서가 풀 밖으로 새는지부터 확인하는 것이 좋습니다. 이 지점이 정리되지 않으면 메모리 사용량을 읽을 때 실제 병목을 놓치기 쉽습니다.
또 하나는 수명 관리입니다. 풀의 수명과 컴파일 아티팩트의 수명을 맞춰 두지 않으면 캡처와 재실행 과정에서 할당자 상태가 어긋날 수 있습니다. 재현 노트에는 PyTorch 버전, 디바이스, 풀 생성 방식, torch.compile 모드, Dynamo·Inductor 로그 한 줄, 빌드 번호를 함께 남겨 두면 비교가 쉬워집니다.
풀을 밖에만 열어 둔 상태와 컴파일 구간 안에서 연 상태를 나란히 비교하면, 임시 버퍼가 기본 할당자로 새는지 빠르게 점검할 수 있습니다.
관련 자료
참고 자료
https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/community
여러분 환경에서는 풀을 컴파일 구간 안에 둘 때와 밖에 둘 때 어떤 차이가 보였는지 궁금합니다.