torch.cuda.use_mem_pool을 컴파일 구간 안에 두면 폴백·외부 커널 할당도 같은 풀로 모입니다 | DAKER 커뮤니티

한 줄 요약
PyTorch 2.14부터 torch.cuda.use_mem_pool을 torch.compile로 올린 구간 안에서도 쓸 수 있습니다. 컨텍스트 안에서 나는 할당—폴백·외부 커널 포함—이 지정한 CUDA 메모리 풀로 모입니다. 풀을 밖에만 열어 두면 컴파일 경로의 임시 버퍼가 기본 할당자로 새어 나가던 상황을 오늘 바로 점검할 수 있습니다.

작은 모델 하나를 torch.compile한 뒤, 같은 스텝을 풀 컨텍스트 밖에서 한 번, 안에서 한 번 돌리십시오. 피크 예약·세그먼트 수·스텝 시간만 비교하면 됩니다. 노트에는 PyTorch 버전, 디바이스, 풀 생성 방식, Dynamo·Inductor 로그 한 줄을 남깁니다.

재현해 볼 실험 조건

1. CUDA 메모리 풀을 만들고, 컴파일된 함수 밖에서만 use_mem_pool을 연 채 스텝을 측정합니다.
2. 같은 풀을 컴파일된 함수 안에서 연 뒤, 폴백·외부 커널이 끼는 경로를 한 번 더 측정합니다.
3. 피크 메모리·세그먼트·스텝 시간을 비교하고, 컴파일 구간 할당이 풀에 붙는지 메모리 스냅샷으로 확인합니다.

어제 다룬 torch.switch·CUDA 그래프 while_loop·NVGEMM과는 다른 축입니다. 오늘은 컴파일 구간과 메모리 풀의 결합만 봅니다.

실무에서 바로 볼 포인트

• 커스텀 풀·공유 버퍼를 쓰는 서빙·배치 파이프라인에서 컴파일 임시 텐서가 풀 밖으로 새는지 먼저 확인합니다.
• 풀 수명과 컴파일 아티팩트 수명을 맞춰 두지 않으면 캡처·재실행 때 할당자 상태가 어긋날 수 있습니다.
• 재현 노트에 빌드 번호와 torch.compile 모드를 고정합니다.

관련 DAKER 학습

• PyTorch 실전 고급 트랙
• DAKER 커뮤니티

설명용 생성 이미지입니다. 동작은 로컬에서 풀 안·밖 컴파일 스텝을 같은 입력으로 비교해 확인하십시오.