Text Generation Inference(TGI)를 활용한 프로덕션 레벨 LLM 추론 가속화
Text Generation Inference(TGI)는 Hugging Face의 Inference API과 Hugging Chat에서 LLM 추론 성능을 강화하는 내부 툴로서 2023년 초부터 개발이 시작되었다. 이후 프로덕션에서의 추론을 위한 관리형 솔루션인 Inference Endpoint 뿐만 아니라, Amazon Sagemaker Hugging Face LLM Inference DLC 등 여러 CSP의 ML 서비스에서 추론 성능을 극대화하기 위해 내장되며, 중요한 내부 구성 요소로서 자리잡게 되었다. FLAN-T5, Llama, Falcon 등 인기 있는 여러 모델에 대해 최적화되어 있으며, 메타의 Llama 2 또한 최적화된 리스트에 포함되었다. 이번 포스팅에서는 TGI가 어떻게 추론 성능을 개선할 수 있는지 각 특징에 대해서 쉽게 풀어보도록 하겠다.
Tensor Parallelism
Tensor Parallelism은 weights, gradient 그리고 optimizer 상태와 같은 텐서를 여러 chunk로 분산하여 다중 GPU 장비에서 분산 저장하고 처리되는 것을 의미한다. Nvidia의 연구진들이 발표한 MegaTron-LM 논문에서 매우 큰 트랜스포머 모델을 훈련하기 위한 테크닉을 설명하며 이러한 개념이 처음 소개되었다.
A(가중치 행렬)을 다중 GPU에서 column 방향으로 분할한 이후 행렬곱을 병렬로 수행하면 XA_n 벡터가 n개 만큼 생성된다. 이는 row 방향으로 쪼개더라도 동일한 출력 벡터 Y로 반환된다.
Megatron-LM 논문에서는 self-attention를 병렬로 다수번 학습시키는 Multi-Head Attention에서도 출력 벡터를 취합하는 마지막까지 GPU 간 동기화 없이 병렬 학습을 지원함을 확인하였다. Q, K, V 행렬이 본질적으로 병렬 형태이기 때문에 더욱 분산 처리에 용이한 구조라고 볼 수 있다.
Quantization with bitsandbytes or gptq
Quantization(양자화)는 실수형 변수를 정수형 변수로 변환하는 과정을 말한다. 주로 딥러닝에서 FP32(32-bit floating point)는 사용하여 가중치를 표현하는데, 이를 FP16(16-bit floating point)이나 INT8(8-bit integer)로 변환하여 모델의 크기를 줄일 수 있게 된다.
Llama 등 파라미터 수가 매우 큰 LLM 모델들이 출시되면서, 13B 크기만 하더라도 26GB의 vram이 필요한 등 가정용 GPU로는 구동하기 힘들었으나, Quantization 기법을 통해 필요한 메모리를 약 4배(to INT8)까지 줄일 수 있게 되었다. 그러나, 모델 압축 과정에서 발생하는 오버헤드로 인해 추론 성능이 감소한다. 양자화는 기본적으로 다른 데이터 유형의 범위로 변환할 때 반올림하기 때문에, 정보 손실, 즉 Quantization Error가 발생한다. Quantization Error에 대해서는 아래 두가지 블로그 글의 예시가 이해가 쉬울 것이다.
- LLM.int8() and Emergent Features
- A Gentle Introduction to 8-bit Matrix Multiplication for transformers at scale using Hugging Face Transformers, Accelerate and bitsandbytes
bitsandbytes는 QLoRA를 논문으로 발표한 Tim Dettmers에 의해 배포된 양자화 라이브러리이며, 뒷단에서 LLM.int8()라고 명명된 양자화 절차가 동작한다. 벡터별 양자화, 혼합 정밀도 분해의 2가지 절차이다.
전체 레이어(텐서)에 양자화를 적용하는 고전적인 양자화 방법에서는 특정 임계치를 초과하는 이상값(Outliers)이 다른 모든 값의 정밀도(역양자화 시 손실되지 않는 정도)를 떨어뜨리게 되며, 심각하게 성능을 저하하게 된다. 벡터별 양자화는 행과 열 각각의 absmax(절대 최대값)을 추출하여 양자화를 진행한다. 전체 모델에 양자화를 적용하지 않고, 행과 열로 나누어 폭발 반경을 최소화한다.
벡터별 양자화도 6.7B 이상의 모델에서 나타나는 이상값 특성을 해결할 수는 없었기에, 이를 해결하기 위해 혼합 정밀도 분해 기법이 도입되었다. 입력 행렬에서 이상값을 포함한 행을 추출하고, 가중치 행렬에서는 이상값을 포함한 열을 추출하여 FP16으로 행렬 곱셈을 수행한다. 이상값의 비율이 전체 차원 중 0.1%에 불과하였기 때문에, 성능 저하를 최소화하여 양자화를 수행할 수 있게 되었다.
GPTQ는 LLM을 양자화하는 또다른 방법이며, Post-Training Quantization(PTQ)를 통해 32비트에서 3~4비트로 정수 양자화를 수행한다. 기존 OBQ 기법에서 가중치를 양자화할 때 Greedy 방식으로 순서를 조정하며 실행시간이 모델 크기에 비례하여 증가했던 것을 개선하여, 모든 행에 대해 고정된 순서로 가중치를 양자화하도록 한다.
위 방식을 적용할 때 발생할 수 있는 메모리 대역폭 제한으로 인한 병목을 줄이기 위해 Lazy Batch 기법이 함께 도입되어 한 번에 여러 열에 업데이트하도록 하였다. 또한 반복 작업으로 인한 누적된 수치 오류의 개선을 위해 Cholesky 분해를 적용하였다. GPTQ는 bitsandbytes과 마찬가지로 AutoGPTQ 라이브러리를 통해 Hugging Face 생태계와 직접적으로 통합되었다.
Optimized transformers code for inference using flash-attention
Flash Attention은 Google Research Brain 팀에 의해서 개발된 메커니즘이다. N을 시퀀스 길이라고 할 때, 전통적인 Attention은 O(N²)의 메모리 복잡성을 가진다. 이로 인해 모델의 효율과 속도가 떨어지는 문제가 있었다. Flash Attention은 메모리 복잡성을 선형 복잡도인 O(N)까지 줄여준다. 이를 통해 훈련과 추론 속도를 가속화하였다.
다른 최적화 기법과는 달리, Flash Attention은 하드웨어적으로 성능을 개선한다. GPU 내 메모리 계층에 따라 용량과 속도가 다르다는 것에서 착안한다. HBM(High Bandwidth Memory)는 메모리 용량은 크지만 처리 속도는 느리고, SRAM은 용량이 작지만 속도가 빠르다.
표준 Attention에서는 HBM에서 쿼리(Query), 키(Key), 값(Value)를 로드하고, SRAM에 올려 연산을 수행하고, 다시 HBM에 해당 정보를 저장한다. Load 후 Write하는 과정이 매 어텐션 단계마다 반복되면서 메모리 병목 현상이 발생한다. Flash Attention은 쿼리, 키, 값을 한 번만 로드하도록 함으로써 이러한 문제를 해결하고 효율적인 메모리 사용을 가능하게 하였다. 내부적으로는 Tiling과 Recomputation이라는 알고리즘을 통해 이를 구현한다.
- Tiling : 순방향 및 역방향 패스에서 Softmax를 작은 블록 단위로 분할하여 HBM에서 큰 Attention 행렬를 사용하지 않도록 하며, 메모리 읽기/쓰기 사용량을 최적화한다.
- Recomputation : 역방향 패스에서, O(N²)의 메모리를 저장하지 않기 위해, 저장된 출력과 softmax normalization statistics를 통해 Attention 행렬을 재계산한다.
Conclusion
지금까지 Hugging Face의 TGI가 Tensor Parallelism, 양자화 기법, Flash Attention을 통해 어떻게 추론 성능을 향상시켰는지에 대해 살펴보았다. 본 포스팅에서는 소개하지 않았지만 safetensors, 토큰 스트리밍, in-flight 배칭 등의 기법 또한 TGI 내에서 지원하고 있다.
ChatGPT 등장 이후 2023년 동안 LLM 모델 관점에서의 발전에 주목이 집중되었지만, TGI와 같은 대규모 모델을 활용하기 위한 최적화 기법과 오픈 소스 생태계 또한 같이 발전해왔다는 사실을 간과하지 않아야 한다.
