PyTorch 2.14의 CUDAGraph 다중 메모리 풀 지원, 무엇이 달라졌나 | DAKER 커뮤니티

PyTorch CUDAGraph 다중 메모리 풀 설명 이미지

한 줄 답: PyTorch 2.14 계열에서 CUDA Graph와 메모리 풀을 다룰 때는, 풀 경계 때문에 그래프를 쪼개 두었던 핫 패스를 다시 점검할 여지가 있습니다. 공식 API는 torch.cuda.graph(..., pool=...)로 풀 토큰·MemPool을 넘길 수 있고, 풀을 공유한 그래프는 캡처 순서와 같은 순서로 리플레이해야 합니다.

CUDA Graph를 실무에 붙이다 보면 성능보다 먼저 구조 제약에 부딪히는 경우가 많습니다. 이 글은 DAKER PyTorch 학습 독자를 위해, 다중 풀·공유 풀 관점에서 무엇을 재현·측정하면 좋은지 정리합니다. 구체 커밋·벤치 숫자는 워크로드마다 다르므로 여기서는 측정 항목만 고정합니다.

무엇이 바뀌었다고 보나

이전에는 서로 다른 메모리 풀에서 나온 할당이 섞이는 구간에서 캡처가 나뉘거나, 그래프를 풀 경계마다 분리하는 쪽으로 설계를 맞추는 경우가 있었습니다. PyTorch 문서의 torch.cuda.graph는 pool 인자로 다른 그래프의 pool() 토큰이나 graph_pool_handle(), MemPool을 넘겨 “같은 풀을 쓸 수 있다”는 힌트를 줍니다.

핵심은 기능을 한 번에 켜는 것이 아니라, 기존에 풀 경계 때문에 분할해 두었던 핫 패스를 다시 볼 수 있게 됐다는 점입니다. 그래프 수가 많아 리플레이 오버헤드가 커졌던 경로라면 특히 먼저 확인해 볼 만합니다.

풀을 공유한 CUDA Graph는 캡처한 순서와 같은 순서로 리플레이해야 합니다. 순서를 바꾸면 메모리 손상이 날 수 있습니다.

먼저 재현할 실험 조건

가장 좋은 출발점은 예전 제약이 드러나던 조건을 그대로 재현하는 것입니다. 서로 다른 메모리 풀에서 텐서를 만드는 짧은 CUDA 구간을 준비한 뒤, 풀별로 그래프를 나눴을 때의 기준선을 먼저 측정합니다.

그다음 같은 구간을 하나의 torch.cuda.CUDAGraph 경로(또는 공유 pool 인자)로 캡처·리플레이해, 풀 경계 오류 없이 끝나는지 비교합니다. 기록할 항목은 풀 개수, 캡처 성공 여부, 리플레이 시간, OOM·격리 오류입니다. 평균만 쓰지 말고 중앙값과 실패 사례를 남기세요.

캡처 중 새 풀 할당이 섞이는 경우와, 미리 워밍업한 뒤 캡처하는 경우는 기록을 분리합니다. 같은 워크로드라도 해석이 달라질 수 있습니다.

실무에서 바로 볼 포인트

모든 경로를 한꺼번에 바꾸기보다, 그래프 개수가 늘어 리플레이 오버헤드가 커진 핫 패스부터 확인하는 편이 현실적입니다. 풀의 수명·파괴 순서를 리플레이 훅·콜백과 한 번에 검증하면 원인 분리가 어렵습니다. 먼저 캡처 성공 여부만 떼어 확인한 뒤, 수명 관리와 주변 로직을 붙이세요.

torch.compile 구간의 메모리 풀 실험과 결과가 겹치면 항목을 섞지 않습니다. CUDAGraph·풀 결과만 따로 표에 정리해 두면 비교가 쉽습니다. 워밍업 연장, TunableOp, 핀드 호스트 같은 인접 주제는 변수에서 잠시 분리하는 편이 결과를 읽기 쉽습니다.

오늘 바로 할 측정

  1. 풀별로 쪼갠 기존 그래프 경로의 리플레이 시간 기준선을 잽니다.
  2. 동일 구간을 공유 pool 인자로 다시 캡처합니다.
  3. 리플레이 순서를 캡처 순서와 같게 유지한 채 성공·시간·OOM을 기록합니다.
  4. 의도적으로 순서를 바꾼 실패 케이스 한 건을 문서에 “하지 말 것”으로 남깁니다.

관련 DAKER 글과 주제를 나누는 법

워밍업 연장(cudagraph_mark_warmup_incomplete), torch.cuda.use_mem_pool와 torch.compile, TunableOp·Grouped GEMM은 각각 다른 축입니다. 오늘 노트 제목에 “multiple pools / shared pool”만 남기고, 다른 실험 로그는 파일명을 분리하세요. 같은 표에 섞으면 어떤 변화가 리플레이 시간을 줄였는지 읽기 어렵습니다.

실패를 문서로 남기는 이유

“공유 풀 + 잘못된 리플레이 순서” 실패를 한 번만 재현해 스크린샷·에러 문자열을 남기면, 팀원이 같은 함정에 다시 빠지지 않습니다. NVIDIA 가이드가 경고하는 지점이 바로 이 순서 제약입니다. 성공 케이스만 올리면 다음날 누군가 순서를 바꿔도 원인을 모릅니다.

측정표 템플릿

열을 이렇게 고정하세요. 날짜 | 풀 개수 | 공유 pool 사용 여부 | 캡처 성공 | 리플레이 순서 준수 | 리플레이 ms(중앙값) | OOM/에러 문자열 | 비고. 행은 “분할 그래프 기준선”과 “공유 pool 시도” 두 줄만 있어도 다음 실험이 빨라집니다. DAKER PyTorch 고급 트랙 실습 노트에 이 표를 그대로 붙여 넣으면 됩니다.

마치며

PyTorch 2.14 문서 기준으로 CUDA Graph의 pool 인자는 다중 그래프·다중 풀 설계를 다시 열어 줍니다. DAKER 학습 트랙의 인접 글(워밍업 연장, use_mem_pool)과 주제를 섞지 말고, 오늘 측정은 “풀 경계로 쪼갠 그래프를 합칠 수 있는가” 한 가지만 답하면 됩니다.

출처