Visualize and understand GPU memory in PyTorch

🖼️ 인포그래픽

Visualize and understand GPU memory in PyTorch 내용을 설명하는 본문 이미지

🖼️ 4컷 인포그래픽

💡 한 줄 요약

PyTorch 메모리 스냅샷으로 GPU 사용량을 단계별로 시각화하고, 학습 과정의 최대 메모리를 구성 요소별로 추정하는 방법을 설명한다.

📌 핵심 요약

PyTorch의 메모리 기록 기능으로 실행 중 발생한 GPU 메모리 이벤트를 profile.pkl 파일에 저장하고, 공식 memory_viz 도구에서 시간 흐름에 따른 할당과 해제를 확인할 수 있다.
단순 선형 계층 예제에서는 모델 파라미터, 입력 텐서, 출력 텐서가 각각 차지하는 용량을 원소 수와 자료형의 바이트 수로 계산하고, 참조와 재할당에 따라 메모리가 유지되거나 해제되는 과정을 추적한다.
실제 언어 모델 학습에서는 모델 파라미터, 활성값, 그래디언트, 옵티마이저 상태와 임시 값이 순전파·역전파·옵티마이저 단계에 맞추어 서로 다른 패턴으로 증가하고 감소한다.
최대 메모리 시점은 항상 순전파가 아니며, 배치 크기가 작아지면 옵티마이저 단계가 최고점을 만들 수 있으므로 모든 단계의 후보 피크를 함께 고려해야 한다.
필요 메모리는 모델 파라미터와 옵티마이저 상태에 활성값 또는 그래디언트·옵티마이저 임시 값 중 더 큰 항목을 더해 추정하며, 활성값은 순전파 훅으로 모듈 출력 크기를 측정할 수 있다.

🧩 주요 포인트

PyTorch의 메모리 기록 기능으로 실행 중 발생한 GPU 메모리 이벤트를 profile.pkl 파일에 저장하고, 공식 memory_viz 도구에서 시간 흐름에 따른 할당과 해제를 확인할 수 있다.
단순 선형 계층 예제에서는 모델 파라미터, 입력 텐서, 출력 텐서가 각각 차지하는 용량을 원소 수와 자료형의 바이트 수로 계산하고, 참조와 재할당에 따라 메모리가 유지되거나 해제되는 과정을 추적한다.
실제 언어 모델 학습에서는 모델 파라미터, 활성값, 그래디언트, 옵티마이저 상태와 임시 값이 순전파·역전파·옵티마이저 단계에 맞추어 서로 다른 패턴으로 증가하고 감소한다.
최대 메모리 시점은 항상 순전파가 아니며, 배치 크기가 작아지면 옵티마이저 단계가 최고점을 만들 수 있으므로 모든 단계의 후보 피크를 함께 고려해야 한다.
필요 메모리는 모델 파라미터와 옵티마이저 상태에 활성값 또는 그래디언트·옵티마이저 임시 값 중 더 큰 항목을 더해 추정하며, 활성값은 순전파 훅으로 모듈 출력 크기를 측정할 수 있다.

🧠 상세 정리

1. CUDA 메모리 부족을 원인 단위로 이해하기

GPU 메모리 부족 오류는 전체 용량, 이미 할당된 메모리, 남은 메모리와 PyTorch가 예약한 메모리를 보여 주지만, 어떤 텐서와 학습 단계가 용량을 소비했는지는 바로 설명하지 않는다. 원문은 단순히 GPU가 가득 찼다는 사실을 확인하는 데서 그치지 않고, 실행 도중 메모리가 언제 증가하고 어떤 참조가 해제를 막는지 시간 순서대로 파악하는 것을 목표로 삼는다. 이를 위해 먼저 작은 선형 계층을 대상으로 메모리 변화를 해석한 뒤, 실제 대규모 언어 모델의 전체 학습 루프로 범위를 넓힌다. 마지막에는 관찰된 최고점만 보는 방식을 넘어 모델 파라미터, 옵티마이저 상태, 활성값, 그래디언트와 임시 값으로 요구량을 분해해 추정하는 방법을 제시한다.

2. PyTorch 메모리 스냅샷 기록과 시각화

PyTorch에서는 torch.cuda.memory._record_memory_history를 호출해 GPU 메모리 이벤트 기록을 시작하고, max_entries로 최대 기록 수를 지정할 수 있다. 실행이 끝나면 torch.cuda.memory._dump_snapshot으로 기록을 profile.pkl 파일에 저장한 뒤, enabled=None을 전달해 기록을 중단한다. 생성된 파일은 https://pytorch.org/memory_viz 에 끌어다 놓아 메모리 할당과 해제의 시간적 흐름을 그래프로 확인할 수 있다. 다만 모든 GPU 메모리 이벤트가 기록되므로 프로파일링 단계 수를 제한해야 하며, 원문의 세 단계 언어 모델 학습 예제만으로도 약 8MB의 스냅샷 파일이 만들어진다.

3. 선형 계층의 메모리 크기 계산

첫 예제는 입력 차원이 10,000이고 출력 차원이 50,000인 선형 계층을 GPU에 생성한 뒤, 크기가 5,000×10,000인 입력으로 순전파를 세 번 수행한다. 모델에는 약 5억 개의 가중치와 50,000개의 편향이 있으며 float32 원소당 4바이트를 적용하면 약 2GB를 차지하고, 이 파라미터 메모리는 실행 내내 유지된다. 각 입력 텐서는 5,000×10,000개의 float32 원소이므로 약 200MB이고, 각 출력 텐서는 5,000×50,000개의 원소이므로 약 1GB다. 시각화 결과의 증가 폭을 이 계산값과 대응시키면 모델 생성, 입력 할당, 순전파 출력 생성이 그래프의 어느 구간인지 수치로 식별할 수 있다.

4. 참조와 재할당이 메모리 해제를 결정하는 방식

두 번째 반복에서 inputs 변수가 새 텐서로 바뀌더라도 첫 번째 입력이 곧바로 해제되지는 않는데, 순전파 계산이 역전파에 필요한 활성값으로 그 입력을 계속 참조하기 때문이다. 두 번째 순전파가 완료되어 저장된 활성값이 새 입력으로 대체되면 첫 번째 입력에 대한 참조가 사라지고 해당 메모리를 해제할 수 있다. 출력도 같은 원리로, 새 결과가 output 변수에 재할당되면 이전 출력은 더 이상 참조되지 않아 메모리가 반환된다. 세 번째 반복에서도 입력 생성, 출력 생성, 이전 활성값 해제와 이전 출력 해제가 반복되며, 코드 실행이 끝나면 남아 있던 메모리도 모두 해제된다. 원문은 torch.no_grad()를 적용해 역전파용 활성값을 보관하지 않을 때 그래프가 어떻게 달라지는지 비교해 보라고 제안한다.

5. 실제 언어 모델 학습의 반복별 메모리 패턴

보다 현실적인 예제에서는 Qwen2.5-1.5B 모델을 GPU에 올리고 AdamW 옵티마이저를 구성한 다음, 배치 크기 16과 시퀀스 길이 256의 임의 입력으로 세 번 학습한다. 각 반복은 모델 순전파와 손실 계산, loss.backward(), optimizer.step(), optimizer.zero_grad() 순서로 진행된다. 메모리 프로파일에는 학습 반복 수와 대응하는 세 개의 큰 피크가 나타나며, 각 피크 내부에는 파라미터와 활성값, 그래디언트, 옵티마이저 관련 메모리의 생성과 해제가 겹쳐 있다. 단순 선형 계층보다 그래프는 복잡하지만 반복마다 유사한 모양이 재현되므로, 특정 모델과 학습 루프가 요구하는 GPU 메모리를 분석하고 예상하는 근거로 활용할 수 있다.

6. 순전파·역전파·옵티마이저 단계의 구성 요소

모델 초기화 단계에서 GPU에 적재된 파라미터는 학습이 끝날 때까지 지속적으로 메모리를 차지한다. 순전파에서는 각 계층의 중간 출력인 활성값이 역전파를 위해 저장되며, 계층을 지날수록 누적되어 손실이 계산되는 지점에서 활성값 영역이 최고점에 도달한다. 역전파가 시작되면 그래디언트가 계산되어 저장되는 동시에 이미 사용한 활성값은 더 이상 필요하지 않아 점차 해제된다. optimizer.step()에서는 첫 실행 시 옵티마이저 상태가 초기화되고, 파라미터를 갱신하는 동안 별도의 중간 값이 일시적으로 생성된다. 갱신이 끝나면 그래디언트와 옵티마이저 임시 값이 해제되고 한 번의 반복이 완료되지만, 옵티마이저 상태 자체는 이후 반복에서도 사용되므로 유지된다.

7. 최대 메모리 지점과 일반화된 요구량 공식

배치 크기가 16인 예제만 보면 활성값이 크게 쌓이는 순전파 최고점이 전체 요구량을 결정하는 것처럼 보이며, 이때 필요한 메모리는 모델 파라미터와 옵티마이저 상태 및 활성값의 합이다. 그러나 배치 크기를 16에서 2로 줄이면 활성값이 감소해 순전파 피크가 낮아지고, 그래디언트와 옵티마이저 임시 값이 함께 존재하는 옵티마이저 단계가 더 높은 피크가 된다. 따라서 하나의 단계만 기준으로 삼으면 설정 변화에 따라 실제 요구량을 잘못 추정할 수 있다. 원문은 이를 모델 파라미터와 옵티마이저 상태의 합에 max(그래디언트+옵티마이저 임시 값, 활성값)을 더하는 형태로 일반화한다. 이 공식은 순전파와 파라미터 갱신 중 어느 쪽이 최고점을 만들더라도 더 큰 후보를 최종 요구량에 반영한다.

8. 파라미터·옵티마이저 상태·활성값 추정

모델 파라미터 메모리는 파라미터 수 N에 원소당 정밀도 P바이트를 곱해 계산하며, 15억 개 파라미터를 float32로 저장하면 약 6GB가 필요하다. AdamW는 각 파라미터에 대해 첫 번째와 두 번째 모멘트를 저장하므로, 원문에서는 옵티마이저 상태 크기를 2×N×P로 표현한다. 활성값은 중간 계층 출력 전체를 포함해 직접 계산하기 더 어렵기 때문에, 모든 하위 모듈에 순전파 훅을 등록하고 출력 텐서의 원소 수와 원소 크기를 곱해 합산하는 방법을 사용한다. Qwen2.5-1.5B 예제에서는 입력 토큰당 5,065,216개의 활성값이 측정되며, 전체 활성값 메모리는 토큰당 활성값 수 A, 배치 크기 B, 시퀀스 길이 L, 정밀도 P를 곱해 추정한다. 다만 이 방식은 모델을 실제로 실행해야 하므로 항상 실용적이지 않으며, 원문은 더 큰 모델일수록 활성값도 많다는 직관을 바탕으로 파라미터 수와 활성값 수의 관계를 질문하는 지점에서 제공된 본문을 마친다.

🧾 핵심 주장 / 시사점

GPU 메모리 최적화의 출발점은 총사용량만 확인하는 것이 아니라, 메모리 피크를 만든 학습 단계와 그때 동시에 살아 있는 텐서 종류를 구분하는 것이다.
배치 크기를 줄이면 활성값 메모리는 감소할 수 있지만 파라미터, 옵티마이저 상태와 그래디언트의 비중은 그대로 남으므로, 최대 피크가 순전파에서 옵티마이저 단계로 이동할 수 있다.
메모리 스냅샷은 실제 할당과 해제 흐름을 검증하는 도구이고, 구성 요소별 공식과 순전파 훅은 실행 전후의 요구량을 설명하고 추정하는 보완 수단이다.