PaperlayerView exact revision on arXiv

Source-linked AI summary

FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU

Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, Ce Zhang

arXiv:2303.06865v2cs.LGcs.AIcs.PF

TL;DR

Large language model inference requires substantial computation and memory, motivating high-throughput generation on limited hardware for latency-insensitive batched tasks. FlexGen aggregates GPU, CPU, and disk resources, searches offloading policies with linear programming, and compresses weights and KV cache to 4 bits. It achieves higher throughput on a single 16GB GPU, including 100× higher maximum throughput with effective batch size 144, while maintaining negligible accuracy loss under the evaluated compression scheme.

  • Problem

    LLM inference has extremely high computational and memory requirements, creating a need for high-throughput generation with limited resources for latency-insensitive batched tasks.

  • Method

    FlexGen aggregates GPU, CPU, and disk resources, uses linear-programming policy search for tensor placement and access, and applies 4-bit compression to weights and KV cache.

  • Results

    100× higher maximum throughput is achieved with effective batch size 144 under a compressed configuration, while 4-bit compression produces negligible accuracy loss compared to FP16.

  • Takeaways & Limitations

    FlexGen supports high-throughput LLM generation on resource-constrained hardware, including a single commodity GPU and multi-scenario benchmarking of a 30B model.

  • Takeaways & Limitations

    The policy search can produce strategies that run out of memory because of relaxed placement variables and difficulty modeling peak memory usage, so manual adjustment may be needed.

Abstract

from arXiv · show

The high computational and memory requirements of large language model (LLM) inference make it feasible only with multiple high-end accelerators. Motivated by the emerging demand for latency-insensitive tasks with batched processing, this paper initiates the study of high-throughput LLM inference using limited resources, such as a single commodity GPU. We present FlexGen, a high-throughput generation engine for running LLMs with limited GPU memory. FlexGen can be flexibly configured under various hardware resource constraints by aggregating memory and computation from the GPU, CPU, and disk. By solving a linear programming problem, it searches for efficient patterns to store and access tensors. FlexGen further compresses the weights and the attention cache to 4 bits with negligible accuracy loss. These techniques enable FlexGen to have a larger space of batch size choices and thus significantly increase maximum throughput. As a result, when running OPT-175B on a single 16GB GPU, FlexGen achieves significantly higher throughput compared to state-of-the-art offloading systems, reaching a generation throughput of 1 token/s for the first time with an effective batch size of 144. On the HELM benchmark, FlexGen can benchmark a 30B model with a 16GB GPU on 7 representative sub-scenarios in 21 hours. The code is available at https://github.com/FMInference/FlexGen

1. Introduction

FlexGen targets latency-insensitive, batched LLM inference on limited hardware by combining offloading, optimized tensor placement, and 4-bit compression. On a single 16GB GPU, it substantially improves throughput over existing offloading systems.

  • LLM inference is constrained by large computational and memory demands, with GPT-175B requiring 325GB merely to load its weights.
  • Latency-insensitive batch workloads can trade latency for higher throughput and lower resource requirements.
  • FlexGen aggregates GPU, CPU, and disk resources while scheduling I/O, compression, and distributed pipeline parallelism.
  • A linear-programming search unifies placement of weights, activations, and KV cache, producing an offloading strategy within 2× of optimal I/O complexity and enabling larger batch sizes.
  • 4-bit compression reduces weights and KV-cache storage without retraining or calibration, with negligible accuracy loss.
  • 100× higher maximum throughput is achieved with effective batch size 144 and 4000-second latency when 4-bit compression is allowed.
  • FlexGen outperforms a decentralized Petals cluster in per-GPU throughput and can achieve lower latency in some cases.

2. Related Work

Prior LLM inference systems largely target latency-oriented workloads on high-end accelerators, while related compression, collaborative inference, and offloading approaches address resource use from different directions.

  • Most specialized LLM inference systems focus on latency-oriented scenarios with high-end accelerators.
  • Offloading systems support commodity-hardware inference but typically inherit training strategies that overlook generative inference’s computational structure.
  • Sparsification and quantization reduce computation or memory footprint, with prior work compressing weights to 3 bits.
  • Memory optimization and offloading have also been studied for training and linear algebra.

3. Background: LLM Inference

LLM generation uses prefill and autoregressive decoding, with model weights and KV cache dominating memory needs. For OPT-175B, the KV cache can exceed model-weight storage at large batch sizes.

  • Generative inference has prefill and decoding stages: prefill creates the KV cache, while decoding updates it to generate tokens step-by-step.
  • The inference workload is parameterized by batch size, input and output sequence lengths, hidden dimensions, and transformer-layer count.
  • OPT-175B requires 325GB for FP16 model weights, while a batch size of 512 requires 1.2TB for peak KV-cache storage.
  • At those settings, the KV cache is 3.8× the model weights and becomes a bottleneck for large-batch, high-throughput inference.
  • Generation throughput is defined as bn/t, where b is effective batch size, n is generated-token count per input, and t is total processing latency.
  • Figure 2 presents the computational graph of LLM inference.

4. Offloading Strategy

FlexGen models offloaded generative inference across GPU, CPU, and disk as a constrained graph traversal, then searches schedules and tensor placements to minimize execution time. Its zig-zag block schedule is within 2× of optimal I/O complexity, while linear programming and profiling select hardware-specific policies.

  • 4.1. Problem Formulation: FlexGen formulates offloaded generative inference as a graph traversal over GPU-batch computations subject to dependency, tensor-loading, and memory-capacity constraints.A valid path computes every square while respecting tensor availability and device capacities; execution time includes computation and inter-device I/O.
  • 4.2. Search Space: Row-by-row traversal repeatedly reloads weights, whereas column-by-column traversal reuses shared layer weights to reduce weight I/O.The column-wise approach must stop before activations and KV cache fill CPU and disk memory, motivating a bounded block schedule.
  • 4.2. Search Space: The zig-zag block schedule has I/O complexity within 2× of the optimal solution.FlexGen implements the simpler block schedule rather than the more difficult optimal schedule.
  • 4.2. Search Space: FlexGen overlaps loading the next layer and batch with storing the previous batch and computing the current batch.The block schedule exposes GPU batch size and the number of GPU batches as search parameters; their product is the effective batch size.
  • 4.3. Cost Model and Policy Search: Policy search enumerates block-size and GPU-batch-size choices, then solves a nine-variable linear program for weight, activation, and KV-cache placement.The cost model uses profiled hardware parameters and can be extended with latency constraints.
  • 4.3. Cost Model and Policy Search: The cost model can produce out-of-memory policies because percentage relaxation and peak-memory modeling do not fully capture fragmentation, so manual tuning may be required.The authors report that manual tuning can obtain a better policy than the optimizer’s output.

5. Approximate Methods

FlexGen adds approximation methods to reduce memory and I/O costs during high-throughput inference. It quantizes weights and KV cache to 4 bits and uses top-10% sparse attention while maintaining similar accuracy or model quality on OPT-175B.

  • Approximation Methods: FlexGen introduces group-wise quantization and sparse attention as approximations for boosting inference throughput with negligible accuracy loss.The framework is designed to plug in multiple approximation methods.
  • Group-wise Quantization: 4-bit quantization compresses both weights and KV cache without retraining or calibration on OPT-175B.Weights use groups along the output-channel dimension, while KV cache uses groups along the hidden dimension, with group size 64.
  • Sparse Attention: Top-10% sparse attention loads only the highest-scoring 10% of attention value-cache entries while maintaining model quality.For each query, FlexGen selects top-K tokens from the K cache and loads the corresponding subset of the V cache.

6. Evaluation

FlexGen is evaluated on resource-constrained LLM generation across hardware configurations, model sizes, latency constraints, compression settings, and benchmark workloads. It outperforms the evaluated offloading baselines and scales throughput through policy search, compression, and parallelism.

  • Maximum throughput benchmark: FlexGen outperforms all baselines in maximum generation throughput across the evaluated cases.OPT-6.7B can use a single GPU with Accelerate and FlexGen, while OPT-30B requires CPU offloading for all evaluated systems.
  • Scaling performance: FlexGen achieves super-linear scaling on decoding throughput with pipeline parallelism across four GPUs.Pipeline parallelism reduces per-machine memory pressure, enabling larger batch sizes or less disk offloading than independent data-parallel execution.
  • Latency-throughput trade-off: 40× higher throughput is achieved than DeepSpeed and Accelerate at the same latency requirement of 5000 seconds without compression.With higher latency and compression, throughput reaches a 100× improvement using effective batch size 144.
  • HELM and data wrangling: FlexGen completes seven representative HELM sub-scenarios for OPT-IML-30B in 21 hours with all system overhead included.The evaluation uses the hardware setup described in Table 1.
  • Accuracy and compression: 4-bit quantization of weights and KV cache shows negligible accuracy loss compared to FP16 on Lambada and WikiText.The evaluated 4-bit method uses group-wise quantization for both weights and KV cache.

7. Conclusion

FlexGen is a high-throughput LLM generation engine designed for latency-insensitive batch-processing tasks in resource-constrained scenarios.

  • Conclusion: FlexGen targets high-throughput generation for latency-insensitive batch-processing tasks under resource constraints.The paper presents it as a generation engine for LLM inference.

A.1. Notations

The appendix introduces the notation table used for the paper's formal descriptions.

  • A.1. Notations: Table 6 provides the notations used in the appendix.The subsection explicitly states that it uses the notations in Table 6.

A.2. Compute Schedule Optimality

The appendix analyzes computation schedules for models that do not fit on a single GPU, focusing on tensor movement and reuse across the computational graph.

  • A.2. Compute Schedule Optimality: The analysis compares the zig-zag block schedule with an I/O-optimal diagonal block schedule.It assumes no CPU computation and considers GPU loading of tensors plus cache and activation offloading.
  • A.2. Compute Schedule Optimality: Weights should be reused across prompt batches because generating one prompt requires swapping weights once per token.This reuse amortizes weight I/O time over the batch.

A.2.1. ZIG-ZAG BLOCK SCHEDULE AND DIAGONAL

The zig-zag schedule organizes computation in blocks under memory constraints, while the diagonal schedule can enlarge block size and improve throughput-related efficiency and latency. The diagonal schedule also introduces implementation challenges involving dynamic KV-cache storage.

  • Schedule construction: The zig-zag schedule computes columns in blocks, with block size bls defined as the number of samples touched by the diagonal.The diagonal schedule illustration processes sub-diagonals within each block after a one-time warm-up phase.
  • Memory and I/O: Under FP16, peak memory includes one layer of weights, activations, and KV caches, yielding a CPU-only block-size constraint from available memory.The cited estimate is peak mem = w + 2h1 · bls + 4h1 · bls · l · (s + n).
  • Diagonal-schedule benefits: The diagonal schedule can enlarge block size by around 2 in some cases while preserving the same I/O efficiency as the zig-zag schedule.The ratio is close to 2 when n ≫ s and close to 1 when s ≫ n.
  • Memory and I/O: A larger diagonal block size leaves activations and KV-cache I/O per token unchanged while proportionally reducing weights I/O per token.Weights I/O can occupy a large portion of total I/O, making block-size expansion practically relevant.
  • Diagonal-schedule benefits: When compute resources avoid offloading, the diagonal schedule can reduce peak memory, enlarge batch size, and increase GPU utilization.At equal throughput, it also reduces average completion latency by half compared with the zig-zag schedule.
  • Implementation difficulty: Implementing the diagonal schedule requires efficient attention computation on non-contiguous memory because continuous KV-cache preallocation removes its memory-saving benefit.Dynamic KV-cache buffer updates are identified as the major implementation difficulty.

A.2.2. PROOF OF THEOREM 4.1

The proof models offloaded inference as state transitions over computation squares and derives memory-based lower bounds on weight loading. It shows that the diagonal schedule attains the asymptotic optimum, while zig-zag is within a factor of two.

  • State model: The proof represents a working state by the ordered column indices of the last computed squares for the GPU batches.A move is defined as computation transitioning between working states.
  • I/O accounting: Any move between computation squares requires loading and offloading the corresponding KV cache, while activation I/O is smaller and can be ignored with bounded multiplicative error.The activations are around 1/(2s + n) of KV-cache I/O in size.
  • State model: A closed sequence of moves returning to the same state computes the same number of squares in every column.This equal-column property supports reducing the optimization to repeated state-to-state compute orders.
  • Lower bound: Memory consumption imposes a lower bound on weight-loading times, because each load can compute only squares from different rows while their caches and activations remain resident.The bound is expressed using total square memory M, available capacity M′, and per-color weight-load time tw.
  • Optimality: The diagonal schedule reaches the lower bound by equalizing peak memory across weight loads, making its I/O complexity asymptotically optimal.Its repeated move sequence uses ⌈M/M′⌉ weight-loading times.
  • Optimality: The zig-zag schedule is within 2× of the optimal solution because its peak memory varies across generated tokens, leaving average memory utilization below capacity.The theorem states the formal approximation guarantee for zig-zag scheduling.

A.3. Cost Model

FlexGen formulates throughput-oriented offloading as a linear programming problem over tensor placement and access policies. The cost model accounts for latency, working memory, and peak memory constraints across GPU, CPU, and NVMe resources.

  • Cost model: The cost model simplifies hardware constants while using piece-wise functions and regularization terms to preserve linearity under changing system loads.Constants include bandwidth and TFLOPS; the resulting optimization remains linear with respect to policy variables.
  • Objective: The optimization maximizes throughput, equivalently minimizing its reciprocal in seconds per token.Free variables are explicitly identified in the formulation.
  • Discretization: Discrete tensor-size choices can approximate the target memory allocation closely enough that the total objective changes by less than 1%.Using finer-grained tensor-level partitioning makes the allocation squares extremely fine-grained.
  • Memory constraints: GPU, CPU, and NVMe peak-memory constraints bound the placement of weights, activations, and attention cache during and after prefill.The formulation separately models working memory and peak-memory usage at each storage level.
Loading 2303.06865v2…