PyTorch 2.14에서 pinned host memory를 메모리 스냅샷에 함께 기록하는 방법 | DAKER 커뮤니티

핀드 호스트 스냅샷

PyTorch 메모리 스냅샷을 볼 때 보통은 디바이스 할당부터 확인하게 됩니다. 그런데 실제로는 DataLoader, H2D 복사, CUDA 그래프와 함께 오래 유지되는 핀드 호스트 버퍼가 문제의 원인인 경우도 적지 않습니다. 디바이스 메모리만 봐서는 이런 증가분이 잘 드러나지 않았습니다.

PyTorch 2.14부터는 torch.cuda.memory._record_memory_history에 record_pinned_host_memory=True를 주면, 디바이스 할당뿐 아니라 핀드 호스트 할당도 스냅샷에 남습니다. 같은 워크로드를 기록 전후로 한 번씩만 비교해도 프로세스 수명 동안 스테이징 버퍼가 불어나는지 바로 확인할 수 있습니다.

PyTorch 2.14부터 record_pinned_host_memory=True를 주면 핀드 호스트 할당도 메모리 스냅샷에 기록됩니다.

왜 이 기록이 필요한가

핀드(페이지 고정) 호스트 메모리는 전송 성능이나 그래프 실행 경로 때문에 한 번 확보한 뒤 오래 유지되는 경우가 많습니다. 특히 DataLoader·H2D 복사·CUDA 그래프에 묶인 핀드 버퍼는 프로세스가 살아 있는 동안 계속 남아 있을 수 있습니다.

문제는 기존처럼 디바이스 메모리만 기록한 스냅샷으로는 이런 누수를 확인하기 어려웠다는 점입니다. 이번 옵션을 켜면 host_segments와 host_traces를 통해 호스트 쪽 할당 흐름도 함께 볼 수 있습니다.

재현해 볼 실험 조건

비교 방법은 단순합니다. 먼저 기존처럼 디바이스만 기록한 스냅샷을 한 장 저장합니다. 그다음 같은 구간을 record_pinned_host_memory=True로 다시 기록하면 됩니다. 필요하면 record_cuda=False로 호스트만 남겨 비교할 수도 있습니다.

이후 스냅샷 JSON에서 host_segments·host_traces의 크기, 개수, 타임스탬프를 비교해 보면 됩니다. 웹 비주얼라이저는 아직 호스트를 그리지 않을 수 있으니, 이 경우에는 해당 필드를 직접 읽는 것이 좋습니다.

웹 비주얼라이저가 호스트 메모리를 아직 그리지 않을 수 있으므로 host_segments와 host_traces를 직접 확인하는 편이 정확합니다.

어제 다룬 메모리 풀·그래프 트리 복제와는 다른 축의 이야기입니다. 여기서는 핀드 호스트 관측 자체에만 집중하면 됩니다.

실무에서 먼저 볼 포인트

우선 그래프 캡처에 고정 호스트 주소가 필요해 핀드 버퍼를 오래 붙잡는 서빙·배치 경로부터 확인하는 것이 좋습니다. 이런 경로는 겉으로는 안정적으로 보여도, 프로세스 수명 동안 호스트 메모리가 조금씩 늘어나는 패턴이 숨어 있을 수 있습니다.

다만 PyTorch 할당자 밖에서 cudaHostRegister만 사용한 버퍼는 이 기록에 잡히지 않을 수 있습니다. 따라서 스냅샷에 보이지 않는다고 해서 호스트 쪽 이슈가 없다고 단정하기는 어렵습니다.

재현 노트에는 빌드 번호, 기록 플래그, 스냅샷 파일 경로를 함께 남겨 두는 편이 좋습니다. 같은 워크로드를 다시 비교할 때 차이를 빠르게 좁히는 데 도움이 됩니다.

관련 DAKER 학습

컴파일 구간 use_mem_pool
MPS 할당 버킷
PyTorch 실전 고급 트랙

참고 자료

https://daker.ai/community/post-mufwpmjb-8b4bb524
https://daker.ai/community/post-mubnk8op-b1f171dd
https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track

여러분은 메모리 스냅샷을 볼 때 디바이스보다 먼저 호스트 쪽 증가분을 의심했던 경험이 있었나요?