Batch Size and LLM Inference Efficiency

#machinelearning #llm #mlops #ai

Having an optimal Batch size can decrease your models cost per token at the time of inference.

This will be an explanation on how Batch size affects the cost at the inference, We will be going deep and building the framework from the ground up. This will be an informal explanation of the concept so don't expect any formal language.

The word plot and graph are used interchangebly in this Article so dont get confused

Dependencies

So first thing we need to understand is what does LLM inference actually depends upon.

In the equation above we can see that the actual time taken to do inference depends on two major factors

1 -> Memory

KV cache & Model weights fetching time

2 -> Compute

The time it takes to perform mathematical operations

The equation says, The time it takes to do inference is greater than or equal to the max value between time it takes for compute and time it takes to fetch the memory.

Time taken to fetch memory (t_memory)

In the equation above, time taken to fetch the memory is the sum of size of the model in bytes and the kv cache in bytes divided by the memory bandwidth

here the kv cache is

B × len_ctx × (bytes/token)

Time taken for compute (t_compute)

The equation t_compute is batch size multiplied by number of active parameter divided by flops

We don't take time it takes to do the attention computation into account because its trivial when you compare it with the weight matrix multiplication cost.

Plotting the graphs

What we have here is the latency plot, we are essentially plotting all the values that we worked with earlier, on this graph.

plot 1.1

t_compute grows linearly as the batch size increases
t_kv-fetch also grows linear.
t_weight-fetch dosen't grow at all it remains constant

plot 1.2

plotting the conjunction of both t_weight-fetch & t_kv-fetch as t_memory
notice that the y intercept of t_memory starts exactly from the t_weight-fetch intercept

plot 1.3

Here we are taking the maximum of t_compute & t_memory
Now this graph tells us, How latency grows with the increase in batch size.
The point worth noting here is that there is a lower bound on latency, because the time it takes to load the model parameters is a constant

we have been plotting the latency (t) as the function of batch size, but to get cost per token we need to plot (t/B) cost as the function of batch size.

plot 2.1

The Graph 2.1 plots (t/B) as a function of (B), and the variables change quite a bit

t_compute was linear but after dividing it by (B) we get a constant
t_KV-fetch was also linear but after dividing it by (B) we get a constant
t_weight-fetch was constant but after dividing it by (B) we get a hyperbolic curve

plot 2.2

Plotting the conjuntion of both t_weight-fetch & t_kv-fetch as t_memory
What happens is that the conjuntion shifts the hyperbola up

plot 2.3

Taking the maximum of the graph 2.2 we get graph 2.3
Now, you can see the cost per token plot

Observations

In the graph 2.3 you can see that the cost decreases as we increase the batch size until it hits the lower bound
If you want you can do a cost and inference time tradeoff where as you decrease the batch size the cost goes up but inference time decreases
There is a lower bound on cost too which is t_compute.
Here is a caveat, As you increase the batch size, at some point t_kv-fetch will surpass t_compute and at that point compute is no longer the lower bound its memory bandwidth and the MFU (model flop utilization) decreases significantly, So the cost/token increases if the batch size is too big

Thanks for the read :)