Run Local LLM Inference10x Faster (244 TOK/s): PyTorch II

Author(s): Dr. Mandar Karhade, MD. PhD.

Originally published on Towards AI.

There are many ways to run, here is a quick overview

TLDR:

The Pytorch team came up with a way of optimizing purely by building ground-up optimizations below

Torch.compile: A compiler for PyTorch modelsGPU quantization: Accelerate models with reduced precision operationsSpeculative Decoding: Accelerate LLMs using a small “draft” model to predict a large “target” model’s outputTensor Parallelism: Accelerate models by running them across multiple devices.

Pytorch achieved it all inside 1000 lines of code.

Here is the performance comparison

Now, let's go through these one by one!

This is the performance of LLaMA-7b.

Using LLaMA-7B, Pytorch analyzed the performance issues that are CPU-bound. That means the overhead is the first target to make it more efficient..

In comes the compiler. By compiling larger regions into a pre-compiled state, the number of CPU calls per operation decreases. This means that this wrapper can now perform without having gaps, as you see below.

To use it, simply use the following code

torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True)

What is happening under the hood?

Firstly, when more tokens are being generated, the kv-cache grows, needing reallocating and copying (expensive computation) each time the cache grows. Pytorch implemented stating large cache to allow maximum size.

Secondly, two strategies are needed to separately compile during the prefill phase. The entire prompt is processed dynamically and tokens are decoded as the code shown… Read the full blog for free on Medium.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI