속도와 메모리를 위한 LLM 최적화

GPT3/4, Falcon, Llama 같은 대규모 언어 모델(LLM)은 인간 중심 작업을 처리하는 능력이 빠르게 발전하며, 현대 지식 기반 산업에서 필수 도구로 자리 잡고 있습니다. 이 문서는 효율적인 LLM 배포를 위한 저정밀도, Flash Attention, 아키텍처 혁신 기법을 다룹니다.

출처: 문서

본문

GPT3/4, Falcon, Llama 같은 대규모 언어 모델(LLM)은 인간 중심 작업을 처리하는 능력이 빠르게 발전하며, 현대 지식 기반 산업에서 필수 도구로 자리 잡고 있습니다. 그러나 이러한 모델을 실제 작업에 배포하는 것은 여전히 어렵습니다.

인간과 유사한 텍스트 이해 및 생성 능력을 보이려면 LLM은 현재 수십억 개의 파라미터로 구성되어야 합니다(Kaplan et al, Wei et. al 참조). 이는 추론에 필요한 메모리 요구를 증폭시킵니다.
많은 실제 작업에서 LLM에 광범위한 컨텍스트 정보를 제공해야 합니다. 이는 추론 중 매우 긴 입력 시퀀스를 처리하는 모델의 능력을 요구합니다.

이러한 과제의 핵심은 LLM의 계산 및 메모리 능력을 확장하는 데 있으며, 특히 확장된 입력 시퀀스를 처리할 때 그렇습니다.

이 가이드에서는 효율적인 LLM 배포를 위한 효과적인 기법을 다룹니다.

저정밀도(Lower Precision): 연구에 따르면 8비트 및 4비트로 감소된 수치 정밀도로 작업하면 모델 성능의 큰 저하 없이 계산상 이점을 얻을 수 있습니다.
Flash Attention: Flash Attention은 어텐션 알고리즘의 변형으로, 더 메모리 효율적인 접근을 제공할 뿐만 아니라 최적화된 GPU 메모리 활용으로 효율성도 높입니다.
아키텍처 혁신: LLM은 추론 중 항상 같은 방식, 즉 긴 입력 컨텍스트로 자동 회귀 텍스트 생성을 수행하기 때문에, 더 효율적인 추론을 가능하게 하는 특수 모델 아키텍처가 제안되었습니다. 모델 아키텍처에서 가장 중요한 발전은 Alibi, Rotary embeddings, Multi-Query Attention (MQA), Grouped-Query-Attention (GQA)입니다.

이 가이드 전반에서 텐서 관점에서 자동 회귀 생성을 분석합니다. 저정밀도 채택의 장단점을 살펴보고, 최신 어텐션 알고리즘을 포괄적으로 탐구하며, 개선된 LLM 아키텍처를 논의합니다. 그 과정에서 각 기능 개선을 보여주는 실용적인 예제를 실행합니다.

1. 저정밀도

LLM의 메모리 요구는 LLM을 가중치 행렬과 벡터의 집합으로 보고 텍스트 입력을 벡터의 시퀀스로 보면 가장 잘 이해할 수 있습니다. 이하에서 weights라는 용어는 모든 모델 가중치 행렬과 벡터를 의미합니다.

이 가이드를 작성하는 시점에 LLM은 최소 수십억 개의 파라미터로 구성됩니다. 각 파라미터는 4.5689 같은 십진수로 만들어지며, 보통 float32, bfloat16, 또는 float16 형식으로 저장됩니다. 이를 통해 LLM을 메모리에 로드하는 데 필요한 메모리 요구를 쉽게 계산할 수 있습니다.

X십억 파라미터를 가진 모델의 가중치를 로드하려면 float32 정밀도에서 대략 4 * X GB의 VRAM이 필요합니다.

오늘날 모델은 완전한 float32 정밀도로 훈련되는 경우가 드물고, 보통 bfloat16 정밀도로, 덜 빈번하게 float16 정밀도로 훈련됩니다. 따라서 경험 법칙은 다음과 같습니다.

X십억 파라미터를 가진 모델의 가중치를 로드하려면 bfloat16/float16 정밀도에서 대략 2 * X GB의 VRAM이 필요합니다.

더 짧은 텍스트 입력(1024 토큰 미만)의 경우 추론의 메모리 요구는 가중치 로드 메모리 요구가 크게 지배합니다. 따라서 지금은 추론의 메모리 요구가 모델을 GPU VRAM에 로드하는 메모리 요구와 같다고 가정하겠습니다.

모델을 bfloat16으로 로드하는 데 대략 얼마나 많은 VRAM이 필요한지 몇 가지 예를 들어 보겠습니다.

GPT3는 2 * 175 GB = 350 GB VRAM 필요
Bloom은 2 * 176 GB = 352 GB VRAM 필요
Llama-2-70b는 2 * 70 GB = 140 GB VRAM 필요
Falcon-40b는 2 * 40 GB = 80 GB VRAM 필요
MPT-30b는 2 * 30 GB = 60 GB VRAM 필요
bigcode/starcoder는 2 * 15.5 = 31 GB VRAM 필요

이 문서를 작성하는 시점에 시장에서 가장 큰 GPU 칩은 80GB VRAM을 제공하는 A100 및 H100입니다. 앞서 나열한 대부분의 모델은 로드만 해도 80GB 이상이 필요하므로 tensor parallelism 및/또는 pipeline parallelism이 반드시 필요합니다.

🤗 Transformers는 이제 각각의 config 클래스에 base_tp_plan이 있는 지원 모델에 대해 tensor parallelism을 지원합니다. Tensor Parallelism에 대해 더 알아보려면 여기를 참조하세요. 또한 tensor parallelism 친화적인 방식으로 모델을 작성하는 데 관심이 있다면 text-generation-inference 라이브러리를 살펴보세요.

네이티브 pipeline parallelism은 기본적으로 지원됩니다. 이를 위해 device="auto"로 모델을 로드하면 여기에서 설명한 대로 다른 레이어들을 사용 가능한 GPU에 자동으로 배치합니다. 그러나 매우 효과적이지만, 이 네이티브 pipeline parallelism은 GPU 유휴 문제를 해결하지 못한다는 점에 유의하세요. 이를 위해서는 여기에서 설명한 더 고급 pipeline parallelism이 필요합니다.

8 x 80GB A100 노드에 접근할 수 있다면 BLOOM을 다음과 같이 로드할 수 있습니다.

!pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

device_map="auto"를 사용하면 attention 레이어가 사용 가능한 모든 GPU에 균등하게 분배됩니다.

이 가이드에서는 단일 40GB A100 GPU 장치 칩에서 실행할 수 있는 bigcode/octocoder를 사용합니다. 앞으로 적용할 모든 메모리 및 속도 최적화는 모델 또는 tensor parallelism이 필요한 모델에도 동일하게 적용된다는 점에 유의하세요.

모델이 bfloat16 정밀도로 로드되므로 위의 경험 법칙을 사용하면 bigcode/octocoder로 추론을 실행하는 메모리 요구가 약 31GB VRAM일 것으로 예상합니다. 시도해 보겠습니다.

먼저 모델과 토크나이저를 로드한 다음 둘 다 Transformers의 pipeline 객체에 전달합니다.

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", dtype=torch.bfloat16, device_map="auto", pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:"

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

출력:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

좋아요. 이제 결과를 직접 사용해 바이트를 기가바이트로 변환할 수 있습니다.

def bytes_to_giga_bytes(bytes):
  return bytes / 1024 / 1024 / 1024

torch.accelerator.memory.max_memory_allocated를 호출해 피크 가속기 메모리 할당을 측정해 봅시다.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

출력:

29.0260648727417

대략적인 계산과 매우 유사합니다! 바이트에서 킬로바이트로 가려면 1000이 아닌 1024를 곱해야 하므로 숫자가 정확하지 않다는 것을 알 수 있습니다. 따라서 대략적인 공식은 "최대 X GB" 계산으로 이해할 수 있습니다. 만약 모델을 완전한 float32 정밀도로 실행하려고 했다면 무려 64GB의 VRAM이 필요했을 것입니다.

오늘날 거의 모든 모델이 bfloat16으로 훈련되므로, GPU가 bfloat16을 지원한다면 모델을 완전한 float32 정밀도로 실행할 이유가 없습니다. Float32는 모델을 훈련하는 데 사용된 정밀도보다 더 나은 추론 결과를 주지 않습니다.

Hub에서 모델 가중치가 어떤 형식으로 저장되어 있는지 확실하지 않다면 체크포인트의 config에서 "dtype"을 항상 확인할 수 있습니다, 예: 여기. 원본 타입이 float32인 경우를 제외하고는 config에 적힌 것과 같은 정밀도 타입으로 from_pretrained(..., dtype=...)에서 모델을 설정하는 것이 좋습니다. float32인 경우 추론에는 float16 또는 bfloat16 중 하나를 사용할 수 있습니다.

피크 할당 GPU 메모리를 정확히 측정할 수 있도록 모든 할당 메모리를 해제하는 flush(...) 함수를 정의해 봅시다.

del pipe
del model

import gc
import torch

def flush():
  gc.collect()
  torch.accelerator.empty_cache()
  torch.accelerator.reset_peak_memory_stats()

다음 실험을 위해 지금 호출해 봅시다.

flush()

Accelerate 라이브러리에서 XPU, MLU, NPU, MPS 등 다양한 하드웨어 백엔드를 고려하는 장치에 구애받지 않는 유틸리티 메서드인 release_memory를 사용할 수도 있습니다.

from accelerate.utils import release_memory
# ...

release_memory(model)

이제 GPU에 32GB의 VRAM이 없다면 어떨까요? 모델 가중치를 8비트 또는 4비트로 양자화해도 성능 큰 손실 없이 사용할 수 있다는 것이 밝혀졌습니다(Dettmers et al. 참조). 모델은 최근 GPTQ 논문 🤯에서 보여준 것처럼 3비트 또는 2비트로도 양자화할 수 있으며 성능 손실이 허용 가능합니다.

너무 많은 세부 사항에 들어가지 않고, 양자화 기법은 모델의 추론 결과를 최대한 정확하게(bfloat16에 최대한 가깝게) 유지하면서 가중치의 정밀도를 낮추는 것을 목표로 합니다. 양자화는 우리가 신경 쓰는 것이 가장 가능성 높은 다음 토큰 집합을 선택하는 것뿐이고 다음 토큰 logit 분포의 정확한 값은 크게 신경 쓰지 않기 때문에 텍스트 생성에서 특히 잘 작동합니다. 중요한 것은 다음 토큰 logit 분포가 대략 동일하게 유지되어 argmax 또는 topk 연산이 같은 결과를 준다는 것입니다.

여기서 자세히 논의하지는 않지만 다양한 양자화 기법이 있으며, 일반적으로 모든 양자화 기법은 다음과 같이 동작합니다.

1. 모든 가중치를 대상 정밀도로 양자화
1. 양자화된 가중치를 로드하고, 입력 시퀀스 벡터를 bfloat16 정밀도로 전달
1. bfloat16 정밀도의 입력 벡터로 계산을 수행하기 위해 가중치를 bfloat16으로 동적으로 비양자화

간단히 말해, 이것은 $X$가 inputs, $W$가 가중치 행렬, $Y$가 출력인 inputs-weight matrix 곱셈이

$$ Y = X * W $$

모든 행렬 곱셈에 대해

$$ Y = X * \text{dequantize}(W) $$

로 바뀌는 것을 의미합니다. 입력이 네트워크 그래프를 통과하면서 모든 가중치 행렬에 대해 비양자화와 재양자화가 순차적으로 수행됩니다.

따라서 추론 시간은 양자화된 가중치를 사용할 때 감소하지 않고 오히려 증가하는 경우가 많습니다. 이론은 충분하니 시도해 봅시다! Transformers로 가중치를 양자화하려면 bitsandbytes 라이브러리가 설치되어 있는지 확인해야 합니다.

!pip install bitsandbytes

그런 다음 from_pretrained에 load_in_8bit=True 플래그를 추가하기만 하면 모델을 8비트 양자화로 로드할 수 있습니다.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

이제 예제를 다시 실행하고 메모리 사용을 측정해 봅시다.

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

출력:

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

좋아요, 이전과 같은 결과를 얻습니다. 정확도 손실이 없습니다! 이번에는 메모리가 얼마나 사용되었는지 살펴봅시다.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

출력:

15.219234466552734

훨씬 적습니다! 15GB 조금 넘게 줄었고, 따라서 4090 같은 소비자 GPU에서도 이 모델을 실행할 수 있습니다. 메모리 효율에서 매우 좋은 이득을 보고 있고 모델 출력의 저하는 거의 없습니다. 그러나 추론 중 약간의 속도 저하도 볼 수 있습니다.

모델을 삭제하고 메모리를 다시 플러시합니다.

del model
del pipe

flush()

4비트 양자화가 주는 피크 GPU 메모리 소비가 무엇인지 살펴봅시다. 모델을 4비트로 양자화하는 것은 이전과 같은 API로 할 수 있습니다. 이번에는 load_in_8bit=True 대신 load_in_4bit=True를 전달합니다.

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=True), pad_token_id=0)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
result

출력:

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n    return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

이전과 거의 같은 출력 텍스트를 볼 수 있습니다. 코드 스니펫 직전의 python만 빠져 있습니다. 메모리가 얼마나 필요했는지 살펴봅시다.

bytes_to_giga_bytes(torch.accelerator.max_memory_allocated())

출력:

9.543574333190918

겨우 9.5GB입니다! 150억 개 이상의 파라미터 모델치고는 정말 많지 않습니다.

여기서 모델의 정확도 저하가 매우 적음을 보지만, 실제로 4비트 양자화는 8비트 양자화나 완전한 bfloat16 추론과 비교해 종종 다른 결과를 줄 수 있습니다. 시도해 보는 것은 사용자의 몫입니다.

또한 4비트 양자화를 위한 더 공격적인 양자화 방법으로 인해 추론 중 $\text{quantize}$와 $\text{dequantize}$가 더 오래 걸려 이전 8비트 양자화보다 여기서 추론이 다시 약간 느렸다는 점에 유의하세요.

del model
del pipe

flush()

전체적으로 OctoCoder를 8비트 정밀도로 실행하면 필요한 GPU VRAM이 32G GPU VRAM에서 15GB로 줄어들었고, 모델을 4비트 정밀도로 실행하면 필요한 GPU VRAM이 9GB 조금 넘게 더 줄어드는 것을 확인했습니다.

4비트 양자화를 사용하면 대부분의 사람들이 쉽게 접근할 수 있는 RTX3090, V100, T4 같은 GPU에서 모델을 실행할 수 있습니다.

양자화에 대한 더 많은 정보와 4비트보다 더 적은 GPU VRAM을 요구하도록 모델을 양자화하는 방법을 보려면 GPT-QModel 구현을 살펴보는 것이 좋습니다.

결론적으로, 모델 양자화는 개선된 메모리 효율을 정확도와 경우에 따라 추론 시간과 맞바꾼다는 것을 기억하는 것이 중요합니다.

GPU 메모리가 사용 사례의 제약이 아니라면 양자화를 살펴볼 필요가 없는 경우가 많습니다. 그러나 많은 GPU는 양자화 방법 없이는 LLM을 실행할 수 없으며, 이 경우 4비트 및 8비트 양자화 기법은 매우 유용한 도구입니다.

더 자세한 사용 정보를 보려면 Transformers 양자화 문서를 살펴보는 것이 좋습니다. 다음으로 더 나은 알고리즘과 개선된 모델 아키텍처를 사용해 계산 및 메모리 효율을 어떻게 개선할 수 있는지 알아봅시다.

2. Flash Attention

오늘날 최고 성능의 LLM은 대부분 피드포워드 레이어, 활성화 레이어, 레이어 정규화 레이어, 그리고 가장 중요한 self-attention 레이어로 구성된 동일한 기본 아키텍처를 공유합니다.

Self-attention 레이어는 모델이 입력 토큰 간의 컨텍스트 관계를 이해할 수 있게 해주므로 대규모 언어 모델(LLM)의 핵심입니다. 그러나 self-attention 레이어의 피크 GPU 메모리 소비는 입력 토큰 수(이하 $N$으로 표기, 시퀀스 길이라고도 함)에 대해 계산과 메모리 복잡도 모두 이차적으로 증가합니다. 이는 더 짧은 입력 시퀀스(최대 1000 입력 토큰)에서는 그다지 눈에 띄지 않지만, 더 긴 입력 시퀀스(약 16000 입력 토큰)에서는 심각한 문제가 됩니다.

자세히 살펴봅시다. 길이 $N$의 입력 $\mathbf{X}$에 대한 self-attention 레이어의 출력 $\mathbf{O}$를 계산하는 공식은 다음과 같습니다.

$$ \textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X} $$

$\mathbf{X} = (\mathbf{x}1, ... \mathbf{x}{N})$은 attention 레이어의 입력 시퀀스입니다. 프로젝션 $\mathbf{Q}$와 $\mathbf{K}$는 각각 $N$개의 벡터로 구성되어 $\mathbf{QK}^T$가 크기 $N^2$가 됩니다.

LLM은 보통 여러 attention head를 가지므로 여러 self-attention 계산을 병렬로 수행합니다. LLM에 40개의 attention head가 있고 bfloat16 정밀도로 실행된다고 가정하면, $\mathbf{QK^T}$ 행렬을 저장하는 메모리 요구를 $40 * 2 * N^2$ 바이트로 계산할 수 있습니다. $N=1000$일 때 약 50MB의 VRAM만 필요하지만, $N=16000$일 때는 19GB의 VRAM이 필요하고, $N=100,000$일 때는 $\mathbf{QK}^T$ 행렬을 저장하는 것만으로 거의 1TB가 필요합니다.

요컨대, 기본 self-attention 알고리즘은 큰 입력 컨텍스트에서 빠르게 엄청나게 메모리 비용이 높아집니다.

LLM이 텍스트 이해와 생성에서 개선됨에 따라 점점 더 복잡한 작업에 적용됩니다. 모델이 한때 몇 문장의 번역이나 요약을 처리했다면 이제 전체 페이지를 관리하며, 광범위한 입력 길이를 처리하는 능력을 요구합니다.

큰 입력 길이에 대한 엄청난 메모리 요구를 어떻게 없앨 수 있을까요? $\mathbf{QK}^T$ 행렬을 없애는 self-attention 메커니즘을 계산하는 새로운 방법이 필요합니다. Tri Dao et al.는 정확히 그러한 새로운 알고리즘을 개발했고 그것을 Flash Attention이라고 불렀습니다.

간단히 말해, Flash Attention은 $\mathbf{V} \times \text{Softmax}(\mathbf{QK}^T)$ 계산을 분해하고, 대신 여러 softmax 계산 단계를 반복하여 출력의 더 작은 청크를 계산합니다.

$$ \textbf{O}i \leftarrow s^a{ij} * \textbf{O}i + s^b{ij} * \mathbf{V}{j} \times \text{Softmax}(\mathbf{QK}^T{i,j}) \text{ for multiple } i, j \text{ iterations} $$

여기서 $s^a_{ij}$와 $s^b_{ij}$는 모든 $i$와 $j$에 대해 다시 계산해야 하는 softmax 정규화 통계입니다.

전체 Flash Attention은 조금 더 복잡하며, 이 가이드에서는 너무 깊이 들어가는 것은 범위를 벗어나므로 크게 단순화되었습니다. 자세한 내용은 잘 쓰인 Flash Attention 논문을 살펴보시기 바랍니다.

여기서 중요한 점은 다음과 같습니다.

softmax 정규화 통계를 추적하고 영리한 수학을 사용함으로써, Flash Attention은 $N$에 따라 선형적으로만 증가하는 메모리 비용으로 기본 self-attention 레이어와 수치적으로 동일한 출력을 제공합니다.

공식을 보면 더 많은 계산을 해야 하므로 Flash Attention이 기본 self-attention 공식보다 훨씬 느릴 것이라고 직관적으로 말할 수 있습니다. 실제로 Flash Attention은 softmax 정규화 통계를 계속 다시 계산해야 하므로 일반 attention보다 더 많은 FLOP을 요구합니다(관심이 있다면 논문 참조).

그러나 Flash Attention은 GPU의 느리고 높은 대역폭 메모리(VRAM)에 대한 요구를 크게 줄이고 더 빠른 온칩 메모리(SRAM)에 집중할 수 있어 추론 시 기본 attention보다 훨씬 빠릅니다.

본질적으로 Flash Attention은 출력 벡터 $\mathbf{O}$를 계산하기 위해 더 느린 VRAM 메모리에 접근하는 대신 모든 중간 쓰기 및 읽기 연산이 빠른 온칩 SRAM 메모리를 사용할 수 있도록 합니다.

실제로 현재로선 Flash Attention을 사용할 수 있다면 사용하지 않을 이유가 전혀 없습니다. 이 알고리즘은 수학적으로 같은 출력을 주고 더 빠르며 더 메모리 효율적입니다.

3. 아키텍처 혁신

지금까지 다음을 통해 계산 및 메모리 효율을 개선하는 방법을 살펴봤습니다.

가중치를 더 낮은 정밀도 형식으로 캐스팅
self-attention 알고리즘을 더 메모리 및 계산 효율적인 버전으로 교체

이제 긴 텍스트 입력이 필요한 작업, 예를 들어:

검색 증강 질의 응답(Retrieval augmented Question Answering),
요약(Summarization),
채팅(Chat)

에 가장 효과적이고 효율적이도록 LLM의 아키텍처를 어떻게 변경할 수 있는지 알아봅시다.

채팅은 LLM이 긴 텍스트 입력을 처리하는 것뿐만 아니라 사용자와 어시스턴트 간의 왕복 대화(ChatGPT처럼)를 효율적으로 처리할 수 있어야 합니다.

한번 훈련되면 기본 LLM 아키텍처는 변경하기 어렵기 때문에, LLM의 작업을 미리 고려하고 그에 따라 모델의 아키텍처를 최적화하는 것이 중요합니다. 큰 입력 시퀀스에서 빠르게 메모리 및/또는 성능 병목이 되는 모델 아키텍처의 두 가지 중요한 구성 요소가 있습니다.

포지셔널 임베딩(positional embeddings)
키-값 캐시(key-value cache)

각 구성 요소를 더 자세히 살펴봅시다.

3.1 LLM의 포지셔널 임베딩 개선

Self-attention은 각 토큰을 다른 모든 토큰과 관계시킵니다. 예를 들어 텍스트 입력 시퀀스 *"Hello", "I", "love", "you"*의 $\text{Softmax}(\mathbf{QK}^T)$ 행렬은 다음과 같을 수 있습니다.

각 단어 토큰은 다른 모든 단어 토큰에 attend하는 확률 질량을 가지며, 따라서 다른 모든 단어 토큰과 관계가 설정됩니다. 예를 들어 *"love"*라는 단어는 *"Hello"*에 5%, *"I"*에 30%, 자신에게 65% attend합니다.

포지셔널 임베딩이 없는 self-attention 기반 LLM은 텍스트 입력들의 서로에 대한 위치를 이해하는 데 큰 어려움을 겪을 것입니다. 이는 $\mathbf{QK}^T$가 계산한 확률 점수가 각 단어 토큰을 서로의 상대적 위치 거리와 무관하게 $O(1)$ 계산으로 다른 각 단어 토큰과 관계시키기 때문입니다. 따라서 포지셔널 임베딩이 없는 LLM에서는 각 토큰이 다른 모든 토큰과 같은 거리를 가지는 것처럼 보이며, 예를 들어 *"Hello I love you"*와 *"You love I hello"*를 구분하는 것은 매우 어려울 것입니다.

LLM이 문장 순서를 이해하려면 추가적인 *신호(cue)*가 필요하며, 일반적으로 포지셔널 인코딩(또는 포지셔널 임베딩)의 형태로 적용됩니다. 포지셔널 인코딩은 각 토큰의 위치를 LLM이 문장 순서를 더 잘 이해하는 데 활용할 수 있는 수치 표현으로 인코딩합니다.

Attention Is All You Need 논문의 저자들은 정현파 포지셔널 임베딩 $\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$을 도입했습니다. 여기서 각 벡터 $\mathbf{p}_i$는 위치 $i$의 정현파 함수로 계산됩니다. 포지셔널 인코딩은 입력 시퀀스 벡터에 단순히 더해져 $\mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N$ = $\mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N$이 되어 모델이 문장 순서를 더 잘 학습하도록 신호를 줍니다.

고정 포지셔널 임베딩을 사용하는 대신, 다른 사람들(Devlin et al. 등)은 학습된 포지셔널 인코딩을 사용했으며 여기서 포지셔널 임베딩 $\mathbf{P}$는 훈련 중에 학습됩니다.

정현파 및 학습된 포지셔널 임베딩은 LLM에 문장 순서를 인코딩하는 지배적인 방법이었지만, 이 포지셔널 인코딩과 관련된 몇 가지 문제가 발견되었습니다.

정현파 및 학습된 포지셔널 임베딩은 둘 다 절대 포지셔널 임베딩입니다, 즉, 각 위치 id: $0, \ldots, N$에 대해 고유한 임베딩을 인코딩합니다. Huang et al.과 Su et al.이 보여준 것처럼, 절대 포지셔널 임베딩은 긴 텍스트 입력에서 LLM 성능이 좋지 않습니다. 긴 텍스트 입력에서는 모델이 절대 위치 대신 입력 토큰이 서로에 대해 가지는 상대적 위치 거리를 학습하는 것이 유리합니다.
학습된 포지셔널 임베딩을 사용할 때 LLM은 고정 입력 길이 $N$으로 훈련해야 하므로, 훈련된 것보다 더 긴 입력 길이로 외삽(extrapolate)하기 어렵습니다.

최근에는 위의 문제를 해결할 수 있는 상대 포지셔널 임베딩이 더 인기를 얻었으며, 가장 주목할 만한 것은 다음과 같습니다.

RoPE와 ALiBi 모두 단어 토큰이 서로 관계를 맺는 곳인 self-attention 알고리즘에서 직접 LLM에 문장 순서를 신호로 주는 것이 가장 좋다고 주장합니다. 더 구체적으로, 문장 순서는 $\mathbf{QK}^T$ 계산을 수정하여 신호를 주어야 합니다.

너무 많은 세부 사항에 들어가지 않고, RoPE는 포지셔널 정보를 쿼리-키 쌍, 예를 들어 $\mathbf{q}_i$와 $\mathbf{x}_j$에 각 벡터의 문장 위치를 나타내는 $i, j$에 대해 각각 각도 $\theta * i$와 $\theta * j$로 각 벡터를 회전시켜 인코딩할 수 있다고 지적합니다.

$$ \mathbf{\hat{q}}_i^T \mathbf{\hat{x}}_j = \mathbf{{q}}i^T \mathbf{R}{\theta, i -j} \mathbf{{x}}_j. $$

여기서 $\mathbf{R}_{\theta, i - j}$는 회전 행렬을 나타냅니다. $\theta$는 훈련 중에 학습되지 않고, 훈련 중 최대 입력 시퀀스 길이에 따라 달라지는 사전 정의된 값으로 설정됩니다.

이렇게 함으로써 $\mathbf{q}_i$와 $\mathbf{q}_j$ 사이의 확률 점수는 $i \ne j$일 때만 영향을 받고, 각 벡터의 특정 위치 $i$와 $j$와 무관하게 상대 거리 $i - j$에만 의존합니다.

RoPE는 오늘날 가장 중요한 여러 LLM에서 사용됩니다, 다음과 같이.

대안으로 ALiBi는 훨씬 더 간단한 상대 위치 인코딩 방식을 제안합니다. 입력 토큰이 서로에 대해 가지는 상대 거리는 softmax 계산 직전에 $\mathbf{QK}^T$ 행렬의 각 쿼리-키 항목에 사전 정의된 값 m으로 스케일된 음의 정수로 더해집니다.

ALiBi 논문에서 보여준 것처럼, 이 간단한 상대 포지셔널 인코딩은 매우 긴 텍스트 입력 시퀀스에서도 높은 성능을 유지할 수 있게 합니다.

ALiBi는 오늘날 가장 중요한 여러 LLM에서 사용됩니다, 다음과 같이.

RoPE와 ALiBi 포지셔널 인코딩은 둘 다 훈련 중에 보지 못한 입력 길이로 외삽할 수 있지만, ALiBi가 RoPE보다 외삽이 기본적으로 훨씬 잘 작동하는 것으로 나타났습니다. ALiBi의 경우 입력 시퀀스의 길이에 맞게 하삼각 위치 행렬의 값을 단순히 늘리면 됩니다. RoPE의 경우 훈련 중에 사용된 것과 같은 $\theta$를 유지하면 훈련 중에 본 것보다 훨씬 긴 텍스트 입력을 전달할 때 결과가 좋지 않습니다, c.f Press et al.. 그러나 커뮤니티는 $\theta$를 적응시키는 몇 가지 효과적인 트릭을 발견하여 RoPE 포지셔널 임베딩이 외삽된 텍스트 입력 시퀀스에서 잘 작동하게 했습니다(여기 참조).

RoPE와 ALiBi는 모두 훈련 중에 학습되지 않고 다음 직관에 기반한 상대 포지셔널 임베딩입니다.

텍스트 입력에 대한 포지셔널 신호는 self-attention 레이어의 $\mathbf{QK}^T$ 행렬에 직접 주어져야 합니다.
LLM은 일정한 상대 거리 포지셔널 인코딩을 학습하도록 장려되어야 합니다.
텍스트 입력 토큰이 서로 멀수록, 그들의 쿼리-값 확률이 낮아져야 합니다. RoPE와 ALiBi 모두 서로 멀리 떨어진 토큰의 쿼리-키 확률을 낮춥니다. RoPE는 쿼리-키 벡터 사이의 각도를 증가시켜 그들의 벡터 곱을 줄여 낮춥니다. ALiBi는 벡터 곱에 큰 음수를 더해 낮춥니다.

결론적으로, 큰 텍스트 입력을 처리해야 하는 작업에 배포하려는 LLM은 RoPE와 ALiBi 같은 상대 포지셔널 임베딩으로 훈련하는 것이 더 좋습니다. 또한 RoPE와 ALiBi가 있는 LLM이 $N_1 = 2048$ 같은 고정 길이로만 훈련되었더라도, 포지셔널 임베딩을 외삽하여 $N_2 = 8192 > N_1$ 같은 $N_1$보다 훨씬 큰 텍스트 입력과 함께 실제로 사용할 수 있습니다.

3.2 키-값 캐시

LLM으로 자동 회귀 텍스트 생성은 입력 시퀀스를 반복적으로 넣고, 다음 토큰을 샘플링하고, 다음 토큰을 입력 시퀀스에 추가하고, LLM이 생성을 완료했음을 나타내는 토큰을 생성할 때까지 계속하는 방식으로 작동합니다.

자동 회귀 생성이 어떻게 작동하는지 더 시각적인 설명은 Transformers의 Generate Text 튜토리얼을 살펴보세요.

실제로 자동 회귀가 어떻게 작동하는지 보여주는 빠른 코드 스니펫을 실행해 봅시다. torch.argmax로 가장 가능성 높은 다음 토큰을 가져올 것입니다.

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits = model(input_ids)["logits"][:, -1:]
  next_token_id = torch.argmax(next_logits,dim=-1)

  input_ids = torch.cat([input_ids, next_token_id], dim=-1)
  print("shape of input_ids", input_ids.shape)

generated_text = tokenizer.batch_decode(input_ids[:, -5:])
generated_text

출력:

shape of input_ids torch.Size([1, 21])
shape of input_ids torch.Size([1, 22])
shape of input_ids torch.Size([1, 23])
shape of input_ids torch.Size([1, 24])
shape of input_ids torch.Size([1, 25])
[' Here is a Python function']

볼 수 있듯이 텍스트 입력 토큰을 방금 샘플링한 토큰만큼 매번 늘립니다.

매우 드문 예외를 제외하고 LLM은 causal language modeling 목표로 훈련되므로 attention 점수의 상삼각 행렬을 마스킹합니다. 이것이 위의 두 다이어그램에서 attention 점수가 비어 있는(즉 0 확률) 이유입니다. causal language modeling에 대한 빠른 요약은 Illustrated Self Attention blog를 참조할 수 있습니다.

결과적으로 토큰은 결코 이후 토큰에 의존하지 않습니다. 더 구체적으로 $\mathbf{q}_i$ 벡터는 $j > i$인 경우 어떤 키, 값 벡터 $\mathbf{k}_j, \mathbf{v}j$와도 관계를 맺지 않습니다. 대신 $\mathbf{q}i$는 이전 키-값 벡터 $\mathbf{k}{m < i}, \mathbf{v}{m < i} \text{ , for } m \in \{0, \ldots i - 1\}$에만 attend합니다. 불필요한 계산을 줄이기 위해, 모든 이전 타임스텝에 대한 각 레이어의 키-값 벡터를 캐시할 수 있습니다.

이하에서 우리는 LLM에 키-값 캐시를 검색하고 각 forward pass에 전달하여 사용하도록 지시할 것입니다. Transformers에서 use_cache 플래그를 forward 호출에 전달하여 키-값 캐시를 검색하고 현재 토큰과 함께 전달할 수 있습니다.

past_key_values = None # past_key_values is the key-value cache
generated_tokens = []
next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)

for _ in range(5):
  next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple()
  next_logits = next_logits[:, -1:]
  next_token_id = torch.argmax(next_logits, dim=-1)

  print("shape of input_ids", next_token_id.shape)
  print("length of key-value cache", past_key_values.get_seq_length())  # past_key_values are of shape [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim]
  generated_tokens.append(next_token_id.item())

generated_text = tokenizer.batch_decode(generated_tokens)
generated_text

출력:

shape of input_ids torch.Size([1, 1])
length of key-value cache 20
shape of input_ids torch.Size([1, 1])
length of key-value cache 21
shape of input_ids torch.Size([1, 1])
length of key-value cache 22
shape of input_ids torch.Size([1, 1])
length of key-value cache 23
shape of input_ids torch.Size([1, 1])
length of key-value cache 24
[' Here', ' is', ' a', ' Python', ' function']

볼 수 있듯이 키-값 캐시를 사용할 때 텍스트 입력 토큰의 길이는 증가하지 않고 단일 입력 벡터로 유지됩니다. 반면 키-값 캐시의 길이는 매 디코딩 단계마다 1씩 증가합니다.

키-값 캐시를 사용한다는 것은 $\mathbf{QK}^T$가 본질적으로 $\mathbf{q}_c\mathbf{K}^T$로 줄어드는 것을 의미합니다. 여기서 $\mathbf{q}_c$는 항상 단일 벡터인 현재 전달된 입력 토큰의 쿼리 프로젝션입니다.

키-값 캐시를 사용하면 두 가지 장점이 있습니다.

전체 $\mathbf{QK}^T$ 행렬을 계산하는 것보다 더 적은 계산을 수행하므로 계산 효율이 크게 증가합니다. 이는 추론 속도의 증가로 이어집니다.
최대 필요 메모리는 생성된 토큰 수에 대해 이차적으로 증가하지 않고 선형적으로만 증가합니다.

더 긴 입력 시퀀스에 대해 동일한 결과와 상당한 속도 향상을 주므로 항상 키-값 캐시를 사용해야 합니다. Transformers는 텍스트 pipeline이나 generate 메서드를 사용할 때 키-값 캐시가 기본적으로 활성화되어 있습니다. 캐시 전용 가이드가 여기에 있습니다.

키-값 캐시를 사용하라고 조언했지만, 사용할 때 LLM 출력이 약간 다를 수 있다는 점에 유의하세요. 이는 행렬 곱셈 커널 자체의 속성입니다. 자세한 내용은 여기에서 읽을 수 있습니다.

3.2.1 다중 턴 대화

키-값 캐시는 여러 번의 자동 회귀 디코딩이 필요한 채팅 같은 애플리케이션에서 특히 유용합니다. 예를 살펴봅시다.

User: How many people live in France?
Assistant: Roughly 75 million people live in France
User: And how many are in Germany?
Assistant: Germany has ca. 81 million inhabitants

이 채팅에서 LLM은 자동 회귀 디코딩을 두 번 실행합니다.

첫 번째에 키-값 캐시는 비어 있고 입력 프롬프트는 "User: How many people live in France?"이며 모델은 매 디코딩 단계마다 키-값 캐시를 늘리면서 텍스트 "Roughly 75 million people live in France"를 자동 회귀적으로 생성합니다.
두 번째에 입력 프롬프트는 "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?"입니다. 캐시 덕분에 처음 두 문장에 대한 모든 키-값 벡터가 이미 계산되어 있습니다. 따라서 입력 프롬프트는 "User: And how many in Germany?"로만 구성됩니다. 단축된 입력 프롬프트를 처리하는 동안 계산된 키-값 벡터는 첫 번째 디코딩의 키-값 캐시에 연결됩니다. 그러면 두 번째 Assistant 응답 "Germany has ca. 81 million inhabitants"가 "User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many are in Germany?"의 인코딩된 키-값 벡터로 구성된 키-값 캐시로 자동 회귀적으로 생성됩니다.

여기서 두 가지를 주목해야 합니다.

모든 컨텍스트를 유지하는 것은 LLM이 대화의 모든 이전 컨텍스트를 이해해야 하므로 채팅에 배포된 LLM에 중요합니다. 예를 들어 위의 예에서 LLM은 사용자가 "And how many are in Germany"를 물을 때 인구를 말하는 것을 이해해야 합니다.
키-값 캐시는 인코딩된 채팅 기록을 처음부터 다시 인코딩하는 대신(예: encoder-decoder 아키텍처를 사용할 때) 계속 성장시킬 수 있게 해주므로 채팅에 매우 유용합니다.

transformers에서 generate 호출은 기본 use_cache=True 외에도 return_dict_in_generate=True가 전달되면 past_key_values를 반환합니다. 아직 pipeline 인터페이스에서는 사용할 수 없다는 점에 유의하세요.

# Generation as usual
prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True)
decoded_output = tokenizer.batch_decode(generation_output.sequences)[0]

# Piping the returned `past_key_values` to speed up the next conversation round
prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here"
model_inputs = tokenizer(prompt, return_tensors='pt')
generation_output = model.generate(
  **model_inputs,
  past_key_values=generation_output.past_key_values,
  max_new_tokens=60,
  return_dict_in_generate=True
)
tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

출력:

 is a modified version of the function that returns Mega bytes instead.

def bytes_to_megabytes(bytes):
   return bytes / 1024 / 1024

Answer: The function takes a number of bytes as input and returns the number of

좋아요, attention 레이어에 대해 동일한 키와 값을 다시 계산하는 데 추가 시간이 소요되지 않습니다! 그러나 한 가지 문제가 있습니다. $\mathbf{QK}^T$ 행렬에 필요한 피크 메모리는 크게 줄지만, 키-값 캐시를 메모리에 보유하는 것은 긴 입력 시퀀스나 다중 턴 채팅에서 매우 메모리 비용이 높을 수 있습니다. 키-값 캐시는 모든 self-attention 레이어와 모든 attention head에 대해 모든 이전 입력 벡터 $\mathbf{x}_i \text{, for } i \in \{1, \ldots, c - 1\}$의 키-값 벡터를 저장해야 한다는 것을 기억하세요.

이전에 사용한 LLM bigcode/octocoder의 키-값 캐시에 저장해야 하는 float 값의 수를 계산해 봅시다. float 값의 수는 시퀀스 길이 곱하기 attention head 수 곱하기 attention head 차원 곱하기 레이어 수의 두 배입니다. 가상의 입력 시퀀스 길이 16000에서 LLM에 대해 계산하면 다음과 같습니다.

config = model.config
2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

출력:

7864320000

대략 80억 개의 float 값입니다! 80억 개의 float 값을 float16 정밀도로 저장하는 데는 약 15GB의 RAM이 필요하며, 이는 모델 가중치 자체의 약 절반입니다! 연구자들은 키-값 캐시를 저장하는 메모리 비용을 크게 줄일 수 있는 두 가지 방법을 제안했으며, 다음 하위 섹션에서 살펴봅니다.

3.2.2 Multi-Query-Attention (MQA)

Multi-Query-Attention은 Noam Shazeer의 Fast Transformer Decoding: One Write-Head is All You Need 논문에서 제안되었습니다. 제목이 말하듯 Noam은 n_head 키-값 프로젝션 가중치를 사용하는 대신 모델의 성능이 크게 저하되지 않으면서 모든 attention head가 공유하는 단일 head-value 프로젝션 가중치 쌍을 사용할 수 있음을 발견했습니다.

단일 head-value 프로젝션 가중치 쌍을 사용하면 키 값 벡터 $\mathbf{k}_i, \mathbf{v}_i$가 모든 attention head에서 동일해야 하며, 이는 캐시에 n_head개 대신 1개의 키-값 프로젝션 쌍만 저장하면 된다는 것을 의미합니다.

대부분의 LLM이 20에서 100개의 attention head를 사용하므로 MQA는 키-값 캐시의 메모리 소비를 크게 줄입니다. 이 노트북에서 사용한 LLM의 경우 입력 시퀀스 길이 16000에서 필요한 메모리 소비를 15GB에서 400MB 미만으로 줄일 수 있습니다.

메모리 절약 외에도 MQA는 다음에서 설명하는 것처럼 계산 효율도 향상시킵니다. 자동 회귀 디코딩에서 큰 키-값 벡터를 매 단계마다 현재 키-값 벡터 쌍과 연결한 다음 $\mathbf{q}_c\mathbf{K}^T$ 계산에 공급해야 합니다. 자동 회귀 디코딩의 경우 지속적인 재로드에 필요한 메모리 대역폭이 심각한 시간 병목이 될 수 있습니다. 키-값 벡터의 크기를 줄이면 접근해야 하는 메모리가 줄어들어 메모리 대역폭 병목이 줄어듭니다. 자세한 내용은 Noam의 논문을 살펴보세요.

여기서 이해해야 할 중요한 점은 키-값 attention head 수를 1로 줄이는 것은 키-값 캐시가 사용될 때만 의미가 있다는 것입니다. 키-값 캐시가 없는 단일 forward pass의 모델 피크 메모리 소비는 각 attention head가 여전히 고유한 쿼리 벡터를 가지므로 각 attention head가 여전히 다른 $\mathbf{QK}^T$ 행렬을 가지기 때문에 변하지 않습니다.

MQA는 커뮤니티에서 널리 채택되었고 이제 가장 인기 있는 많은 LLM에서 사용됩니다.

또한 이 노트북에서 사용한 체크포인트 bigcode/octocoder도 MQA를 사용합니다.

3.2.3 Grouped-Query-Attention (GQA)

Google의 Ainslie et al.이 제안한 Grouped-Query-Attention은 MQA를 사용하면 일반적인 멀티 키-값 head 프로젝션을 사용하는 것보다 종종 품질 저하를 초래할 수 있음을 발견했습니다. 이 논문은 쿼리 head 프로젝션 가중치 수를 덜 급격하게 줄이면 더 많은 모델 성능을 유지할 수 있다고 주장합니다. 단일 키-값 프로젝션 가중치만 사용하는 대신 n < n_head 키-값 프로젝션 가중치를 사용해야 합니다. n을 2, 4 또는 8 같은 n_head보다 훨씬 작은 값으로 선택하면 MQA의 메모리와 속도 이점을 거의 모두 유지하면서 모델 용량을 덜 희생하므로 논란이 있지만 성능 손실이 더 적습니다.

또한 GQA의 저자들은 기존 모델 체크포인트를 원래 사전 훈련 계산의 5%만으로 GQA 아키텍처를 갖도록 *업트레이닝(uptrain)*할 수 있음을 발견했습니다. 원래 사전 훈련 계산의 5%는 여전히 엄청난 양일 수 있지만, GQA 업트레이닝을 통해 기존 체크포인트를 더 긴 입력 시퀀스에 유용하게 만들 수 있습니다.

GQA는 최근에야 제안되었기 때문에 이 노트북을 작성하는 시점에는 채택이 덜 되었습니다. GQA의 가장 주목할 만한 적용은 Llama-v2입니다.

결론적으로, LLM이 자동 회귀 디코딩으로 배포되고 채팅의 경우처럼 큰 입력 시퀀스를 처리해야 한다면 GQA 또는 MQA 중 하나를 사용하는 것이 강력히 권장됩니다.

결론

연구 커뮤니티는 점점 더 커지는 LLM의 추론 시간을 가속화하는 새롭고 영리한 방법을 끊임없이 고안하고 있습니다. 예를 들어, 유망한 연구 방향 중 하나는 speculative decoding으로, "쉬운 토큰"은 더 작고 빠른 언어 모델이 생성하고 "어려운 토큰"만 LLM 자신이 생성합니다. 더 자세히 들어가는 것은 이 노트북의 범위를 벗어나지만, 이 훌륭한 블로그 게시물에서 읽을 수 있습니다.

GPT3/4, Llama-2-70b, Claude, PaLM 같은 대규모 LLM이 Hugging Face Chat이나 ChatGPT 같은 채팅 인터페이스에서 그렇게 빠르게 실행될 수 있는 이유는 크게 위에서 언급한 정밀도, 알고리즘, 아키텍처 개선 덕분입니다. 앞으로 GPU, TPU 같은 가속기는 더 빨라지고 더 많은 메모리를 허용하겠지만, 항상 최고의 알고리즘과 아키텍처를 사용하여 최상의 결과를 얻도록 해야 합니다 🤗

속도와 메모리를 위한 LLM 최적화