Transformers documentation
속도와 메모리를 위한 LLM 최적화
속도와 메모리를 위한 LLM 최적화
GPT3/4, Falcon, Llama와 같은 대규모 언어 모델(LLM)은 인간 중심의 작업을 처리하는 능력이 빠르게 발전하면서 현대 지식 기반 산업에서 필수적인 도구로 자리 잡고 있습니다. 그러나 이러한 모델을 실제 작업에 배포하는 것은 여전히 도전적입니다:
- 거의 인간 수준의 텍스트 이해 및 생성 능력을 갖추기 위해, LLM은 현재 수십억 개의 파라미터로 구성되어야 합니다(자세한 내용은 Kaplan et al, Wei et. al 참고). 이로 인해 추론을 위한 메모리 요구량이 크게 증가합니다.
- 많은 실제 작업에서 LLM에는 방대한 컨텍스트 정보가 제공되어야 합니다. 이를 위해 모델은 추론 단계에서 매우 긴 입력 시퀀스를 처리할 수 있는 능력이 필요합니다.
이러한 과제들의 핵심은 LLM의 계산 및 메모리 처리 능력을 향상시키는 데 있으며, 특히 방대한 입력 시퀀스를 처리할 때 그렇습니다.
이 가이드에서는 효율적인 LLM 배포를 위한 효과적인 기법들을 살펴봅니다:
낮은 정밀도: 연구에 따르면, 8비트 및 4비트와 같이 수치 정밀도를 낮추어 작동하면 모델 성능의 상당한 저하 없이 계산상의 이점을 얻을 수 있습니다.
Flash Attention: Flash Attention은 메모리 효율적인 접근 방식을 제공할 뿐만 아니라 최적화된 GPU 메모리 활용도를 통해 효율성을 높이는 attention 알고리즘의 변형입니다.
아키텍처 혁신: LLM은 추론 시 항상 동일한 방식으로 배포되는데, 이는 긴 입력 컨텍스트를 가진 자기회귀 텍스트 생성입니다. 이러한 맥락에서 더 효율적인 추론을 가능하게 하는 특화된 모델 아키텍처가 제안되었습니다. 모델 아키텍처 측면에서 가장 중요한 발전은 Alibi, Rotary embeddings, Multi-Query Attention (MQA), Grouped-Query-Attention (GQA)입니다.
이 가이드에서는 텐서 관점에서 오토레거시(autoregressive) 생성을 분석합니다. 낮은 정밀도를 채택하는 것의 장단점을 깊이 있게 다루고, 최신 어텐션 알고리즘을 종합적으로 살펴본 뒤, 개선된 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 라이브러리를 참고해 보세요.
Naive pipeline parallelism은 기본적으로 지원됩니다. 이를 위해 device="auto"로 모델을 로드하면 여기에서 설명한 대로 서로 다른 레이어가 사용 가능한 GPU에 자동으로 배치됩니다. 그러나 매우 효과적이긴 하지만, 이 naive pipeline parallelism은 GPU가 유휴 상태(idling)로 있는 문제를 해결하지 못합니다. 이를 해결하려면 여기에서 설명한 대로 더 고급의 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"를 사용하면 어텐션 레이어가 사용 가능한 모든 GPU에 균등하게 분배됩니다.
이 가이드에서는 bigcode/octocoder를 사용합니다. 이 모델은 단일 40 GB A100 GPU 장치 칩에서 실행할 수 있기 때문입니다. 참고로, 앞으로 적용할 모든 메모리 및 속도 최적화는 모델 또는 텐서 병렬성이 필요한 모델에도 동일하게 적용됩니다.
모델이 bfloat16 정밀도로 로드되므로, 위에서 언급한 경험칙에 따라 bigcode/octocoder로 추론을 실행하는 데 필요한 메모리는 약 31 GB 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 정밀도로 실행하려고 했다면, 무려 64 GB의 VRAM이 필요했을 것입니다.
요즘 거의 모든 모델이 bfloat16로 학습되므로, GPU가 bfloat16를 지원하는 경우 모델을 float32 정밀도로 실행할 이유가 없습니다. Float32는 모델 학습에 사용된 정밀도보다 더 나은 추론 결과를 제공하지 않습니다.
Hub에 모델 가중치가 어떤 형식으로 저장되어 있는지 확실하지 않다면, 항상 체크포인트의 config에서 "dtype" 항목을 확인하면 됩니다. 예를 들어 여기와 같이요. from_pretrained(..., dtype=...)로 모델을 로드할 때는 config에 명시된 것과 동일한 정밀도 유형으로 설정하는 것이 권장되며, 원래 유형이 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에 32 GB의 VRAM이 없다면 어떻게 될까요? 모델 가중치를 8-bit 또는 4-bit로 양자화해도 성능에 유의미한 손실 없이 사용할 수 있음이 밝혀졌습니다(자세한 내용은 Dettmers et al. 참고). 최근 GPTQ paper에서 보았듯이, 모델은 허용 가능한 수준의 성능 손실로 3비트 또는 2비트까지 양자화할 수 있습니다 🤯.
자세한 내용은 생략하고, 양자화 스킴은 모델의 추론 결과가 가능한 한 정확하도록(a.k.a bfloat16에 가능한 한 가깝도록) 가중치의 정밀도를 낮추는 것을 목표로 합니다.
양자화는 텍스트 생성에 특히 효과적입니다. 왜냐하면 우리가 신경 쓸 것은 가장 가능성이 높은 다음 토큰의 집합을 선택하는 것뿐이며, 다음 토큰 logit 분포의 정확한 값에는 크게 신경 쓰지 않기 때문입니다.
중요한 것은 다음 토큰 logit 분포가 대략적으로 동일하게 유지되어 argmax 또는 topk 연산이 동일한 결과를 내는 것입니다.
여러 가지 양자화 기법이 있으며, 여기서는 자세히 다루지 않겠지만, 일반적으로 모든 양자화 기법은 다음과 같이 작동합니다:
- 모든 가중치를 대상 정밀도로 양자화
- quantized 가중치를 로드하고, bfloat16 정밀도의 벡터 입력 시퀀스를 전달합니다
- bfloat16 정밀도로 입력 벡터와 계산을 수행하기 위해 가중치를 동적으로 bfloat16으로 역양자화합니다
요약하자면, 이는 입력-가중치 행렬 곱셈을 의미하며, 은(는) 입력, 은(는) 가중치 행렬이고 은(는) 출력입니다:
은(는) 다음으로 변경됩니다.
모든 행렬 곱셈에 대해 수행됩니다. 입력이 네트워크 그래프를 통과하는 동안 모든 가중치 행렬에 대해 역양자화 및 재양자화가 순차적으로 수행됩니다.
따라서 양자화된 가중치를 사용할 때 추론 시간이 줄어들지 않고 오히려 증가하는 경우가 많습니다.
이론은 여기까지 하고 직접 시도해 보겠습니다. Transformers로 가중치를 양자화하려면 bitsandbytes 라이브러리가 설치되어 있는지 확인해야 합니다.
!pip install bitsandbytes
이후 from_pretrained에 load_in_8bit=True 플래그를 추가하기만 하면 8-bit 양자화로 모델을 로드할 수 있습니다.
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
훨씬 적습니다! 이제 15 GB를 약간 상회하는 수준까지 줄어들었으므로, 4090과 같은 소비자용 GPU에서도 이 모델을 실행할 수 있습니다. 메모리 효율성이 크게 향상되었고, 모델 출력의 품질 저하가 거의 없습니다. 다만 추론 과정에서 약간의 지연이 발생함을 확인할 수 있습니다.
모델을 삭제하고 메모리를 다시 정리합니다.
del model
del pipeflush()
4-bit 양자화가 얼마나 GPU 메모리 사용량을 줄이는지 확인해 보겠습니다. 모델을 4-bit로 양자화하는 방법은 이전과 동일한 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-bit 양자화는 8-bit 양자화나 전체 bfloat16 추론과 다른 결과를 초래할 수 있습니다. 사용자는 직접 시도해 보시기 바랍니다.
또한, 4-bit 양자화에는 더 공격적인 양자화 방법이 사용되므로 추론 속도가 8-bit 양자화보다 다소 느린 점도 참고하세요. 이는 및 가 추론 중에 더 오래 걸리기 때문입니다.
del model
del pipeflush()
전체적으로, OctoCoder를 8-bit 정밀도로 실행하면 필요한 GPU VRAM이 32GB에서 15GB로 줄었고, 모델을 4-bit 정밀도로 실행하면 필요한 GPU VRAM이 9GB를 약간 넘는 수준까지 추가로 줄어든다는 것을 확인했습니다.
4-bit 양자화를 사용하면 대부분의 사용자가 쉽게 접근할 수 있는 RTX3090, V100, T4와 같은 GPU에서 모델을 실행할 수 있습니다.
양자화에 대한 자세한 정보와 4비트보다 더 적은 GPU VRAM 메모리를 요구하도록 모델을 양자화하는 방법을 확인하려면 GPT-QModel 구현을 살펴보는 것을 권장합니다.
결론적으로, 모델 양자화는 메모리 효율성을 향상시키는 대신 정확도 및 일부 경우 추론 시간을 희생한다는 점을 기억하는 것이 중요합니다.
GPU 메모리가 사용 사례에서 제약 조건이 아니라면 양자화를 살펴볼 필요가 없는 경우가 많습니다. 그러나 많은 GPU는 양자화 방법 없이는 LLM을 실행할 수 없으며, 이 경우 4비트 및 8비트 양자화 방식은 매우 유용한 도구입니다.
자세한 사용법에 대해 더 알고 싶으시다면 Transformers Quantization Docs를 참고하시는 것을 강력히 권장합니다. 다음으로, 더 나은 알고리즘과 개선된 모델 아키텍처를 사용하여 계산 및 메모리 효율성을 어떻게 향상시킬 수 있는지 살펴보겠습니다.
2. Flash Attention
오늘날 최고의 성능을 내는 LLM들은 대부분 피드포워드 레이어, 활성화 레이어, 레이어 정규화 레이어, 그리고 가장 중요한 셀프-어텐션 레이어로 구성된 동일한 기본 아키텍처를 공유합니다.
Self-attention 레이어는 입력 토큰 간의 문맥적 관계를 이해할 수 있게 해주므로 Large Language Models(LLM)의 핵심입니다. 그러나 self-attention 레이어의 GPU 메모리 사용량 피크는 제곱적으로 증가하며, 이는 계산 및 메모리 복잡성 모두에서 입력 토큰 수(또는 시퀀스 길이) 라고도 하며, 이후 . 로 표기합니다)에 따라 발생합니다. 이는 짧은 입력 시퀀스(입력 토큰 1000개 이하)에서는 크게 두드러지지 않지만, 긴 입력 시퀀스(입력 토큰 약 16000개)에서는 심각한 문제가 됩니다.
자세히 살펴보겠습니다. 입력 에 대한 self-attention 레이어의 출력 의 길이가 일 때, 출력을 계산하는 공식은 다음과 같습니다:
은(는) attention layer의 입력 시퀀스가 됩니다. 투영 과(와) 는 각각 개의 벡터로 구성되며, 그 결과 의 크기는 .
LLM은 일반적으로 여러 개의 attention head를 가지며, 이를 통해 여러 self-attention 계산을 병렬로 수행합니다. LLM이 40개의 attention head를 가지고 bfloat16 정밀도로 실행된다고 가정하면, 다음을 저장하는 데 필요한 메모리 용량을 계산할 수 있습니다. 행렬은 바이트입니다. 인 경우 VRAM 약 50MB만 필요하지만, 인 경우 VRAM 19GB가 필요하고, 인 경우 다음을 저장하는 데만 거의 1TB가 필요합니다. 행렬입니다.
요약하자면, 기본 self-attention 알고리즘은 큰 입력 컨텍스트에서 메모리 비용이 급격히 증가하여 사용이 불가능해집니다.
LLM의 텍스트 이해 및 생성 능력이 향상됨에 따라, 점점 더 복잡한 작업에 적용되고 있습니다. 과거에는 몇 문장의 번역이나 요약만 처리하던 모델이 이제는 전체 페이지를 다루며, 방대한 입력 길이를 처리할 수 있는 능력이 요구되고 있습니다.
긴 입력 길이에 따른 과도한 메모리 요구 사항을 어떻게 제거할 수 있을까요? 행렬을 없애는 새로운 self-attention 메커니즘 계산 방식이 필요합니다. Tri Dao 외 이러한 새로운 알고리즘을 정확히 개발하여 Flash Attention.
요약하자면, Flash Attention은 계산을 분리하고, 대신 여러 softmax 계산 단계를 반복하여 출력의 더 작은 청크를 계산합니다:
여기서 그리고 는 매번 다시 계산해야 하는 softmax 정규화 통계량이며 입니다. .
Flash Attention 전체는 다소 복잡하며, 이 가이드의 범위를 벗어나기 때문에 여기서는 크게 단순화되어 있습니다. 자세한 내용은 잘 작성된 Flash Attention 논문을 참고하시기 바랍니다.
여기서 핵심은 다음과 같습니다:
softmax 정규화 통계를 추적하고 몇 가지 정교한 수학적 기법을 활용함으로써 Flash Attention은 수치적으로 동일한 출력을 제공하며, 메모리 비용은 .
수식을 보면, 더 많은 계산이 필요하므로 Flash Attention이 기본 self-attention 수식보다 훨씬 느릴 것이라고 직관적으로 말할 수 있습니다. 실제로 softmax 정규화 통계가 지속적으로 재계산되어야 하므로 Flash Attention은 일반 attention보다 더 많은 FLOPs가 필요합니다(관심 있으시면 논문에서 자세한 내용을 확인하세요).
그러나 Flash Attention은 GPU의 느리고 대역폭이 높은 메모리(VRAM)에 대한 요구를 크게 줄이고, 대신 칩 온보드 메모리(SRAM)에 집중할 수 있는 능력 덕분에 기본 attention 방식에 비해 추론 시 훨씬 빠릅니다.
결론적으로, Flash Attention은 출력 벡터를 계산하기 위해 느린 VRAM 메모리에 접근해야 하는 대신, 모든 중간 쓰기 및 읽기 연산을 빠른 온칩(on-chip) SRAM 메모리를 사용하여 수행할 수 있도록 보장합니다. .
실제로, Flash Attention이 사용 가능한 경우 이를 사용하지 않을 이유는 전혀 없습니다. 이 알고리즘은 수학적으로 동일한 출력을 제공하며, 더 빠르고 메모리 효율적입니다.
3. 아키텍처 혁신
지금까지 다음을 통해 계산 및 메모리 효율성을 개선하는 방법에 대해 살펴보았습니다:
- weights를 더 낮은 정밀도 형식으로 캐스팅
- 메모리와 연산 효율이 더 높은 버전으로 self-attention 알고리즘을 대체
이제 긴 텍스트 입력이 필요한 작업에서 가장 효과적이고 효율적이도록 LLM의 아키텍처를 어떻게 변경할 수 있는지 살펴봅시다. 예를 들어:
- 검색 증강 질의응답,
- 요약,
- 채팅
chat는 LLM이 긴 텍스트 입력을 처리할 수 있을 뿐만 아니라, 사용자와 어시스턴트 간의 왕복 대화(예: ChatGPT)를 효율적으로 처리할 수 있는 능력도 요구한다는 점에 유의하세요.
학습이 완료된 후, 기본 LLM 아키텍처를 변경하기가 어렵기 때문에 LLM의 작업에 대해 사전에 고려하고 이에 따라 모델 아키텍처를 최적화하는 것이 중요합니다. 모델 아키텍처에는 긴 입력 시퀀스에서 빠르게 메모리 및/또는 성능 병목 현상을 유발하는 두 가지 중요한 구성 요소가 있습니다.
- 위치 임베딩
- 키-값 캐시
각 구성 요소를 자세히 살펴보겠습니다
3.1 LLM의 위치 임베딩 개선
Self-attention은 각 토큰을 다른 토큰들과의 관계로 표현합니다. 예를 들어, 텍스트 입력 시퀀스의 ““Hello”, “I”, “love”, “you” 행렬은 다음과 같은 형태일 수 있습니다:

각 단어 토큰은 다른 모든 단어 토큰에 주의를 기울일 확률 질량을 가지며, 따라서 다른 모든 단어 토큰과 관계를 맺게 됩니다. 예를 들어, 단어 “love”는 단어 “Hello”에 5%, “I”에 30%, 그리고 자신에게 65%의 주의를 기울입니다.
A LLM based on self-attention, but without position embeddings would have great difficulties in understanding the positions of the text inputs to each other. This is because the probability score computed by relates each word token to each other word token in computations regardless of their relative positional distance to each other. Therefore, for the LLM without position embeddings each token appears to have the same distance to all other tokens, e.g. differentiating between “Hello I love you” and “You love I hello” would be very challenging.
LLM이 문장 순서를 이해하려면 추가적인 큐(cue)가 필요하며, 이는 일반적으로 위치 인코딩(positional encodings)(또는 위치 임베딩(positional embeddings)이라고도 함)의 형태로 적용됩니다. 위치 인코딩은 각 토큰의 위치를 LLM이 문장 순서를 더 잘 이해할 수 있도록 활용할 수 있는 수치적 표현으로 인코딩합니다.
논문 Attention Is All You Need 의 저자들은 사인 위치 임베딩을 도입했습니다 . 여기서 각 벡터 는 해당 위치 . 에 대한 사인 함수로 계산됩니다. 위치 인코딩은 단순히 입력 시퀀스 벡터에 더해집니다 = 이를 통해 모델이 문장 순서를 더 잘 학습하도록 유도합니다.
고정된 위치 임베딩을 사용하는 대신, 다른 연구자들(예: Devlin et al.) )은 학습된 위치 인코딩을 사용하여, 위치 임베딩 이 학습 과정에서 학습되도록 했습니다.
사인 함수 기반 위치 임베딩과 학습된 위치 임베딩은 LLM에 문장 순서를 인코딩하는 데 주로 사용되었지만, 이러한 위치 인코딩과 관련된 몇 가지 문제가 발견되었습니다:
- 사인 파형 위치 임베딩과 학습된 위치 임베딩은 모두 절대 위치 임베딩이며, i.e. 각 위치 id에 고유한 임베딩을 인코딩합니다: . 에서 보았듯이 Huang et al. 그리고 Su et al., 절대 위치 임베딩은 긴 텍스트 입력에서 LLM의 성능을 저하시킵니다. 긴 텍스트 입력의 경우, 모델이 토큰의 절대 위치가 아니라 서로 간의 상대적 위치 거리를 학습하는 것이 유리합니다.
- 학습된 위치 임베딩을 사용할 경우, LLM은 고정된 입력 길이 $N$으로 학습되어야 하며, 이는 학습된 길이보다 긴 입력 길이에 대한 외삽을 어렵게 만듭니다.
최근, 위에서 언급한 문제들을 해결할 수 있는 상대 위치 임베딩이 더 널리 사용되고 있으며, 특히 주목할 만한 것은 다음과 같습니다:
둘 다 RoPE 와 ALiBi 는 단어 토큰이 서로 관계를 맺는 곳이 바로 self-attention 알고리즘이므로, 문장 순서를 LLM에 직접 전달하는 것이 가장 좋다고 주장합니다. 더 구체적으로 말하면, 문장 순서는 계산을 수정하여 전달해야 합니다.
세부 사항에 깊이 들어가지 않고, RoPE 는 위치 정보를 쿼리-키 쌍에 인코딩할 수 있다고 언급합니다. e.g. 각 벡터를 각도 로 회전시켜 각각 으로 회전시키며, 각 벡터의 문장 위치를 나타냅니다:
이는 회전 행렬을 나타냅니다. 은 학습 되지 않으며, 대신 학습 시 최대 입력 시퀀스 길이에 따라 미리 정의된 값으로 설정됩니다.
이렇게 하면 과 사이의 확률 점수는 일 때에만 영향을 받으며, 각 벡터의 구체적인 위치와 관계없이 상대적 거리 와 에만 의존합니다. .
RoPE는 다음과 같은 오늘날 가장 중요한 LLM 중 여러 개에서 사용됩니다:
대안으로, ALiBi 는 훨씬 더 간단한 상대 위치 인코딩 방식을 제안합니다. 입력 토큰 간의 상대적 거리를 사전에 정의된 값 m 으로 스케일링한 음의 정수를 행렬에 softmax 계산 직전에 각 query-key 항목에 더합니다.

ALiBi 논문에서 보여준 바와 같이, 이 단순한 상대 위치 인코딩은 매우 긴 텍스트 입력 시퀀스에서도 모델이 높은 성능을 유지할 수 있게 합니다.
ALiBi는 다음과 같은 오늘날 가장 중요한 LLM 중 여러 개에서 사용됩니다:
둘 다 RoPE 및 ALiBi 위치 인코딩은 학습 중에 보지 못한 입력 길이로 외삽할 수 있으며, ALiBi 는 RoPE. 에 비해 기본적으로 외삽이 훨씬 잘 작동하는 것으로 나타났습니다. ALiBi의 경우, 입력 시퀀스의 길이에 맞게 하삼각 위치 행렬의 값을 단순히 증가시키면 됩니다. RoPE, 학습 시 사용했던 을 그대로 사용하면 학습 시보다 훨씬 긴 텍스트 입력을 전달할 때 결과가 좋지 않습니다. c.f Press et al.. 그러나 커뮤니티는 $\theta$를 조정하여 RoPE 위치 임베딩이 외삽된 텍스트 입력 시퀀스에서 잘 작동하도록 하는 몇 가지 효과적인 트릭을 발견했습니다( 여기).
RoPE와 ALiBi 모두 학습 과정에서 학습되지 않는 상대 위치 임베딩이며, 다음 직관에 기반합니다:
- 텍스트 입력에 대한 위치 정보는 self-attention 레이어의 행렬에 직접 주어져야 합니다.
- LLM은 일정한 상대 거리 위치 인코딩을 학습하도록 유도되어야 합니다.
- 입력 토큰끼리 거리가 멀수록, 해당 토큰들의 query-value 확률은 낮아집니다. RoPE와 ALiBi 모두 서로 멀리 떨어진 토큰들의 query-key 확률을 낮춥니다. RoPE는 query-key 벡터 사이의 각도를 증가시켜 벡터 내적을 감소시킴으로써 확률을 낮추고, ALiBi는 벡터 내적에 큰 음수를 더함으로써 확률을 낮춥니다.
요약하자면, 대량의 텍스트 입력을 처리해야 하는 작업에 배포할 LLM은 RoPE나 ALiBi와 같은 상대 위치 임베딩으로 학습하는 것이 더 좋습니다. 또한 RoPE와 ALiBi를 사용한 LLM이 예를 들어 과 같은 고정된 길이로만 학습되었더라도, 위치 임베딩을 외삽하여 $N_1$보다 훨씬 큰 텍스트 입력, 예를 들어 $N_2 = 8192 > N_1$을 실제 환경에서 사용할 수 있습니다.
3.2 키-값 캐시
LLM을 사용한 자기 회귀적 텍스트 생성은 입력 시퀀스를 넣고, 다음 토큰을 샘플링한 후, 해당 토큰을 입력 시퀀스에 추가하는 과정을 반복하여 수행됩니다. 이 과정은 LLM이 생성이 종료되었음을 나타내는 토큰을 생성할 때까지 계속됩니다.
자동 회귀 생성이 어떻게 작동하는지 더 시각적으로 이해하려면 Transformer’s Generate Text Tutorial을 참고하세요.
auto-regressive가 실제로 어떻게 작동하는지 보여주는 간단한 코드 스니펫을 실행해 보겠습니다. 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']
text input tokens가 매번 방금 샘플링된 토큰만큼 증가하는 것을 확인할 수 있습니다.
매우 적은 예외를 제외하고, LLM은 causal language modeling objective를 사용하여 학습되며, 따라서 attention score의 상삼각 행렬을 mask합니다. 이는 위의 두 다이어그램에서 attention score가 비어 있는(즉, a.k.a 확률이 0인) 이유입니다. causal language modeling에 대한 간단한 복습을 원하시면 Illustrated Self Attention blog를 참고하세요.
그 결과 토큰은 절대 나중 토큰에 의존하지 않으며, 더 구체적으로 벡터는 인 경우 어떤 키, 값 벡터와도 관련이 되지 않습니다. 대신 는 이전 키-값 벡터 $\mathbf{k}${m < i}, \mathbf{v}${m < i} \text{ , for } m \in {0, \ldots i - 1}$에만 주의를 기울입니다. 불필요한 계산을 줄이기 위해 각 레이어의 모든 이전 타임스텝에 대한 키-값 벡터를 캐싱할 수 있습니다.
다음에서는 LLM이 각 forward pass마다 key-value cache를 가져와 전달하도록 지시합니다.
Transformers에서는 forward 호출에 use_cache 플래그를 전달하여 key-value cache를 가져올 수 있으며, 이후 이를 현재 토큰과 함께 전달할 수 있습니다.
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']
알 수 있듯이, key-value 캐시를 사용할 때 텍스트 입력 토큰의 길이는 늘지 않으며 단일 입력 벡터로 유지됩니다. 반면, key-value 캐시의 길이는 디코딩 단계마다 1씩 증가합니다.
키-값 캐시를 활용하면 은 기본적으로 로 축소되며 은 현재 전달된 입력 토큰의 쿼리 투영으로 항상 단일 벡터입니다.
key-value cache를 사용하면 두 가지 장점이 있습니다:
- 전체 행렬을 계산하는 것보다 수행되는 연산이 적어 계산 효율이 크게 향상됩니다. 이로 인해 추론 속도가 빨라집니다.
- 생성된 토큰 수에 따라 필요한 최대 메모리가 제곱으로 증가하지 않고 선형으로만 증가합니다.
키-값 캐시를 항상 사용해야 합니다. 이는 동일한 결과를 보장하고 긴 입력 시퀀스에서 상당한 속도 향상을 가져오기 때문입니다. Transformers는 텍스트 파이프라인 또는
generate메서드를 사용할 때 기본적으로 키-값 캐시를 활성화합니다. 캐시에 전념한 전체 가이드를 여기에서 확인할 수 있습니다.
키-밸류 캐시를 사용하도록 권장했음에도 불구하고, 캐시를 사용할 때 LLM의 출력이 약간 달라질 수 있습니다. 이는 행렬 곱셈 커널 자체의 특성으로, 자세한 내용은 여기에서 확인할 수 있습니다.
3.2.1 다중 라운드 대화
key-value cache는 여러 번의 auto-regressive decoding이 필요한 채팅과 같은 애플리케이션에서 특히 유용합니다. 예를 들어 살펴보겠습니다.
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은 오토레그레시브 디코딩을 두 번 실행합니다:
- 처음에는 key-value 캐시가 비어 있으며, 입력 프롬프트는
"User: How many people live in France?"이고, 모델은 각 디코딩 단계마다 key-value 캐시를 늘리면서"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"라고 물을 때 인구(population)를 언급하고 있다는 점을 이해해야 합니다. - key-value cache는 인코더-디코더 아키텍처를 사용하는 경우처럼 매번 처음부터 채팅 히스토리를 다시 인코딩해야 하는 대신, 인코딩된 채팅 히스토리를 지속적으로 확장할 수 있게 해주므로 채팅에 매우 유용합니다.
transformers에서 generate 호출 시 return_dict_in_generate=True를 전달하면 기본값인 use_cache=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 layer에서 동일한 key와 values를 다시 계산하는 데 추가 시간이 소요되지 않습니다! 다만 한 가지 주의할 점이 있습니다. 행렬에 필요한 피크 메모리는 크게 줄지만, 긴 입력 시퀀스나 멀티턴 챗의 경우 key-value 캐시를 메모리에 유지하는 것이 매우 많은 메모리를 차지할 수 있습니다. key-value 캐시는 모든 이전 입력 벡터 모든 self-attention layer와 모든 attention head의 key-value 벡터를 저장해야 한다는 점을 기억하세요.
이전에 사용했던 LLM bigcode/octocoder의 key-value 캐시에 저장해야 하는 float 값의 수를 계산해 보겠습니다.
float 값의 수는 시퀀스 길이의 2배에 어텐션 헤드 수, 어텐션 헤드 차원, 레이어 수를 곱한 값입니다.
가상의 입력 시퀀스 길이 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 정밀도로 저장하려면 약 15 GB의 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개의 key-value projection 가중치를 사용하는 대신 모든 attention head에서 공유되는 단일 head-value projection 가중치 쌍을 사용해도 모델의 성능이 크게 저하되지 않음을 발견했습니다.
단일 head-value projection weight 쌍을 사용하면 key value 벡터 는 모든 attention head에서 동일해야 하며, 이는 캐시에
n_head개의 key-value projection 쌍 대신 1개만 저장하면 된다는 것을 의미합니다.
대부분의 LLM은 20개에서 100개 사이의 attention head를 사용하므로, MQA는 key-value cache의 메모리 사용량을 크게 줄여 줍니다. 따라서 이 notebook에서 사용한 LLM의 경우, 입력 시퀀스 길이가 16000일 때 필요한 메모리 사용량을 15 GB에서 400 MB 미만으로 줄일 수 있습니다.
메모리 절감 외에도, MQA는 다음에서 설명하듯 계산 효율성도 향상시킵니다. 오토레거시 디코딩에서는 큰 키-밸류 벡터를 다시 로드하여 현재 키-밸류 벡터 쌍과 연결한 뒤 의 계산에 입력해야 합니다. 오토레거시 디코딩에서 이 상시 재로딩에 필요한 메모리 대역폭이 심각한 시간 병목이 될 수 있습니다. 키-밸류 벡터의 크기를 줄이면 접근해야 하는 메모리 양이 감소하여 메모리 대역폭 병목이 완화됩니다. 자세한 내용은 Noam의 논문.
여기서 이해해야 할 중요한 점은, key-value attention head의 수를 1로 줄이는 것이 key-value cache를 사용하는 경우에만 의미가 있다는 것입니다. key-value cache 없이 단일 forward pass를 수행할 때 모델의 피크 메모리 사용량은 모든 attention head가 고유한 query 벡터를 가지므로 각 attention head가 여전히 서로 다른 행렬을 가지기 때문에 변하지 않습니다.
MQA는 커뮤니티에서 널리 채택되어 현재 가장 인기 있는 LLM 중 많은 모델에서 사용되고 있습니다:
또한, 이 노트북에서 사용된 체크포인트 bigcode/octocoder는 MQA를 사용합니다.
3.2.3 그룹 쿼리 어텐션(GQA)
Google의 Ainslie et al.이 제안한 Grouped-Query-Attention은, 일반적인 멀티 키-밸류 헤드 프로젝션을 사용하는 것에 비해 MQA를 사용하면 종종 품질 저하가 발생할 수 있음을 발견했습니다. 이 논문은 쿼리 헤드 프로젝션 가중치의 수를 덜 극단적으로 줄임으로써 더 많은 모델 성능을 유지할 수 있다고 주장합니다. 단일 키-밸류 프로젝션 가중치를 사용하는 대신, n < n_head개의 키-밸류 프로젝션 가중치를 사용해야 합니다. n을 n_head보다 훨씬 작은 값, 예를 들어 2, 4 또는 8로 선택하면 MQA에서 얻을 수 있는 메모리 및 속도 향상 대부분을 유지하면서 모델 용량을 덜 희생하므로, 논쟁의 여지가 있겠지만 성능 저하도 줄일 수 있습니다.
또한, GQA의 저자들은 기존 모델 체크포인트를 원래 사전 학습 연산량의 5%만 사용하여 GQA 아키텍처를 갖도록 업트레이닝(uptrained)할 수 있음을 발견했습니다. 원래 사전 학습 연산량의 5%가 여전히 상당한 양일 수 있지만, GQA 업트레이닝(uptraining)은 기존 체크포인트가 더 긴 입력 시퀀스에 유용하도록 합니다.
GQA는 비교적 최근에 제안된 기법으로, 이 노트북 작성 시점에서는 도입 사례가 적습니다. GQA의 가장 대표적인 적용 사례는 Llama-v2입니다.
결론적으로, LLM이 자동 회귀 디코딩(auto-regressive decoding)으로 배포되고 채팅과 같이 긴 입력 시퀀스를 처리해야 하는 경우라면 GQA 또는 MQA를 사용하는 것이 강력히 권장됩니다.
결론
연구 커뮤니티는 점점 더 커지는 LLM의 추론 시간을 단축하기 위한 새로운 방법들을 끊임없이 개발하고 있습니다. 예를 들어, speculative decoding와 같은 유망한 연구 방향 중 하나는 더 작고 빠른 언어 모델로 “쉬운 토큰”을 생성하고, LLM 자체로는 “어려운 토큰”만 생성하는 방식입니다. 세부 사항은 이 노트북의 범위를 벗어나지만, 이 좋은 블로그 게시물에서 자세히 읽을 수 있습니다.
GPT3/4, Llama-2-70b, Claude, PaLM과 같은 대규모 LLM이 Hugging Face Chat이나 ChatGPT와 같은 채팅 인터페이스에서 매우 빠르게 실행되는 이유는 앞서 언급한 정밀도, 알고리즘, 아키텍처의 개선 덕분인 부분이 큽니다. 앞으로 GPU, TPU와 같은 가속기는 더욱 빨라지고 더 많은 메모리를 허용할 것이지만, 그럼에도 불구하고 최고의 성능을 얻기 위해 항상 사용 가능한 최적의 알고리즘과 아키텍처를 사용하는지 확인해야 합니다 🤗
Update on GitHub