Production을 위한 LLM 최적화 기법들 - from 허깅페이스 블로그

 

블로그 주소: https://huggingface.co/blog/optimize-llm

 

LLM들을 배포하기 위해서는 여러 난관들이 존재한다. LLM은 파라미터가 매우 커서 메모리를 많이 소모하고, LLM에 context를 제공하기 위해 긴 input sequence가 필요하다.

 

LLM을 효율적으로 배포하기 위해 크게 3가지 기법들이 사용된다.

 

1. Lower precision

 

LLM들의 파라미터는 숫자 형태이며, float32, bfloat16, float16 등의 형태로 저장된다. 최근에는 float32 형태로 저장되는 모델은 거의 없으며 보통 bfloat16이나 float16 형태로 저장된다. 이를 통해 메모리를 반으로 줄일 수 있다. 그리고 input sequence가 짧을 수록 모델 파라미터의 메모리 비중이 높아진다.

 

보통 LLM들은 가장 큰 GPU 메모리인 80G에 다 담기지 않아서 tensor parallelism이나 pipeline parallelism이 필요하다.

 

Transformers library는 모델 구조를 변형해야 하는 tensor parallelism을 지원하지 않는다. 그 대신 TGI library를 통해 모델의 tensor parallelism을 지원한다.

 

기본적으로 naive pipeline parallelism을 지원한다. 간단히 device="auto" 로 모델을 로드하면 각 layer들을 사용 가능한 GPU에 자동으로 배정해준다. 이 방법은 효과적이기는 하지만 GPU idling 문제가 발생한다. 이를 개선하기 위해 좀 더 발전된 기법이 있다.

 

모델의 크기를 더 줄일 수도 있다. 모델을 8-bit, 4-bit로 줄여도 큰 성능 하락이 없다고 하며, GPTQ 논문에 따르면 2,3bit까지 줄일 수도 있다. 이런 quantization은 text generation에 특히 잘 맞는다. 그 이유는 text generation이 정확한 logit distribution value보다는 가장 확률이 높은 token을 선택하는 것에 집중하기 때문이다. 

 

Quantization은 다음과 같은 과정을 거친다.

1. 모든 weight들을 특정 precision으로 quantize

2. quantized weight를 불러오고, input sequence의 vector들을 bfloat16 precision으로 변환한다.

3. Weight들을 dynamical하게 bfloat16으로 변환한다.

4. Input sequence에 대해 계산한 후, 다시 quantize한다.

 

이러면 메모리는 적게 차지하지만, weight들을 quantize, dequantize하는데 시간이 소요되어 전체적으로 inference 시간이 늘어나게 된다.  

 

 

2. Flash Attention

 

Attention 메카니즘의 self-attention을 계산하기 위해서는 input sequence 길이의 제곱에 비례하는 메모리가 필요하다. Flash Attention은 V x Softmax(Q, K^T)를 작은 chunk로 쪼개 부분들을 순서대로 계산하고, 각 부분들의 정보를 이용해 최종적인 output을 계산한다. 이 과정을 통해 FLOP이 늘어나긴 하지만 chunk 단위로 계산하여 GPU의 구성 요소 중 VRAM보다 크기는 작지만 계산은 빠른 SRAM을 이용할 수 있게 되어 실제 계산 과정은 더 빨라진다.

 

Huggingface의 BetterTransformers library를 이용하면 모델에 쉽게 Flash Attention을 적용할 수 있다.

 

 

3. Architectures

 

LLM은 한 번 학습시키면 구조를 바꾸기 어렵다. 그래서 LLM이 수행할 task를 고려하여 모델의 구조를 정해야 한다. 모델 구조 중에서 긴 input sequence를 처리할 때 bottleneckdㅣ 될 수 있는 요소들이 2가지 있다.

 

3.1 Positional Embeddings

 

Token들의 position 정보를 나타내기 위해서 sinusoidal 또는 learned position embedding이 주로 쓰이는데, 이 방법들은 아래와 같은 문제를 갖고 있다.

 

1) 두 방법 모두 absolute positional embedding을 갖는다. Input text가 길 때는 모델이 input text 간의 relative position 거리를 학습하는 게 더 좋다.(Improve Transformer Models with Better Relative Position Embeddings, RoFormer: Enhanced Transformer with Rotary Position Embedding)

 

2) Learned positional embedding을 사용할 경우, LLM이 input 길이를 고정시켜야 하며, 학습에 사용됐던 input보다 더 긴 input이 들어올 경우 extrapolate하기 어렵다.

 

이런 문제들을 해결하기 위해

1) Rotary Position Embedding(RoPE)

2) ALiBi

가 제안되었다.

 

두 방법 모두 QK^T 계산 결과를 변형해 문장의 순서 정보를 담는다. 이를 통해 training 때 학습했던 input보다 더 긴 input이 들어와도 position 정보를 extrapolate할 수 있다. Extrapolate 능력은 ALiBi가 RoPE보다 좋다. 두 방법 모두 query와 key 간의 사이가 멀 수록 probability 값이 낮아진다.

 

3.2 Key-Value Cache

 

이전 timestep의 key-value vector를 캐싱하면 중복되는 연산(key, value projection)을 줄여 속도를 높일 수 있다.

next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True)

이러면 새로 생성된 token에 대해서만 qK^T를 계산하면 되어서 매 번 QK^T를 계산할 필요가 없다. 그리고 V도 이전 timestep에서 projection 했던 matrix에 새로 생성된 토큰에 대한 벡터 하나만 추가해서 계산하면 된다. q는 새로 만들어진 q만 필요하기 때문에 query-key-value cache가 아니라 key-value cache라고 했나보다. 따라서 input 길이에 따른 메모리도 quadratic이 아닌, linear하게 증가하게 된다.

 

이 key-value cache를 유지하는 것도 메모리가 많이 든다. 그래서 이를 줄이려는 방법론이 제시되었다.

 

Multi-Query-Attention(MQA)

Single-head projection weight pair만 갖는다. 이것을 모든 attention head에 공유한다. 또한 이 기법은 key-value caching을 활용할 때만 유효하다. Key-value caching을 쓰지 않고 매번 K,V를 불러와서 계산하면, 한 개의 head(projection)가 있을 때와 여러 개의 head가 있을 때, 메모리는 QKV를 연산할 만큼만 사용할 것이기에 결국 같은 메모리를 사용한다.

 

Grouped-Query-Attention(GQA)

Single-head projection weight pair만 사용할 경우 성능이 비교적 많이 떨어지고, head 수를 2,4,8 정도로 하여 절충한 방법이다.