Inference Optimization
Model Parallelism for Inference â Serving Models Too Large for One GPU
When a model exceeds single-GPU memory, model parallelism splits it across multiple GPUs. This guide covers tensor parallelism, pipeline parallelism, and hybrid strategies for inference.
- Tensor Parallelism â Split individual layers across GPUs
- Pipeline Parallelism â Split model stages across GPUs
- Expert Parallelism â Distribute MoE experts across GPUs
The model does not care how many GPUs it runs on â it only cares about serving responses.
Model Parallelism for LLM Inference
Modern LLMs like LLaMA-2 70B (140GB in FP16) and GPT-4 (~360GB) exceed the memory of any single GPU. Model parallelism distributes the model across multiple GPUs to enable inference on models that cannot fit in one device.
Tensor Parallelism for Inference
How It Works
import torch.distributed as dist
from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel
def tensor_parallel_inference(model, world_size):
"""Apply tensor parallelism for inference."""
# Define parallel plan
tp_plan = {
"self_attn.q_proj": ColwiseParallel(),
"self_attn.k_proj": ColwiseParallel(),
"self_attn.v_proj": ColwiseParallel(),
"self_attn.o_proj": RowwiseParallel(),
"mlp.gate_proj": ColwiseParallel(),
"mlp.up_proj": ColwiseParallel(),
"mlp.down_proj": RowwiseParallel(),
}
# Parallelize the model
model = parallelize_module(model, tp_plan)
return model
Communication Pattern
Pipeline Parallelism for Inference
How It Works
Hybrid Parallelism
TP + PP Combination
Expert Parallelism for MoE Models
def expert_parallel_forward(x, expert_indices, expert_params):
"""Forward pass with expert parallelism."""
# Route tokens to appropriate GPUs
local_experts = [i for i in range(num_experts) if expert_owner[i] == local_rank]
# Gather tokens for local experts
local_tokens = all_to_all_dispatch(x, expert_indices, local_experts)
# Compute local expert outputs
local_outputs = []
for expert_id in local_experts:
expert_output = expert_forward(local_tokens[expert_id], expert_params[expert_id])
local_outputs.append(expert_output)
# Scatter outputs back
output = all_to_all_combine(local_outputs, expert_indices)
return output
Serving Frameworks
Tensor Parallelism in Practice
| Framework | TP Support | Max TP | Communication | Notes |
|---|---|---|---|---|
| vLLM | Yes | 8 GPUs | NCCL | Best for general serving |
| TensorRT-LLM | Yes | 8 GPUs | NCCL | Fastest inference |
| SGLang | Yes | 8 GPUs | NCCL | Best for structured generation |
| TGI | Yes | 8 GPUs | NCCL | Production-ready |
Loading Models with TP
# vLLM tensor parallel serving
from vllm import LLM, SamplingParams
# Serve LLaMA-2 70B with TP=2
llm = LLM(
model="meta-llama/Llama-2-70b-chat-hf",
tensor_parallel_size=2,
gpu_memory_utilization=0.9,
max_model_len=4096,
)
# Generate
sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
outputs = llm.generate(["What is machine learning?"], sampling_params)
Communication Optimization
Overlapping Communication and Computation
def overlapped_tp_forward(layer, x, comm_group):
"""Tensor parallel forward with communication overlap."""
# Start async reduce-scatter for previous layer output
handle = dist.reduce_scatter_async(x, group=comm_group, async_op=True)
# Compute current layer while communication proceeds
output = layer(x)
# Wait for communication to complete
handle.wait()
return output
Practice Exercises
-
Parallelism Design: Design a parallelism strategy for serving a 405B parameter model on 16 GPUs. Justify your choice of TP and PP degrees.
-
Latency Analysis: Calculate the inference latency for a 70B model with TP=2 vs TP=4 on 4 GPUs. Assume 100GB/s interconnect bandwidth.
-
Memory Analysis: If a 70B model requires 140GB in FP16, how much memory does each GPU need with TP=2, TP=4, and TP=8?
-
Communication Volume: For a 70B model with TP=4, calculate the total communication volume per token generation step.
Key Takeaways
What to Learn Next
-> Distributed Training for LLMs Parallelism strategies for training.
-> Mixture of Experts MoE architectures and expert routing.
-> Flash Attention and Memory Efficiency IO-aware attention algorithms.
-> KV Cache Optimization Reducing memory usage of the key-value cache.
-> LLM Inference Optimization Broader inference optimization strategies.
-> Building Production LLM Applications End-to-end production systems.