PyTorch에서 GPU 메모리 시각화 및 이해하기 – Hugging Face 블로그 요약

TL;DR

이 포스트는 PyTorch에서 GPU 메모리 스냅샷을 기록하고 시각화하는 방법을 보여주고, 메모리 프로파일의 각 부분(모델 파라미터, 옵티마이저 상태, 활성화 함수 값, 그래디언트, 옵티마이저 중간 단계)이 무엇을 나타내는지 설명하며, 훈련 중 피크 메모리 사용량을 추정하기 위한 일반화된 공식을 제공합니다.

🔎 PyTorch 시각화 도구

PyTorch는 GPU 메모리 사용량을 기록하고 시각화할 수 있는 내장 도구를 제공합니다. 코드 실행 전에 torch.cuda.memory._record_memory_history를 호출하고 실행 후에 torch.cuda.memory._dump_snapshot를 호출하면 https://pytorch.org/memory_viz 에서 볼 수 있는 profile.pkl 파일을 얻을 수 있습니다. 시각화 도구는 메모리 할당 및 해제 타임라인을 표시합니다.

단순한 nn.Linear 레이어를 사용한 예시에서 그래프는 다음과 같이 나타냅니다:

  • 모델 생성 시 가중치와 편향(float32)을 위해 약 2 GB를 할당합니다.
  • 각 입력 텐서는 약 200 MB를 추가합니다.
  • 각 순전파(forward pass) 출력은 약 1 GB를 추가합니다.
  • 이전 단계의 활성화 함수 값(activations)은 역전파에 더 이상 필요하지 않을 때까지 유지되며, 그 이후에는 메모리가 해제됩니다.
  • 변수를 재할당하면 이전에 참조되었던 텐서가 해제됩니다.

📊 훈련 중 메모리 시각화

대규모 언어 모델(Qwen/Qwen2.5-1.5B)과 AdamW 옵티마이저를 사용하는 실제 훈련 루프의 경우, 메모리 프로파일은 훈련 반복당 한 번씩 총 세 번의 스파이크(급증)를 보여줍니다.

프로파일은 다음과 같이 구분할 수 있습니다:

  1. 모델 초기화 – 모델 파라미터(파란색)가 메모리를 차지하며 훈련이 끝날 때까지 할당된 상태로 유지됩니다.
  2. 순전파(Forward Pass) – 활성화 함수 값(주황색)이 레이어별로 계산되어 저장됩니다. 손실(loss) 계산 시점에 정점을 찍습니다.
  3. 역전파(Backward Pass) – 그래디언트(노란색)가 계산됩니다. 활성화 함수 값이 폐기되면서 주황색 영역이 줄어듭니다.
  4. 옵티마이저 단계(Optimizer Step) – 옵티마이저 상태(초록색)는 한 번 초기화됩니다. 옵티마이저는 그래디언트를 사용하여 파라미터를 업데이트하며, 이때 옵티마이저 중간 단계(빨간색)를 일시적으로 저장합니다. 업데이트가 끝나면 그래디언트와 중간 단계 데이터가 해제됩니다.

이 패턴이 각 반복마다 반복되어 관찰된 스파이크를 생성합니다.

📐 메모리 요구사항 추정

피크 메모리 사용량은 프로파일에서 가장 높은 지점이며, 배치 크기에 따라 순전파 또는 옵티마이저 단계 중에 발생할 수 있습니다.

두 경우를 모두 아우르는 일반적인 식은 다음과 같습니다:

Total Memory = Model Memory + Optimizer State + max(Gradients + Optimizer Intermediates, Activations)

각 항의 정의는 다음과 같습니다.

모델 파라미터

Model Memory = N × P

  • N = 파라미터 수
  • P = 바이트 단위 정밀도 (예: float32의 경우 4)

Qwen2.5-1.5B 예시(1.5B 파라미터, float32)의 경우: Model Memory = 1.5 × 10⁹ × 4 bytes = 6 GB.

옵티마이저 상태

파라미터당 두 개의 모멘트를 저장하는 AdamW의 경우: Optimizer State Size = 2 × N × P

그래디언트

Gradients Memory = N × P (모델 파라미터와 크기가 동일).

옵티마이저 중간 단계

Optimizer Intermediates Memory = N × P (모델 파라미터와 크기가 동일).

활성화 함수 값

활성화 함수 메모리는 배치 크기(B), 시퀀스 길이(L), 토큰당 활성화 함수 수(A)에 따라 달라집니다. Activation Memory = A × B × L × P

A는 forward hooks를 통해 직접 측정할 수 있지만, 이 포스트에서는 모델 전반에 걸친 선형 회귀를 통해 도출된 휴리스틱을 제공합니다:

A = 4.6894 × 10⁻⁴ × N + 1.8494 × 10⁶

이 휴리스틱을 사용하면 전체 순전파를 실행하지 않고도 활성화 함수 메모리를 추정할 수 있습니다.

전체 메모리 공식 (결합)

구성 요소 공식들을 대입하면 다음과 같습니다:

Total Memory = N×P + 2×N×P + max(N×P, N×P, A×B×L×P)
             = 3×N×P + max(N×P, A×B×L×P)

max 항은 항상 최소 N×P이므로, 식은 다음과 같이 단순화됩니다:

Total Memory = 3×N×P + A×B×L×P

활성화 함수 값이 지배적일 때의 식입니다. 그렇지 않고 옵티마이저 단계가 지배적이라면 총합은 4×N×P가 됩니다.

이 포스트에는 N, P, B, L 값을 입력하여 추정치를 얻을 수 있는 간단한 도구가 포함되어 있습니다.

🚀 다음 단계

메모리 프로파일을 이해하면 배치 크기 축소, 그래디언트 체크포인팅(gradient checkpointing) 사용 또는 혼합 정밀도(mixed-precision) 훈련 전환과 같이 사용량을 줄이는 방법을 모색할 수 있는 역량을 갖추게 됩니다. TRL 문서의 "Reducing Memory Usage" 섹션에 있는 제안 사항들은 모든 PyTorch 기반 훈련에 폭넓게 적용될 수 있습니다.

🤝 감사의 말

블로그 포스트에 대한 피드백과 제안을 주신 Kashif Rasul 님께 감사드립니다.

Sources