LLM 인퍼런스 훑어보기 (5) - continuous batching

이전 포스트에서는 kernel funsion에 대해 알아보는 내용을 다뤘었습니다. 커널 퓨전의 의미, 구현 방법, 그리고 활용 사례에 대해 자세히 알아보고 싶으신 분들은 이전 포스트를 참고하시면 좋을 것 같습니다.

LLM 인퍼런스 훑어보기 (4) - kernel fusion

 

오늘은 LLM 인퍼런스에서 Batching 기법에 대해서 다루고자 합니다.

 

---

 

 

Batch

ML에서 "batch"는 일반적으로 훈련 데이터를 나누는 단위를 나타냅니다. 전체 데이터셋을 한 번에 모델에 입력하는 것이 아니라 작은 일괄(batch) 단위로 나누어 모델을 훈련합니다. 이렇게 함으로써 메모리를 효율적으로 사용하고 계산 속도를 높일 수 있습니다. Batch size가 크면 더 많은 데이터를 한 번에 처리할 수 있지만 더 많은 메모리가 필요하고 계산 비용이 높아집니다. 작은 batch size는 더 적은 메모리를 사용하지만 계산 속도가 느려질 수 있습니다. 적절한 batch size를 선택하는 것은 모델의 성능에 영향을 미칠 수 있습니다.

inference 단계에서도 "batch"라는 용어가 사용됩니다. Inference에서의 batch는 모델이 동시에 처리하는 입력 데이터의 묶음을 나타냅니다. 여러 입력을 한 번에 모델에 전달하여 병렬처리를 활용하고 성능을 향상시키는 데 사용됩니다. 배치 인퍼런스를 사용하면 모델이 각 입력에 대한 예측을 하나씩이 아니라 여러 입력에 대한 예측을 동시에 수행할 수 있습니다. 이는 대량의 데이터에 대한 예측을 빠르게 생성하는 데 도움이 됩니다.

 

LLM 인퍼런스에서 Batching 기법의 중요성

LLM 인퍼런스에서 Batching 기법이 왜 중요한지 알아보기 위해서, LLM 인퍼런스 과정에 대해서 좀 더 깊게 살펴보겠습니다.

먼저 "캘리포니아의 수도는 무엇입니까?"라는 프롬프트를 입력하게 되는데 이러한 입력 프롬프트를 처리하는 단계를 prefil phase라고 합니다. prefil phase에서는 프롬프트 전체의 attention matrix를 계산하며 이 attention matrix는 이 후 후속 토큰들을 생성하는데에 계속 사용됩니다. 후속 토큰을 생성하는 것을 decode 라고 하는데, decode 파트에서는 새 토큰에 대한 attention만 계산하면 됩니다. 따라서 prefile phase는 이후 각 후속 토큰을 생성하는 시간만큼을 소요하게 됩니다. 그렇지만 prefill phase는 입력값들이 서로 독립적으로 계산될 수 있기 때문에 GPU의 병렬 컴퓨팅을 사용하여 효율적으로 계산할 수 있습니다.

앞선 포스트들에서도 다뤘듯이 LLM 인퍼런스는 memory-IO 바운드입니다. 현재는 컴퓨팅 코어가 1MB의 데이터에 대해 LLM 계산을 수행하는 것보다 GPU의 컴퓨팅 코어에 1MB의 데이터를 로드하는 데 더 많은 시간이 걸립니다. 즉 LLM 인퍼런스의 throughput은 high-bandwidth GPU 메모리에 얼마나 큰 batch 크기를 넣을 수 있느냐에 따라 크게 좌우됩니다. 

LLM 인퍼런스에서 소비되는 GPU 메모리 양은 모델 크기 + 토큰 시퀀스 길이에 따라 결정됩니다. 일반적으로 13B parameter 모델은 시퀀스의 각 토큰의 state를 위해 약 1MB를 소비하는 것으로 추정됩니다. A100 40GB GPU라고 가정하면, 26GB는 모델 파라미터를 저장하고 나머지 14GB는 토큰을 저장할 수 있습니다. 토큰당 1MB를 소모하므로 아무리 많아도 14k개의 토큰만 가질 수 있습니다. 시퀀스 길이를 512로 제한한다면 최대 28 시퀀스들을 처리할 수 있고, 2048로 제한한다면 고작 7개 시퀀스 밖에 처리하지 못한다는 의미입니다.

이러한 메모리 한계점들을 해소하기 위해서 여러 방법들이 나왔고, AutoGPTQ와 같은 quantization 전략들이 대표적입니다. 물론 모델 파라미터의 수정만이 가능한 것은 아니고 FlashAttention처럼 어텐션 연산에서 memory-IO를 덜 쓰도록 하는 방법론들도 있습니다.

오늘 포스트에서 주요하게 다룰 batching 기법은 continuous batching 또는 in-flight batching이라고 하며, 이러한 방법은 모델에 대한 수정 없이도 메모리를 더 효율적으로 사용할 수 있는 방법입니다.

 

Naive batching / static batching

LLM 인퍼런스에서 사용하던 전통적인 방법은 static batching라고 부르며, 이러한 방법은 배치의 크기가 인퍼런스가 완료되기 전까지 고정적으로 유지됩니다. 아래 일러스트는 LLM 인퍼런스에서 static batching의 모습입니다.

https://www.anyscale.com/blog/continuous-batching-llm-inference

static batching을 사용하여 4개의 시퀀스를 완료하는 시나리오라고 가정했을 때, 왼쪽의 그림은 첫 번째 이터레이션을 나타냅니다. 첫 번째 이터레이션에서는 각 시퀀스는 노란 프롬프트 토큰들에서부터 하나의 파란 토큰을 생성합니다. 어느정도 이터레이션이 완료된 후에는, 완료된 시퀀스들은 각각 서로 다른 크기를 가집니다. 각각의 시퀀스는 서로 다른 이터레이션에서 각각의 빨간 EOS 토큰을 생성하기 때문입니다. 3번 시퀀스가 2번의 이터레이션에 종료되었다고 하더라도, static batching 전체 시퀀스의 생성이 완료될 때까지 GPU가 충분히 활용되지 않습니다.

이러한 특성은 기존 딥러닝 모델과 달리 LLM 인퍼런스의 iterative한 특성으로 인해 발생하는 문제입니다. 직관적으로, 이는 요청이 배치 처리에서 더 일찍 "완료"될 수 있지만, 리소스를 해제하고 완료 상태가 서로 다른 시퀀스들에 새 요청을 추가하는 것이 까다롭기 때문입니다. 그렇기 때문에 LLM 인퍼런스에서는 배치 내에서의 생성된 시퀀스의 길이에 따라 GPU util이 다를 수 있습니다. 모든 입력과 생성 시퀀스들이 같은 길이를 갖는 경우에는 static batching은 최대의 GPU util을 활용할 수 있습니다.

그러나 일반적으로 chatbot 등과 같은 LLM의 사용처들에서는 입력이나 출력 시퀀스의 길이를 고정할 수 없습니다. 이러한 점을 해소하기 위해 나온 방법론이 바로 continuous batching입니다.

 

Continuous batching

업계에서는 위와 같은 비효율성을 인식하고, 더 나은 접근 방식을 고려했습니다.

Orca라는 방법론이 가장 먼저 이러한 문제를 해결하기 위해 제시되었습니다. Orca는 배치 안의 모든 시퀀스들이 생성 완료되는 것을 기다리는 대신에, iteration-level 스케쥴링을 구현했습니다. iteration-level 스케쥴링에서는 배치 사이즈가 각 이터레이션마다 결정됩니다. 결과적으로 배치 안의 한 시퀀스가 생성 완료되면, 새로운 시퀀스가 완료된 시퀀스의 자리에 주입될 수 있습니다. 이러한 방식은 static batching 방법보다 더 높은 GPU util을 가져갈 수 있습니다.

https://www.anyscale.com/blog/continuous-batching-llm-inference

continuous batching을 사용하여 4개의 시퀀스를 완료하는 시나리오라고 가정했을 때, 왼쪽의 그림은 첫 번째 이터레이션을 나타냅니다. 첫 번째 이터레이션에서는 각 시퀀스는 노란 프롬프트 토큰들에서부터 하나의 파란 토큰을 생성합니다. 어느정도 이터레이션이 완료된 후에는 오른쪽 그림처럼 batch가 형성됩니다. 한 시퀀스가 EOS 토큰을 출력하면, 새로운 시퀀스가 그 자리에 주입됩니다.

이러한 continuous batching은 현재 연속 일괄 처리를 구현하는 두 가지 기본 접근 방식이 있습니다. TGI 및 vLLM에서는 생성 단계는 생성을 계속하기 전에 프롬프트 처리(TGI에서는 infill이라고 함)를 수행하기 위해 선점됩니다. Orca에서는 이러한 단계가 구분되지 않습니다. 오르카는 총 시퀀스 수가 고정된 한계에 도달하지 않는 한 실행 중인 배치에 프롬프트를 추가합니다. 다양한 수준의 이러한 접근 방식은 모두 긴 프롬프트를 처리하기 위해 생성을 지연시켜야 합니다.

 

Continuous batching 지원하는 framework

continuous batching은 LLM 인퍼런스의 효율성 극대화에 효과가 좋은 만큼, 이것을 지원하는 framework들이 많아졌습니다. 대표적으로는 vLLM, TGI가 있으며, 최근에는 TensorRT-LLM, Deepspeed-fastgen 등에서도 지원하기 시작했습니다. 본 포스트에서는 각 프레임워크들에 대해서 깊게 다루지 않겠습니다. 추후에 기회가 된다면 각 프레임워크들을 비교하는 글을 게시할 수도 있을 것 같습니다.

 

vLLM

https://github.com/vllm-project/vllm

GitHub - vllm-project/vllm: A high-throughput and memory-efficient inference and serving engine for LLMs

 

TGI

https://github.com/huggingface/text-generation-inference

GitHub - huggingface/text-generation-inference: Large Language Model Text Generation Inference

 

TensorRT-LLM

https://github.com/NVIDIA/TensorRT-LLM

GitHub - NVIDIA/TensorRT-LLM: TensorRT-LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs)

 

Deepspeed-fastgen

https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-fastgen

 

---

 

이렇게 이번 포스트에서는 LLM 인퍼런스에서 batching 기법의 중요성을 알아보았고, GPU util을 극대화하기 위한 batching 기법인 continuous batching에 대해서 알아보았습니다. vLLM과 TGI에서 먼저 도입되어 각광받았고 현재는 대부분의 framework들이 구현하여 지원하고 있어서 한번쯤은 고려해야하는 기능이 된 만큼, 이러한 내용을 숙지하면 LLM 서비스 제작에 큰 도움이 될 것 같습니다.

다음 포스트에서는 다른 메모리 효율화 방법들에 대해서 다뤄볼 예정입니다.

감사합니다.