PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficiently Scaling Transformer Inference

Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, Jeff Dean

arXiv:2211.05102v1cs.LGcs.CL

TL;DR

Efficient generative inference for large Transformer models remains difficult because token generation is sequential and long contexts increase memory and inference time. The paper develops analytical partitioning strategies and low-level optimizations for TPU v4 slices, achieving low latency, high MFU, and practical long-context inference.

  • Problem

    Efficient deployment is challenging because generative inference proceeds one token at a time, while long attention contexts substantially increase inference time.

  • Method

    The paper uses an analytical framework to select partitioning strategies based on model size, sequence length, hardware, and latency-throughput requirements.

  • Results

    The approach achieves 29ms per token during generation and 76% MFU during large-batch input processing on PaLM 540B, while supporting 2048-token contexts.

  • Takeaways & Limitations

    Scaling inference beyond a single server to 64 or more chips and using appropriate partitioning makes tight-latency and long-context inference practical for 500B+ parameter models.

  • Takeaways & Limitations

    Dense Transformer inference remains fundamentally limited by FLOP count and communication volume, motivating sparsity, adaptive computation, and communication-compression techniques.

Abstract

from arXiv · show

We study the problem of efficient generative inference for Transformer models, in one of its most challenging settings: large deep models, with tight latency targets and long sequence lengths. Better understanding of the engineering tradeoffs for inference for large Transformer-based models is important as use cases of these models are growing rapidly throughout application areas. We develop a simple analytical model for inference efficiency to select the best multi-dimensional partitioning techniques optimized for TPU v4 slices based on the application requirements. We combine these with a suite of low-level optimizations to achieve a new Pareto frontier on the latency and model FLOPS utilization (MFU) tradeoffs on 500B+ parameter models that outperforms the FasterTransformer suite of benchmarks. We further show that with appropriate partitioning, the lower memory requirements of multiquery attention (i.e. multiple query heads share single key/value head) enables scaling up to 32x larger context lengths. Finally, we achieve a low-batch-size latency of 29ms per token during generation (using int8 weight quantization) and a 76% MFU during large-batch-size processing of input tokens, while supporting a long 2048-token context length on the PaLM 540B parameter model.

1 INTRODUCTION

This paper addresses the difficulty of deploying very large Transformer models for generative inference, where sequential token generation, memory demands, and application-specific latency or throughput requirements create engineering tradeoffs. It develops analytical partitioning principles and low-level optimizations, achieving strong latency and MFU results on PaLM models.

  • Generative inference is difficult because decoding depends sequentially on previously generated tokens, unlike highly parallel Transformer training.
  • Interactive workloads emphasize tight latency, whereas offline scoring or distillation emphasizes high throughput and low cost per token.
  • Large models require substantial parameter and KV-cache memory, while attention inference cost grows quadratically with input sequence length.
  • The paper analytically selects multi-axis tensor partitioning, batch size, and chip configuration for application requirements instead of relying on exhaustive black-box search.Fine-grained collective-operation control and low-level scheduling optimizations are also used.
  • The primary goal is an engineering framework explaining how optimal partitioning changes with model size, sequence length, chip count, and latency-throughput targets.
  • 29ms per token generation latency and 76% MFU were achieved on PaLM 540B with a 2048-token context.The generation result used int8 weight quantization; the MFU result applies to large-batch input processing.

2 INFERENCE COST TRADEOFFS

Inference cost reflects distinct prefill and decode behavior, memory and compute demands, communication overhead, and application-specific latency-throughput tradeoffs. Long contexts particularly stress KV-cache memory, while larger batches generally improve throughput and cost efficiency.

  • Inference latency comprises prefill, which processes initial input tokens, and decode, which autoregressively generates output tokens.Because prefill parallelizes across input tokens while decode loops sequentially over output steps, they require separate performance analyses.
  • The model must be partitioned across chips when it exceeds one chip’s memory, introducing chip-to-chip communication costs.
  • Weights and KV-cache tensors dominate focused memory costs because they transfer from HBM to compute cores once per forward pass.
  • A decoder-only model with N parameters requires 2N matmul FLOPs per forward-pass token, while attention FLOPs are typically smaller for large models.
  • Communication becomes an increasingly important bottleneck as chip count grows because its time decreases less quickly than weight-loading and non-attention compute time.
  • Lower latency generally requires more chips or smaller batches, but smaller batches reduce MFU and increase cost per token.
  • For 500B+ models with multihead attention, a batch size of 512 and context length of 2048 produce a 3TB KV cache, three times the parameter size.The cache must be loaded from off-chip memory for every generated token.
  • When latency is unimportant, larger batches typically improve MFU and minimize total cost per token for offline inference.

3 PARTITIONING FOR INFERENCE

The section develops partitioning strategies for large Transformer inference, matching layouts to chip topology, batch size, sequence length, and latency–throughput tradeoffs. It covers weight-stationary and weight-gathered feedforward layouts, plus batch-sharded multiquery attention to reduce KV-cache costs.

  • Large models are partitioned across many TPU v4 chips because weights and activation tensors must fit memory and meet latency requirements, at the cost of chip-to-chip communication.
  • 3.2.1 Feedforward layer, 1D weight-stationary layout: 1D weight-stationary sharding partitions each feedforward weight matrix along one axis and aggregates intermediate results across chips.A consecutive-matmul partitioning trick can avoid communication between the two matmuls in a Transformer MLP block.
  • 3.2.1 Feedforward layer, 1D weight-stationary layout: Communication becomes a bottleneck as chip count grows because memory and compute latency decrease more quickly than activation-aggregation latency.
  • 3.2.2 Feedforward layer, 2D weight-stationary layout: 2D weight-stationary sharding partitions each weight matrix across both dimensions, reducing communication scaling to O(1/√nchips) while preserving compute cost.For dff = 4dmodel, it becomes more communication-efficient than 1D weight-stationary when nchips > 16.
  • 3.2.3 Feedforward layer, weight-gathered layout: Weight-gathered layouts keep activations increasingly stationary and transfer weights across chips as batch size and sequence length grow.The communication-optimal configuration switches from 2D weight-stationary at low tokens per batch to weight-gathered layouts at larger batches.
  • 3.3 Partitioning the attention layer: Multiquery attention partitions Q, K, and V over the batch dimension, reducing per-chip KV-cache loading cost by a factor of nchips.The strategy adds an all-to-all resharding cost but enables larger batch sizes and sequence lengths; the reported savings are an order of magnitude versus multihead attention.

4 CASE STUDY FOR PALM MODELS

The PaLM case study evaluates partitioning and inference optimizations across model sizes, batch sizes, context lengths, and chip configurations. It shows that partitioning choices depend on workload phase and that optimized multiquery attention substantially extends feasible context lengths.

  • Methodology: The study evaluates PaLM models using up to 256 TPU v4 chips, including a 540B model whose attention heads were padded from 48 to 64 for more effective partitioning.The padding added 18B parameters and incurred a 3% MFU cost, which was more than recovered through improved partitioning.
  • Feedforward partitioning: The 2D weight-stationary layout performs better for decoding as chip count increases because it scales asymptotically better, although both layouts eventually become communication-limited.This evaluation used batch size 512 to balance latency and MFU.
  • Feedforward partitioning: 76% MFU is achieved by weight-gathered layouts at high prefill batch sizes, while 2D weight-stationary layouts are preferable at low batch sizes.The optimal layout switches as batch size in tokens grows, and large batches would exhaust memory without multiquery attention.
  • Attention partitioning: 32–64 times longer context lengths fit with the optimized multiquery layout than with multihead attention or the baseline multiquery variant.The optimized layout addresses KV-cache memory pressure at large batch sizes and context lengths.
  • Attention partitioning: Multiquery attention scales to 8192–32,768 tokens, with attention consuming 8–31% of total runtime during generation.Its speed advantage grows with context length because loading the KV cache becomes a larger share of inference time.
  • End-to-end results: 28.5ms/token is achieved with int8 weights at batch size 64 on PaLM 540B, compared with 36.9ms/token using bfloat16 weights.The low-latency advantage of int8 quantization is strongest when weight loading dominates cost.

5 FASTERTRANSFORMER BENCHMARKS

The paper compares its implementation with FasterTransformer across varied inference configurations using MFU to normalize throughput across different hardware. Its PaLM implementation achieves the best absolute latency, while its approaches generally provide stronger MFU and scalability.

  • Benchmark setup: MFU is used to normalize throughput for differences in chip count and chip FLOPS between the TPU v4 and FasterTransformer benchmark setups.FasterTransformer uses 16–32 NVIDIA A100s, whereas this work uses 64 Google TPU v4 chips for the comparison.
  • Results: The PaLM 540B implementation achieves the best absolute latency, while the Megatron implementation achieves the best MFU for all but one latency target.The benchmark compares three FasterTransformer configurations with similarly sized PaLM and Megatron models.
  • Results: 10% MFU is the maximum reported advantage of the PaLM implementation over the Megatron implementation, primarily attributed to parallel attention and feedforward layers.The advantage of multiquery attention is not noticeable because the benchmark context length is too short.
  • Scalability: 44% MFU is achieved with 64-way tensor parallelism, compared with FasterTransformer’s maximum of 33% at 32-way and 46% at 16-way tensor parallelism.The comparison suggests better scaling of the paper’s 2D weight-stationary approach at higher tensor-parallel degrees.

6 RELATED WORK

Related work spans partitioning systems for efficient large-model training and inference, architectural approaches to improve Transformer efficiency, and model quantization. The paper combines these directions with its own inference partitioning techniques.

  • Parallelism approaches: Prior systems such as NeMo Megatron, GSPMD, and Alpa propose partitioning approaches for efficient large-model training, while FasterTransformer benchmarks multi-GPU, multi-node inference.FasterTransformer combines tensor parallelism and pipeline parallelism for inference.
  • ML inference efficiency: Other work improves Transformer inference through architectural changes such as efficient attention layers, distillation, and related model-compression methods.The paper positions its approach alongside these architecture- and compression-focused directions.
  • ML inference efficiency: The paper reuses prior model quantization work for additional inference speedups, and its techniques could be coupled with other model-compression methods.The stated combination is presented as compatible with further compression approaches.

7 CONCLUSIONS

The paper argues that practical large-model inference requires scaling beyond single-server systems and using partitioning suited to latency and context requirements. It also identifies dense-model FLOP count and communication volume as fundamental performance constraints.

  • Conclusions: Scaling inference to 64+ chips enables the best latencies for 500B+ parameter models under challenging application requirements.The conclusion frames this as going beyond the traditional single-server inference paradigm.
  • Conclusions: Appropriate partitioning with multiquery attention reduces long-context memory costs and makes long-context inference practical.The proposed partitioning strategies are stated to generalize to single- and multi-node NVLink GPU topologies.
  • Limitations and future directions: FLOP count and communication volume can fundamentally limit inference performance in dense Transformer models.The paper points to sparsity, adaptive computation, and compressed chip-to-chip communication as routes for further gains.

A PARTITIONING STRATEGIES: DERIVING COMMUNICATION COSTS

The partitioning cost model derives communication times for collective operations from transferred tensor sizes and network bandwidth. It simplifies the formulas for large partition counts and relates all-reduce cost to all-gather cost.

  • All-gather communication transfers output chunks of size D across K partitions and K−1 interconnect links.
  • The cost model applies to most real-world network topologies, not only TPU torus topologies.
  • Reduce-scatter uses the larger input buffer to determine D, unlike all-gather, which uses the smaller output buffer.
  • All-reduce communication time is twice the communication time of all-gather.
  • For K ≫ 1, the (K−1)/K factor is approximated as 1, making reduce-scatter proportional to per-chip input and all-gather proportional to per-chip output.

A.2.1 Feedforward layer, 2D weight-stationary layout

The 2D weight-stationary layout partitions feedforward weights across model and feedforward dimensions, then sizes the torus axes to minimize communication time.

  • Weights use an ExFyz layout, partitioned across dmodel into X parts and dff into Y × Z parts with X × Y × Z = nchips.
  • With dff = 4 × dmodel, minimum communication occurs at X = 0.5 × √nchips and YZ = 2 × √nchips.
  • The resulting communication-time expression scales with 8BLE√nchips divided by network bandwidth.

A.2.2 Feedforward layer, weight-gathered layout

Weight-gathered layouts begin from the 2D weight-stationary arrangement, gather weights over selected torus axes, and trade weight communication against activation communication as batch size changes.

  • Weight-gathered layouts preserve the ExFyz starting arrangement and all-gather weight tensors over X and Y before einsums, with communication volume EF/Z.
  • The number of chips participating in weight gathering is N = X, XY, or XYZ for X-, XY-, or XYZ-weight-gathered layouts.
  • Weight communication scales as 2EF × N/nchips divided by network bandwidth.
  • Activation communication scales as 2BLE × N divided by network bandwidth.
  • The communication-optimal layout switches as batch size grows: 2D weight-stationary is best at low tokens per batch, while weight-gathered layouts become optimal at larger batches.

C MFU VS LATENCY TRADEOFF

MFU varies across inference phases, model sizes, latency regimes, and partitioning choices. Larger models generally achieve higher MFU, but model parallelism can reverse this pattern at long-latency decodes.

  • Decode MFU is typically much lower than prefill MFU, while prefill jumps mark transitions from 2D weight-stationary to XYZ weight-gathered layouts.
  • Larger models generally achieve higher MFU because larger matrix multiplies are more efficient.
  • At long-latency decodes, PaLM 62B achieves higher MFU than PaLM 540B, associated with 8-way versus 64-way model parallelism.

D FULL COMPARISON TO FASTERTRANSFORMER

The paper compares PaLM 540B and Megatron-Turing NLG 530B inference on 64 TPU v4 chips against FasterTransformer across three benchmark configurations. Results are summarized using globally defined latency–MFU Pareto frontiers, with batch sizes below 4 excluded because the strategy provides no speedup there.

  • Experimental setup: 64 TPU v4-chip runs compare PaLM 540B and Megatron-Turing NLG 530B implementations with FasterTransformer baselines.The comparison covers both model architectures and is preceded by their hyperparameter differences.
  • Benchmark configurations: Three benchmarks use 20 input and 8 output tokens, 60 input and 20 output tokens, or 128 input and 8 output tokens.Results are reported in Tables D.2–D.4, with all times in milliseconds.
  • Pareto analysis: Pareto-frontier annotations identify results that globally balance latency and MFU across all 500B-class results or across Megatron-Turing NLG results.The frontier is not a per-row comparison; a result is retained when no other result is both faster and higher-MFU.
  • Scope boundary: Batch sizes below 4 are omitted because the multiquery-attention partitioning strategy achieves no speedup below the minimum TPU v4 torus-axis size.The reported comparisons therefore do not cover smaller batch sizes.
  • Latency–MFU visualization: Figure C.1 plots MFU against latency at a 2048-token context length for generation and input-token processing.The left panel measures latency per generated token for 64 generated tokens after context processing; the right measures processing time for 2048 input tokens without output generation.
Loading 2211.05102v1…