LLM 추론 시 GPU 메모리 사용량 알아보기

대규모 언어 모델의 개발과 활용이 점점 더 보편화되면서, 모델의 효율성을 극대화하는 것이 핵심 과제로 떠오르고 있는 추세이다. 특히 추론 단계에서 메모리 사용량을 정확하게 추정하여, 배포하고 운영하는 것은 서비스 측면에서 가장 필수적인 요소이다. 모델의 성능이 아무리 뛰어나도 배포 환경에서 실행되지 않는다면 무용지물이다. DeepSeek-R1 모델은 685B 파라미터 (즉, 6850억 개의 파라미터)에 달하는 것을 클라우드 환경에서 상당한 자원을 요구할 것이며, 엣지 디바이스에서는 실행 자체가 거의 불가능할 것이다. 이에 따라 메모리 사용량을 미리 분석하여 모델을 경량화하거나 최적화 기법을 적용해 배포할 수 있도록 만드는 것이 가장 중요하다고 생각한다. 

 

따라서, 본 게시물에서는 LLM을 사용할 때, 메모리 사용량을 어떻게 추정할 수 있는지에 대해 알아보고자 하며, 메모리를 절약할 수 있는 FlashAttention, KV Cache 등과 같은 방법에 대해서도 소개하고자 한다. 

---

1. 메모리를 잡아 먹는 귀신들 파헤치기 

(i) Model Parameters 

LLM은 보통 수십억에서 수천억 개의 파라미터를 가진 거대한 모델이다. 이런 모델에서 메모리 사용량의 주된 원인이 파라미터 이라고 해도 무방하다. 이는 LLM의 구조와 사용 방식 때문에 이러한 현상은 더 두드러진다. 

• 모델 크기: 그 누구나 아는 이야기지만, LLM은 트랜스포머 아키텍처를 기반으로 하는데, 이 구조는 수많은 레이어(Attention Layer, Feed-Forward Network Layer)와 그에 따른 가중치로 이루어져 있다. 예를 들어, 각 파라미터가 4바이트 라면, 10억 개 파라미터만 있어도 4GB 메모리가 사용된다. 그럼 6,000억 개 파라미터 이라면, 24TB가 필요한 셈이다... 
• 추론 단계: 문장을 생성하거나 질문을 처리할 때, 모델은 입력 토큰을 기반으로 다음 토큰을 예측하게 된다. 이러한 과정에서는 모든 파라미터가 필요하다. 즉, 특정 레이어만 쓰는 게 아니라, 모든 트랜스포머의 모든 레이어를 통과하며 계산하기 때문에 모든 파라미터를 사용해야한다.

(ii) Activations

Activation(활성화 값)은 모델이 입력 데이터를 처리하면서 각 레이어에서 생성되는 중간 결과의 텐서이다. 트랜스포머 모델에서는 특히 Self-Attention과 Feed-Forward Network를 통과할 때 마다 이러한 텐서가 생성된다. 하지만 이렇게 생성된 텐서들은 메모리에 큰 영향을 주게 되는데, 이는 트랜스포머의 Auto-Regressive한 특성 때문이다. 조금 더 직관적으로 보자면, i번째 토큰을 생성할 때 중간 텐서(활성화값)가 만들어지고, 이 텐서는 i+1번째 토큰을 예측하려면 이전 토큰의 정보가 메모리에 계속 남아 있어야 한다. 

 

활성화값은 입력 크기(시퀀스 길이, 배치 크기)와 모델 구조(레이어 수, 히든 사이즈)에 따라 그 크기가 달라집니다. 예를 들어, 6,000억 개 파라미터를 가진 모델을 사용한다고 가정하면, 파라미터 자체는 고정된 6,000억 개(약 24TB)에 해당하지만, 활성화값은 입력 크기나 모델 구조에 따라 이보다 훨씬 더 많은 메모리를 차지할 수도 있다.

 

특히, 이러한 메모리의 부담은 Self-Attention 메커니즘이 크게 기여한다. Self-Attention은 시퀸스 내 모든 토큰 간의 관계를 계산하면서 추가적인 텐서를 생성하기 때문이다. 

(iii) Optimizing Memory Consumption 

위에서 언급한 바와 같이, self-attention 메커니즘이 메모리를 증가시키는 가장 주된 원인이라고 해도 무방하다. 이에 따라, 추론 과정에서 메모리 사용량을 줄이기 위해 몇 가지의 연구를 소개하고자 한다. 

• Grouped Query Attention(GQA): Multi-Head Attention(MHA)은 여러 개의 헤드가 독립적으로 Key와 Value를 계산해서 Attention Score를 도출해내는 방식이다. 반면, GQA는 여러 헤드를 그룹으로 묶고, 각 그룹마다 하나의 Key와 Value를 공유하는 형식으로 진행된다. 이를 쉽게 이해하기 위해 서점의 점원에 비유해볼 수 있다. MHA는 여러 점원이 각기 다른 기준으로 책을 추천하는 상황과 비슷하다. 예를 들어, 한 점원은 자연어 처리, 다른 점원은 선형대수학, 또 다른 점원은 수학 책을 추천한다고 떠올려 보자. 이렇게 하면 다양한 관점에서 책 정보를 얻을 수 있지만, 각 점원이 모든 책의 정보를 따로 외우거나 기록해야 하니 메모리와 노력이 많이 든다. 즉, 연산과 메모리 사용량이 커진다는 뜻이다. 반대로 GQA는 점원들을 몇 개의 그룹으로 나눠서 각 그룹이 하나의 주제를 맡아 책을 추천하는 방식이라고 볼 수 있다. 예를 들어, 점원 1과 2는 '수학'이라는 주제로 묶여 같은 책을 공유하고, 점원 3과 4는 '딥러닝' 주제를 공유하는 식이다. 이렇게 하면 각 점원이 모든 책을 개별적으로 기억할 필요 없이 그룹 내에서 정보를 공유하니 필요한 메모리와 연산이 줄어든다. 결과적으로 GQA는 MHA에 비해 메모리 사용량을 줄이면서도 다양한 관점을 유지할 수 있는 효율적인 방법이라고 할 수 있다. 
• KV Cache: KV Cache의 개념은 생각보다 단순하다. 여러 번 언급 되었던 내용이지만, 트랜스포머의 디코더 기반 언어 모델은 거의 대부분 Auto-Regressive한 특성을 가진다. 다음 토큰을 예측하기 위해서는 이전의 토큰에 대한 정보가 필요하다. 하지만 여기에서의 문제점은 i+1번째 토큰을 예측하기 위해서는 i번째의 토큰의 대한 정보가 필요하다는 것이다. 조금 더 구체적으로 말하자면, 이전 토큰들에 대한 Key와 Value 값 텐서를 매번 새로 계산해야한다. 이에 KV Cache의 개념은 텐서의 값을 매번 계산하는 대신 저장해두고 재사용 하는 것이다.
• Flash Attention: Flash Attention은 Self-Attention을 계산하는 방식을 최적화한 것이다. 이는 시퀸스의 토큰을 작은 블록으로 나눠서 Attention Score을 계산하게 되는데, 이때 이러한 텐서들이 메모리에 저장되지 않는다는 것이다. 대신 필요한 부분만 블록 단위로 계산하고, 그 결과를 쌓아가면 최종 출력을 도출하는 방식이다. 이러한 방식을 통해 메모리 사용량을 크게 줄이며, 연산 속도를 극대화 시킬 수 있다.

---

2. Model Weights 

(i) 기존의 LLM  가중치 저장 방식

대규모 언어 모델(LLM)이 급속하게 발전하면서 모델의 효율성과 성능을 동시에 잡는 방법에 대한 연구가 꾸준하고 진행되고 있다. 너무나도 광범위하지만, 그 중에서도 모델의 가중치(weight) 및 파라미터(parameter)를 어떤 형식으로 저장하는 거에 대해서 알아보고자 한다. 전통적으로 LLM은 32-bit floating point 형식을 사용하였지만, 최근에는 16-bit floating point형식이 주목 받고 있었다.

 

하지만, 2025년도에 접어들면서, 한 단계 더 나아간 8-bit 형식이 본격적으로 소개되었다. 이미 LLaMA3.1:405B과 DeepSeek-V3와 같은 최근에 유명한 모델들은 8-bit 방식을 채택하면서 이에 대한 변화의 물결을 이끌어 가고 있다. 그럼 예시를 통해 32-bit, 16-bit, 8-bit 차이에 대해 아래와 같은 예시를 통해 조금 직관적으로 알아보도록 하겠습니다. 

 

예시 1) 10억 파라미터 모델 (1B) 

• Float32 (32비트): 각 파라미터가 32비트, 즉 4바이트를 사용합니다.
  • 총 메모리: 10억 × 4바이트 = 40억 바이트 = 약 4GB
• Float16 (16비트): 각 파라미터가 16비트, 즉 2바이트를 사용합니다.
  • 총 메모리: 10억 × 2바이트 = 20억 바이트 = 약 2GB
• Float8 (8비트): 각 파라미터가 8비트, 즉 1바이트를 사용합니다.
  • 총 메모리: 10억 × 1바이트 = 10억 바이트 = 약 1GB

float8은 float32에 비해 메모리 사용량이 25% 절감되었으며, float16에 비해서도 50%로 줄었습니다. 1GB면 요즘 스마트폰에서도 충분히 다룰 수 있는 크기라는 점에서, float8의 효율성이 크게 눈에 띈다. 그렇다면, 조금 더 직관적인 예시를 LLaMA3.1:405B 모델을 통해 알아보도록 하겠습니다. 

 

예시 2) 4,050억 파라미터 모델 (LLaMA3.1:405B) 

• Float32 (32비트): 각 파라미터가 32비트, 즉 4바이트를 사용합니다.
  • 총 메모리: 4,050억 × 4바이트 = 1조 6,200억 바이트 = 약 1,620GB (=1.62TB)
• Float16 (16비트): 각 파라미터가 16비트, 즉 2바이트를 사용합니다.
  • 총 메모리: 4,050억 × 2바이트 = 8,100억 바이트 = 약 810GB
• Float8 (8비트): 각 파라미터가 8비트, 즉 1바이트를 사용합니다.
  • 총 메모리: 4,050억 × 1바이트 = 4,050억 바이트 = 약 405GB

추가 예시에서 보았듯이, float8을 사용하면 4,050억 파라미터 모델의 사용량이 405GB으로 줄어든다. 이는 float32의 1.62TB에 비해 무려 75% 감소 하였다. 하지만... 405GB라는 크기는 일반적인 PC 혹은 단일 GPU로는 여전히 감당하기 버겁다. 물론, H100와 같은 고성능 GPU를 여러 개를 병렬로 연결하다면 가능하겠지만, 이 글을 읽고 있는 일반적인 대학원생이라면 불가능에 가깝다. 이에 따라, 우리는 Quantization(양자화)를 통해 메모리 사용량을 더 줄어야한다. 

(ii) Quantization to a Lower Precision 

대규모 언어 모델의 메모리 사용량을 줄이는 데 있어 양자화(Quantization)를 빼놓을 수 없다. 특히 2025년 1월말 쯤 중국과학기술대학교와 마이크로소프트 연구진들은 "Optimizing Large Language Model Training Using FP4 Quantization" 논문을 발간하면서, 4비트로 LLM을 구현하는 것도 가능하다는 것을 제안하였다. 본 논문의 연구진들은 전체 모델의 FP4 양자화를 목표로 하였으며, 이를 수행하기 위해 mixed-precision training과 vector-wise quantization 방법을 통합하여 안정성을 확보하였다. 이에 FP4 정밀도로 BF16 및 FP8와 유사한 정확도를 달성하였다.

 

"Optimizing Large Language Model Training Using FP4 Quantization" 논문에서 제안한  mixed-precision training과 vector-wise quantization 방법을 채택하여 안정성을 확보하였다고 한다. 하지만, 대부분의 양자화 기법에서는 모델의 정확도를 유지하기 위해 모델 전체를 양자화 하지 않는다. 즉, 조금 더 쉽게 이야기 해보자면, 모델의 핵심적인 역할을 하는 입력(Token Embedding)과 출력(Language Modeling Head)를 양자화하면 안된다는 이야기이다. 대체적으로 이런 레이어들은 16비트 정밀도로 유지하여, 토큰의 세부적인 정보까지 포함해야한다.

---

3. Memory Consumption of LLM Inference 

2022년도 NVIDIA에서 발표한 "Reducing Activation Recomputation in Large Transformer Models" 논문의 주요 논점은 트랜스포머 기반 언어 모델에서 메모리 효율성을 극대화하고 학습 속도를 개선하기 위해 Activation Recomputation 줄이는 방법에 초점을 맞추고 있다. 

 

여기에서 Activation Recomputation 이라는 것은 생각보다 단순하다. 모델이 입력 데이터를 처리하며 각 레이어에서 활성화 값(activation)을 생성하게 되는 데, 이러한 값은 역전파에서 기울기를 계산하는 데 필요하다. 하지만 최근 AI 모델들의 파라미터 수는 기하급수적으로 증가하게 되면서 GPT-4, GLaM, Switch Transformer 모델처럼 1조 개 이상의 파라미터를 가지는 트랜스포머 모델들이 대거 등장하고 있다. 하지만, 1조 개 이상이 되는 모델들의 모든 레이어의 활성화 값을 메모리에 저장하는 건 사실상 불가능한 셈이다.

 

따라서 "이런 모델을 어떻게 실질적으로 훈련할 수 있을까?"라는 질문이 매우 중요해졌다. 이에 위 논문에서 제안한 방법을 기반으로 트랜스포머 레이어의 메모리 소비량에 대해 알아보고자 한다.  일반적으로 트랜스포머는 Self-Attention Block과 MLP 블록으로 구성되어 있으며, 각각 두 개의 Layer-Norm으로 연결된다. 각 구성 요소의 메모리 소비량은 다음과 아래와 같이 추정할 수 있다.

 

들어가기에 앞서, 메모리 소비량을 추정하기 위해 알아야하는 것은 다음과 같다:

• s: maximum sequence length
• b: batch size
• h: hidden dimensions
• a: number of attention heads 

(i) Self-Attention Block

Self-Attention Mechanism과 Linear Projection으로 구성되며, 메모리 소비량은 10sbh + 4as^2b으로 계산된다. 이에 대해서 조금 더 세부적으로 알아보고자 한다. 

• Linear Projection retaining its input activations(2sbh):  셀프 어텐션의 선형 투영은 입력 데이터를 메모리에 저장하며, 이는 b × s × h 크기의 입력을 받는다. 이 입력은 쿼리(Q), 키(K), 값(V)로 변환하기 위해 선형 투영을 거치며, 이때 출력도 b × s × h 크기를 가진다. 두 개를 합치면 2sbh가 된다. 여기서 출력 값을 굳이 저장해야 할까라는 의구심이 들 수 있지만, 선형 투영의 출력(Q, K, V)은 다음 단계인 어텐션 스코어 계산(Q × K^T)으로 전달되어 loss에 영향을 주고, 이를 역전파에서 가중치와 입력 기울기를 계산하기 위해 저장해야 한다.
• Linear Projection and Self-Attnetion Input Activations (2sbh each): 이 내용은 위와 비슷해 보이지만 약간 다르다. 여기서는 '선형 투영'과 '셀프 어텐션 입력 활성화'를 별개의 두 항목으로 나눠 각각 2sbh가 필요하다고 본다. 위 내용은 선형 투영을 하나의 통합된 과정으로 계산한 반면, 이 내용은 Q, K, V 각각의 투영 과정을 세분화하여 독립적으로 계산한 것을 의미한다. 그럼 굳이 이렇게 나눈 이유는 다음과 같다. 입력(X)이 Q, K, V에 공유되면서 한 번만 저장하면 되지만, 최적화가 안되었거나 병렬 처리에서 메모리가 더 많이 필요할 수 있는 상황이 있을 수 있기 때문에 이렇게 나눴다고 한다. 
• Query and Key Matrices (4sbh): 위 내용과 굉장히 비슷하기 때문에 설명 생략 
• Softmax (2as^2b): 어텐션 스코어를 계산하면 b × a × s × s 크기의 행렬이 나오지만, 이 또한 역전파를 위해 기울기도 저장해야 하기 때문에 2as^2b가 된다. 
• Attention applied over Self-Attention's Values(2as^2b+2sbh): 위 내용과 굉장히 비슷하기 때문에 설명 생략 

따라서, 전체적으로 Self-Attention 블록에 필요한 메모리는 10sbh + 4as^2b 이다. 

(ii) MLP Block  and Layer-norms 

• MLP 블록에서는 대략 10sbh 메모리가 요구되며, 이에 대한 구체적인 설명은 생략하고자 한다. 또한, 레이어 마다 2 개의 layer-norm이 있으므로 총 4sbh 메모리가 필요하다고 한다. 

(ii) For an Entire Layer 

한 트랜스포머 레이어의 활성화 값을 저장하는 데 필요한 다음과 같이 계산할 수 있다. 

 

그럼 예시로 LLaMA3.3:70B 모델을 배치 크기 1 및 최대 토큰 길이의 수를 2,048로 설정하게 된다면, 아래와 같이 트랜스포머 레이어의 활성화 값을 계산할 수 있다.