GPU 메모리 사용
GPU 메모리 사용
혼합 정밀도로 배치 크기 16에서 40억(4B) 파라미터 모델을 훈련하려면 대략 85GB의 GPU 메모리가 필요합니다. GPU 메모리의 무엇이 차지하고 있는지와 훈련 중 발생하는 계산 연산을 이해하면 메모리 사용을 줄일 수 있는 위치를 파악하는 데 도움이 됩니다.
출처: 문서
본문
혼합 정밀도로 배치 크기 16에서 4B 파라미터 모델을 훈련하려면 대략 85GB의 GPU 메모리가 필요합니다. GPU 메모리가 무엇으로 채워지는지와 훈련 중 발생하는 계산 연산을 이해하면 메모리 사용을 줄일 수 있는 위치를 찾는 데 도움이 됩니다.
┌─────────────────────────── TENSORS ──────────────────────────────────┐
│ │
│ MODEL WEIGHTS ████████████████████████ 6 bytes/param │
│ (mixed precision) ├── fp16 copy (2B) — forward/backward pass │
│ └── fp32 copy (4B) — stable weight updates │
│ │
│ OPTIMIZER STATES ████████████████████████████████ 8 bytes/param │
│ (Adam) ├── fp32 momentum (4B) │
│ └── fp32 variance (4B) │
│ ↳ quantized Adam (bitsandbytes) → 2 bytes/param │
│ │
│ GRADIENTS ████████████████ 4 bytes/param (fp32) │
│ └── computed in backward pass, per parameter │
│ │
│ ACTIVATIONS ████ varies — batch × seq_len × depth × hidden │
│ (forward cache) └── cached for backward; can OOM even if model │
│ fits (long seqs / large batches) │
│ bf16/fp16 if using mixed precision │
│ │
│ TEMPORARY TENSORS ▓ short-lived (softmax, matmul scratch) │
│ └── peak spikes can cause OOM │
│ │
│ OTHER OVERHEAD ▒ beam search caches, large embedding tables │
└──────────────────────────────────────────────────────────────────────┘
GPU 메모리는 저장된 텐서와 이를 처리하는 연산이라는 두 가지 범주를 담고 있습니다.
텐서
훈련은 저장해야 하는 다양한 유형의 텐서를 필요로 하고 생성합니다.
-
모델 가중치는 GPU에 저장됩니다. 혼합 정밀도 훈련에서는 가중치의 복사본 두 개가 필요합니다. 하나는 forward/backward pass용 fp16, 다른 하나는 안정적인 가중치 업데이트를 위한 "메인 복사본"인 fp32입니다. 이는 파라미터당 6바이트에 해당합니다.
-
Adam 같은 옵티마이저 상태는 파라미터당 모멘텀(momentum)과 분산(variance)이라는 두 개의 추가 텐서를 저장하며, 둘 다 fp32입니다. 이는 파라미터당 추가로 8바이트입니다.
bitsandbytes의 양자화된 Adam 같은 다른 옵티마이저를 사용하면 파라미터당 2바이트로 압축할 수 있습니다.
-
그래디언트 텐서는 backward pass에서 각 파라미터에 대해 계산됩니다. 이는 fp32로 유지되므로 파라미터당 4바이트입니다.
-
forward 활성화는 forward pass에서 계산되어 backward pass에서 그래디언트를 계산하기 위해 캐시됩니다. 이러한 활성화는 배치 크기, 시퀀스 길이, 모델 깊이, hidden 크기에 따라 크기가 달라집니다. 그래서 모델 자체가 맞더라도 배치 크기나 시퀀스 길이로 인해 GPU 메모리가 고갈될 수 있습니다.
-
임시 텐서는 softmax나 행렬 곱셈 같은 연산에 의해 생성되어 각 연산 후 해제됩니다. 단일 연산의 피크가 매우 집약적이면, 메모리가 부족해지는 원인이 되는 일시적인 스파이크가 발생할 수 있습니다.
-
일부 기능은 자체 오버헤드를 추가합니다. 예를 들어 beam search는 여러 출력을 유지하고, 큰 어휘의 임베딩 테이블은 매우 커질 수 있습니다.
연산
훈련 연산에는 세 가지 주요 유형이 있습니다.
-
행렬 곱셈(matmul)은 주요 연산 유형으로, linear 레이어, QKV 프로젝션, 어텐션 출력 프로젝션, FFN 레이어를 포함합니다. 주요 메모리 소비자는 어텐션 점수 행렬인데, 이는 시퀀스 길이의 제곱에 비례해 커집니다. 이것이 긴 시퀀스가 매우 비싼 이유이며, 전체 행렬을 메모리에 구체화하지 않기 위한 다양한 attention backends가 존재하는 이유입니다.
-
softmax나 레이어 정규화 같은 축소(reduction) 연산은 전체 텐서를 읽고, 차원에 걸쳐 통계를 계산한 다음, 다시 텐서를 읽어 적용합니다. 이는 연산마다 여러 번 메모리에 접근해야 합니다.
-
활성화와 dropout 같은 요소별(element-wise) 연산은 각 요소에 독립적으로 함수를 적용하며, 메모리 사용은 텐서 크기에 비례합니다.
다음 단계
- GPU 메모리가 부족하지 않으면서 더 큰 유효 배치 크기로 훈련을 시뮬레이션하는 방법은 Gradient accumulation 가이드를 참조하세요.
- 일부 중간 활성화만 저장하여 메모리 사용을 줄이는 방법은 Gradient checkpointing 가이드를 참조하세요.
- 더 낮은 정밀도 데이터 타입을 사용해 메모리를 줄이고 훈련을 가속화하는 방법은 혼합 정밀도 훈련 가이드를 참조하세요.
- 커스텀 퓨전 커널로 훈련을 가속화하는 방법은 Kernels 가이드를 참조하세요.