[pytorch] PyTorch CUDA 메모리 스냅샷 최적화 — 트레이스 선택적 포함

PR 링크: pytorch/pytorch#172672 상태: Merged | 변경: +207 / -43

들어가며

PyTorch의 torch.cuda.memory_snapshot()은 CUDA allocator의 전체 상태를 반환하는 디버깅 도구다. 그러나 이 함수는 매 호출마다 모든 trace entry의 timestamp를 TSC에서 마이크로초로 변환하는 작업을 수행한다. 메모리 기록이 활성화된 상태에서 수만 건의 trace가 쌓이면, 단순히 segment 정보만 필요한 경우에도 수십 밀리초의 오버헤드가 발생한다. 이 PR은 include_traces 파라미터를 추가하여 trace 수집을 선택적으로 건너뛸 수 있게 한다.

핵심 코드 분석

1. C++ 핵심 — NativeCachingAllocator::snapshot

Before:

SnapshotInfo snapshot(MempoolId_t mempool_id) override {
    auto tsc_to_ns = clock_converter.makeConverter();
    auto tsc_to_us = [=](approx_time_t t_approx) {
        return tsc_to_ns(t_approx) / 1000;
    };
    SnapshotInfo result;
    annotation_buffer.getEntries(result.external_annotations);
    for (auto& ae : result.external_annotations) {
        ae.time_.t_ = tsc_to_us(ae.time_.approx_t_);
    }
    for (auto& da : device_allocator) {
        result.device_traces.emplace_back(da->trace(tsc_to_us));
        auto snap = da->snapshot(mempool_id);
        result.segments.insert(result.segments.end(), snap.begin(), snap.end());
    }

After:

SnapshotInfo snapshot(MempoolId_t mempool_id, bool include_traces) override {
    SnapshotInfo result;
    if (include_traces) {
        auto tsc_to_ns = clock_converter.makeConverter();
        auto tsc_to_us = [=](approx_time_t t_approx) {
            return tsc_to_ns(t_approx) / 1000;
        };
        annotation_buffer.getEntries(result.external_annotations);
        for (auto& ae : result.external_annotations) {
            ae.time_.t_ = tsc_to_us(ae.time_.approx_t_);
        }
        for (auto& da : device_allocator) {
            result.device_traces.emplace_back(da->trace(tsc_to_us));
            auto snap = da->snapshot(mempool_id);
            result.segments.insert(result.segments.end(), snap.begin(), snap.end());
        }
    } else {
        // Fast path: skip traces and annotations entirely
        for (auto& da : device_allocator) {
            auto snap = da->snapshot(mempool_id);
            result.segments.insert(result.segments.end(), snap.begin(), snap.end());
        }
    }

include_traces=false일 때 clock_converter.makeConverter() 호출 자체를 건너뛴다. 이 converter 생성이 TSC calibration을 포함하므로 비용이 크다. Trace iteration과 annotation 수집도 완전히 생략된다.

2. Python API 변경

def memory_snapshot(mempool_id=None, include_traces=True):
    if mempool_id is None:
        return torch._C._cuda_memorySnapshot((0, 0, include_traces))["segments"]
    else:
        return torch._C._cuda_memorySnapshot(
            (mempool_id[0], mempool_id[1], include_traces)
        )["segments"]

기존 2-tuple (id1, id2)에 bool을 추가한 3-tuple (id1, id2, include_traces)로 C++ 바인딩에 전달한다. MemPool.snapshot()에도 동일한 파라미터가 추가되었다.

3. C++ 바인딩 호환성 처리

if (size == 2) {
    // (int, int) - mempool_id only
    mempool_id = c10::cuda::MempoolId_t(
        THPUtils_unpackLong(id1), THPUtils_unpackLong(id2));
} else if (size == 3) {
    // (int, int, bool) - mempool_id + include_traces
    mempool_id = c10::cuda::MempoolId_t(
        THPUtils_unpackLong(id1), THPUtils_unpackLong(id2));
    include_traces = (traces.get() == Py_True);
}

기존 2-tuple 호출과의 하위 호환성을 유지한다.

왜 이게 좋은가

테스트 코드에 포함된 벤치마크에 따르면, 1000개의 tensor를 할당/해제한 후:

• Mempool snapshot: trace 포함 대비 trace 미포함 시 수배~수십 배 빨라짐
• Global snapshot: 동일한 패턴의 속도 향상

실제 학습 루프에서 메모리 모니터링 목적으로 snapshot을 주기적으로 호출하는 경우, trace가 불필요하다면 include_traces=False로 오버헤드를 거의 0에 가깝게 줄일 수 있다.

정리

• 기존 API의 기본 동작(include_traces=True)을 변경하지 않아 하위 호환성이 완벽하다.
• 성능 최적화의 핵심은 "하지 않아도 되는 일을 하지 않는 것"이라는 원칙을 잘 보여준다.
• CUDA allocator의 virtual interface를 수정해야 하므로, CUDAPluggableAllocator, CudaMallocAsyncAllocator, HIP 래퍼 등 모든 구현체에 시그니처 변경을 전파한 점이 꼼꼼하다.

참고 자료

• PyTorch CUDA Memory Management 문서 — CUDA 메모리 관리 개요
• torch.cuda.memory_snapshot API — 공식 API 레퍼런스

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