Offline Batch Inference for large models

Context:

• LLMs are being used to improve multiple ML based products— ads & search, content recommendations, detecting violating components etc. A common ML lifecycle step in all these is applying LLMs to the entire corpus of content (ads, search queries, document text) and evaluating the predictions. Applying LLMs is resource intensive.
• Let’s consider Llama–7B model, the latency is around ~25 ms/token on 1 A100 GPU [1]. If you have a corpus of 100 million queries of average length of 8 tokens, we are looking at ~30 days of compute on a 8 A100 GPU box.
  Now consider experimenting with multiple LLMs and multiple products with customizations (search queries, ads content, post titles etc) — voila, we have a problem at hand.

Designing an efficient and scalable system for batch inference with large language models on large datasets is a challenging task. Here are some key considerations to keep in mind when designing such a system:

Model batch inference architecture [6]

1. Compute split over CPU / GPU: Accelerators (GPUs/TPUs) often provide an order of magnitude improvement in latency and throughput specially at large batch sizes. In order to optimize for high utilization of GPUs, one must carefully split the batch processing pipeline compute in CPU/GPU. For example, standard feature processing, light-weight sub-model inference can be done over CPU cluster and large model inference can be done over GPUs.
2. Distributed batch inference over multi-GPU and multi-node: To further speedup batch inference, model can be loaded on multiple GPUs on single or multiple nodes to do distributed batch inference. This is similar to distributed data parallel training paradigm and can provide linear scale up.
3. Choice of GPUs: When choosing GPUs for batch inference, it is important to consider the trade-off between cost and performance. While high-end GPUs like the NVIDIA A100 offer excellent performance, they may not be necessary for all use cases. Cheaper GPUs such as the A10G may be sufficient for inference tasks that require less memory. [5]

Choice of GPUs for inference workloads[7]

4. GPU utilization: Higher GPU utilization results in efficient usage of GPU cluster. Some quick parameters to optimize — batch size, collating predictions across several batches in memory before flushing out to store, dynamically constructing batches from multiple workloads if possible.

5. Model optimizations: There are several active workstreams in the ML community to optimize inference models. Freezing weights, quantization, enabling graph mode, ops fusion, using more performant attention mechanisms (BetterTransformer[2]) can have significant improvements (2–10x) in latency / throughput or both depending on model architecture. Tensorrt[3] and Onnx[4] can be integrated in your batch processing pipeline to provide speedups.

CPU vs GPU compute split

In summary, designing an efficient and scalable system for batch inference with large language models on large datasets requires careful consideration of compute resources, distribution strategies, and hardware choices. It is well worth the effort!

References:

1. The Path to Achieve Ultra-Low Inference Latency With LLaMA 65B on PyTorch/XLA | PyTorch
2. https://pytorch.org/blog/a-better-transformer-for-fast-transformer-encoder-inference/
3. https://docs.nvidia.com/deeplearning/tensorrt/
4. https://onnx.ai/
5. https://technical.city/en/video/Tesla-A100-vs-A10G
6. https://www.anyscale.com/blog/offline-batch-inference-comparing-ray-apache-spark-and-sagemaker
7. https://towardsdatascience.com/a-complete-guide-to-ai-accelerators-for-deep-learning-inference-gpus-aws-inferentia-and-amazon-7a5d6804ef1c