[MLOps, LLMOps] In-flight batching

서빙에서 batch 최적화 기법으로 Adaptive batching을 소개한 적이 있습니다.

https://seokhyun2.tistory.com/91

[MLOps] BentoML - Adaptive Batching

Adaptive batching의 경우, 아래 이미지와 같이 batch 하나에 대해서 longest 기준에 맞춰 padding을 해서 inference를 하게 됩니다.

Adaptive batcing - padding

Transformer의 encoder 기반 모델들은 한 번에 추론을 하기 때문에 저런 방식이 적합하지만, decoder 기반의 생성형 모델들의 경우에는 하나씩 생성하기 때문에 위의 방식이 비효율적이게 됩니다.

 

생성형 모델의 경우, 하나의 토큰씩 생성하다가 end 토큰이 나오게 되면 더 이상 생성할 필요가 없는 구성이기 때문인데요.

생성형 모델에서 batch로 inference를 하게 된다면 아래와 같은 상황이 나오면  3개를 생성하면 input2에 대해서는 더 이상 생성할 필요가 없고, 5개를 생성하면 input3에 대해서는 생성할 필요가 없는 데 계속 의미 없이 inference를 하게 되는 것입니다.

 

이런 상황에서 의미 없는 inference 대신에, 새로운 input에 대한 생성을 할 수 있도록 구현한 기법이 바로 In-flight batching 입니다.

In-flight bathing

위 그림과 같이, input2에 대한 inference가 끝나면 그 뒤에는 새로운 input4에 대한 생성을 시작하고, input3에 대한 inference가 끝나면 그 다음 새로운 input5에 대한 생성을 시작하여 GPU의 효율을 최대한으로 높이는 방식입니다.

 

그래서 In-flight batching은, inference 속도가 빨라지는 것은 아니고 batch inference를 수행함에 있어서 불필요한 inference를 하지 않고 생성이 끝나면 새로운 input에 대한 inference를 수행함으로써 GPU의 효율성을 극대화하는 serving 최적화 기법입니다.

In-flight batching을 적용하면, 생성형 모델 inference 서버의 처리량을 올릴 수 있게 됩니다.

 

In-flight batching은 현재 TensorRT-LLM에서 '--use_inflight_batching'라는 간단한 옵션만으로 사용할 수 있게 제공하고 있습니다.

그 외에도 다양한 최적화 기법을 제공하기 때문에 생성형 모델 서버를 운영하신다면, TensorRT-LLM 도입을 고려해볼 필요가 있는 것 같습니다.