LLM 모델 메모리 사용량 계산하는 방법
LLM이 급격히 떠오르는 요즘, 프로젝트에서 활용하기 위해서는 보유하고 있는 GPU에 맞는 LLM 모델을 잘 선정하는 것이 중요하다. 여기서는 LLM 모델이 잡아먹는 메모리를 계산하는 방법을 간략하게 소개하고, 추가적으로 LLM 모델을 돌리기 위해 필요한 사양을 알려주는 사이트를 소개한다.
들어가기에 앞서
메모리 사용량 계산을 위해서는 LLM 모델 학습 시 소요되는 자원과, 추론 시 소요되는 자원에 대한 이해가 필요하다.
모델 학습 시 고려해야 하는 사항은 다음과 같다.
- Weights (Model Parameters)
- Optimizer
- Gradient
- Activations & Overhead
모델 추론 시 고려해야 하는 사항은 다음과 같다.
- Weights (Model Parameters)
- Total KV Cache Memory
- Activations & Overhead
여기서 Optimizer는 모델 추론에 관여하지 않고, KV Cache는 모델 학습에 관여하지 않는 것을 알 수 있다.
이외에도 자료형은 메모리 사용량에 많은 영향을 미친다.
FP32은 4 Bytes, Mixed Precision은 3 Bytes 및 FP16의 경우 2 Bytes를 필요로 한다. FP16의 경우 precision에 대한 loss를 최소화하면서도 메모리 절약에 특화되어 있으며, 필요에 맞게 자료형을 선택해서 계산해보자
(여기서는 FP16 사용)
모델 학습에 필요한 메모리 사용량 계산
1. Weights (Model Parameters)
FP16 자료형을 기준으로 계산을 하면, 345 million param을 가진 small LLM: 345 million x 2 bytes = 690 MB
llama2-13b 모델: 13 billion x 2 bytes = 26 GB의 메모리를 필요로 한다.
2. Optimizer
FP16 같은 낮은 정밀도로 계산할 경우, 수치적으로 불안해질 수 있기 때문에 대부분 FP32 형식을 사용한다.
- AdamW: 모델 파라미터와 같은 크기의 1, 2차 모멘텀 추적 변수를 추가 유지(모델 파라미터 크기의 3배 메모리)
- SGD: 추가 모멘텀 변수 저장 x (모델 파라미터 크기와 동일)
- LAMB: AdamW, Adam과 비슷한 수준의 메모리 필요
- Adagrad: 2차 모멘텀 변수 유지 (모델 파라미터 크기의 2배 메모리)
- RMSProp: 2차 모멘텀 변수 유지 (모델 파라미터 크기의 2배 메모리)
3. Gradient
Gradient는 모델 파라미터와 동일한 크기의 메모리가 필요하다.
4. Activations & Overhead
Activations는 신경망이 추론 기간동안 결과물로 반환하는 값이고, Temporary Buffer는 계산을 간섭하는데 사용된다.
이들은 크기가 작음에도 불구하고, 모델 파라미터 크기의 5~10%를 사용한다.
5. 메모리 사용량 계산
7B 모델을 학습한다고 했을 때, 필요한 GPU 메모리 사용량을 계산해보자. 모델 학습 시에는 KV Cache를 고려하지 않아도 된다.
- Weights = 14GB (7B x 2 Bytes)
- Optimizer (AdamW) = 84GB (7B x 4 Bytes x 3)
- Gradient = 14GB (7B x 2 Bytes)
- Activatios and Overhead = 0.1 x 14GB = 1.4GB
Total Memory Required: 14GB + 84GB + 14GB + 1.4GB = 132.4GB의 메모리가 필요하다.
모델 추론에 필요한 메모리 사용량 계산
1. KV Cache
모델 추론시에는 Optimizer와 Gradient에 대해서는 고려하지 않아도 되며, KV Cache만 추가적으로 고려하면 된다.
KV Cache는 문장을 구성하는 각 토큰을 생성할 때 필요한 중간 값을 저장한다. 간단히 말하면, 모델은 한 번에 하나의 토큰을 생성하며 문장을 만드는데, 문맥을 유지하기 위해서는 이전 토큰을 기억하는 것이 필요하다. KV Cache는 생성된 토큰들의 Key, value 벡터를 저장함으로써 모델이 다시 계산하는 과정 없이 이전 값을 참조할 수 있도록 도와주는 역할을 한다.
KV Cache의 메모리 사용량 증가를 초래하는 요인으로는 2가지가 있다. 하나는 토큰의 길이로, 토큰이 많아질수록 KV Cache에 더 많은 값들이 저장되기 때문에 메모리 사용량이 증가한다. 또 다른 하나는 다중 사용자로, 여러 요청을 동시에 처리하는 것은 KV Cache 메모리 사용량을 가중시킨다. KV Cache는 Key Vector와 Value Vector로 구성되어 있으며, 토큰 당 총 벡터는 다음과 같이 계산할 수 있다.
L (Number of Layers, depth of the model) x H (Hidden Size, dimensionality)
Llama2-13b 모델이 40층 레이어(L=40), 5120(H=5120) 차원으로 이루어져 있다고 가정해보자.
해당 모델의 Key Vector total element는 40 x 5120 = 204,800이며, 메모리 사이즈는 204,800 x 2 bytes = 409,600 bytes (400KB)를 차지한다. Value Vector 또한 마찬가지로 400 KB를 필요로 하며, 토큰 당 KV Cache 값을 계산하면 800KB가 된다.
출력으로 2000토큰을 생성한다고 하면, 800KB / tok x 2000 = 1.6GB 메모리가 필요하고, 10개의 동시 요청이 존재한다면 16GB의 KV Cache 메모리를 사용한다.
2. 메모리 사용량 계산
8192 Token을 출력하는 Llama-2 13B 모델의 최종 메모리 사용량은 다음과 같이 구할 수 있다.
- Weights = 13 Billion x 2 Bytes = 26GB
- Total KV Cache Memory = 800 KB x 8192 Tokens = 6.6GB
- Activations and Overhead = 0.1 x (26GB + 6.6GB) = 3.2GB
Total Memory Required: 26GB + 6.6GB + 3.2GB = 35.8GB 의 메모리가 필요하다.
만약 10개의 동시 요청이 들어온다고 가정하면, Total KV Cache Memory 값을 다음처럼 바꿔주면 된다.
800KB x 8192 Tokens x 10 Concurrent Request = 66GB, Activations & Overheads: 0.1 x (26GB + 66GB) = 9.2GB
Can it run llm ?
사실 이런 복잡한 계산을 직접하지 않더라도, 우리들의 친구 허깅페이스에 돌리고자 하는 모델과 GPU 환경, 학습할 파라미터 수를 입력하면 관련 정보를 알려주는 페이지(Vokturz/can-it-run-llm)가 있다. 아직 모델을 본격적으로 학습하지 않아 정확한지 확인해보지는 못했지만, 앞으로 자주 사용할 것 같다.
Ref.
https://kipp.ly/transformer-inference-arithmetic/
https://blog.eleuther.ai/transformer-math/#total-inference-memory