LLM 추론 메모리 계산기 쉽게 이해하기 — 개념과 원리
대규모 언어 모델을 GPU에서 돌리려면 먼저 메모리에 들어가는지 알아야 합니다. 메모리를 가장 많이 쓰는 것은 크기가 고정된 모델 가중치와, 문맥 토큰이 늘고 배치에 요청이 추가될 때마다 커지는 키-값(KV) 캐시입니다. 이 계산기는 모델 크기와 구조로 두 값을 추정해 GPU, 양자화 수준, 최대 문맥 길이를 고르는 데 도움을 줍니다.
계산 방법
- 가중치(GB) = 파라미터 수(10억 단위) × 파라미터당 바이트 (FP16/BF16 2바이트, INT8 1바이트, 4비트 0.5바이트)
- KV 캐시(GB) = 2 × 층 수 × KV 헤드 수 × 헤드 차원 × 문맥 토큰 수 × 배치 크기 × 2바이트 ÷ 10⁹
- 합계 = 가중치 + KV 캐시
앞의 2는 키와 값 벡터 두 종류를, 뒤의 2는 FP16 값 하나의 바이트 수를 뜻합니다. GB는 10⁹바이트입니다.
원리
트랜스포머는 문장을 생성할 때 앞선 모든 토큰을 참조합니다. 매 단계마다 지난 토큰의 키와 값을 다시 계산하지 않고, 모든 층과 모든 KV 헤드에 대해 저장해 둡니다. 그래서 이 저장 공간은 문맥 길이와 배치 크기에 정비례해 늘어납니다. 최근 모델 중 상당수는 여러 쿼리 헤드가 하나의 KV 헤드를 함께 쓰는 그룹 쿼리 어텐션(GQA)을 쓰며, 이 경우 캐시가 KV 헤드 수에 비례해 줄어듭니다. 그래서 전체 어텐션 헤드 수가 아니라 KV 헤드 수를 넣어야 합니다.
예제 풀이
- 입력값: 파라미터 80억 개, FP16(2바이트), 32층, KV 헤드 8개, 헤드 차원 128, 문맥 8,192토큰, 배치 1입니다.
- 가중치 = 8 × 2 = 16GB입니다.
- 토큰 하나당 KV 캐시 = 2 × 32 × 8 × 128 × 2 = 131,072바이트(128KiB)입니다.
- 여기에 8,192토큰을 곱하면 1,073,741,824바이트 = 1.0737GB(정확히 1GiB)입니다.
- 합계 = 16 + 1.0737 = 17.07GB입니다.
| 파라미터 수 (10억) | 합계 (GB) |
|---|---|
| 2 | 5.073741824 |
| 4 | 9.073741824 |
| 6 | 13.07374182 |
| 8 | 17.07374182 |
| 10 | 21.07374182 |
| 12 | 25.07374182 |
| 16 | 33.07374182 |
결과 해석
활성값, CUDA 컨텍스트, 프레임워크 부담으로 10~20% 정도를 더하면 이 예는 약 19~20GB가 됩니다. FP16으로는 16GB 그래픽카드에 들어가지 않지만, 8비트나 4비트로 양자화하면(가중치 8GB 또는 4GB) 여유 있게 들어갑니다. KV 캐시는 배치 크기에 정비례하므로 8K 토큰 요청 16개를 동시에 처리하면 캐시만 약 17GB가 필요합니다.
흔한 실수
- GQA 모델에서 KV 헤드 수 대신 쿼리 헤드 수를 넣는 경우
- KV 캐시가 동시 요청마다 따로 든다는 점을 잊는 경우
- GB(10⁹)와 GiB(2³⁰)를 혼동하는 경우
- 가중치를 4비트로 줄이면 KV 캐시도 줄어든다고 생각하는 경우(따로 양자화하지 않으면 FP16입니다)
자주 묻는 질문
층 수, KV 헤드 수, 헤드 차원은 어디서 확인하나요?
모델 설정 파일에서 num_hidden_layers, num_key_value_heads, head_dim(또는 은닉 크기 ÷ 어텐션 헤드 수) 같은 항목으로 확인할 수 있습니다.
문맥이 길어지면 생성 속도도 느려지나요?
네. 새 토큰마다 저장된 모든 토큰을 참조해야 하므로 토큰당 어텐션 비용이 문맥 길이에 따라 늘어납니다.
KV 캐시도 양자화할 수 있나요?
일부 추론 엔진은 8비트 KV 캐시를 지원하며, 약간의 정확도 손실로 캐시 메모리를 대략 절반으로 줄입니다.
출처: A. Vaswani et al., Attention Is All You Need, NeurIPS 2017 · J. Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023