How to scale AI models across multiple GPUs for inference

TL;DR

Distributed inference scales AI models by spreading work across many GPUs. It uses data, pipeline, tensor, and expert parallelism, plus separate prefill and decode pools to manage KV cache and memory bandwidth, preserving throughput and reducing latency under real user load.

Transcript

The largest AI models today can have over a trillion parameters, which would take something like two terabytes of memory to store. But the largest GPU on the market holds about 288GB, which is nowhere near enough to store these huge frontier models. So how are we running these models behind the chatbots everyone uses reliably for millions of people... Read More

Key Insights

• AI models with trillions of parameters cannot fit on a single GPU due to memory limits.
• Distributed inference uses multiple GPUs to store and run model weights and KV cache.
• Data parallelism replicates the full model across GPUs and routes requests to them.
• Pipeline parallelism splits model layers across GPUs to form an assembly-line processing flow.
• Tensor parallelism divides individual layers across GPUs to compute partial results.
• Mixture of Experts models deploys experts across GPUs to handle diverse inputs efficiently.
• Disaggregation splits prefill and decode workloads onto separate GPU pools to optimize bottlenecks.
• An orchestration layer coordinates routing, load balancing, and fault tolerance across GPUs.

Questions & Answers

Q: A real search query question about the video (e.g. 'How to...', 'What is...'). Question 1

What is distributed inference and why is it necessary for AI models that exceed the memory capacity of a single GPU, requiring multiple GPUs to store weights and KV cache while handling real user traffic without bottlenecks?

Q: Question 2

How does data parallelism work in distributing a large model across GPUs, and why does it not require cross-GPU coordination during inference? The approach mirrors sending requests through different doors to achieve parallel processing.

Q: Question 3

What is pipeline parallelism and how does it enable running large models by splitting layers across GPUs, ensuring continuous flow of requests to keep all GPUs busy rather than waiting?

Q: Question 4

What is tensor parallelism, and what is the challenge of combining partial results when GPUs compute slices of the same layer? This method requires high bandwidth, low latency interconnects inside a server.

Q: Question 5

How do mixture of experts models employ expert parallelism across GPUs, and what tradeoffs come with dispatching tokens to multiple experts and aggregating results for each token?

Q: Question 6

What is prefill and decode in the context of LLM inference, and why is it beneficial to separate these stages onto dedicated GPU pools to optimize compute and memory bandwidth?

Q: Question 7

Why is memory bandwidth a critical bottleneck in decode, and how does KV cache growth influence hardware utilization during inference, requiring careful hardware scheduling?

Q: Question 8

What role does the orchestration layer play in production distributed inference, and how does it balance load and handle failures to prevent dropped requests while serving hundreds or thousands of users?

Summary & Key Takeaways

• In production large AI models require multiple GPUs to serve real user traffic without bottlenecks, using several forms of parallelism. The system also separates prefill and decode tasks to optimize memory use and throughput. An orchestration layer routes requests and balances load across GPUs to prevent failures from impacting performance.

• The video explains four main parallelism strategies (data, pipeline, tensor, expert) and how they coordinate within and across servers. It emphasizes the importance of KV cache management and memory bandwidth in achieving low latency at scale. It also mentions disaggregation as a method to optimize different phases of inference.

• A comprehensive approach combines multiple techniques in a multi-dimensional parallelism framework, supported by an orchestration layer that ensures reliable, scalable production serving and handles hardware failures without dropping requests.