PaLM: Scaling Language Modeling with Pathways
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, et al.
TL;DR: Google introduces PaLM, a massive 540-billion parameter autoregressive Transformer trained across 6,144 TPU v4 chips using the Pathways system. By avoiding pipeline parallelism and modifying the Transformer architecture, they achieve an unprecedented 46.2% Model FLOPs Utilization while unlocking breakthrough, discontinuous multi-step reasoning capabilities.
6,144 TPU v4 Chips
Trained across two TPU Pods without complex pipeline parallel bubbles.
46.2% MFU
High hardware efficiency achieved via Parallel Layers and Multi-Query Attention.
Emergent Reasoning
Discontinuous jumps in performance on complex logic and mathematics tasks.
Why This Matters
As Large Language Models (LLMs) grow, the engineering limits of hardware systems often dictate research directions. Traditionally, training models with hundreds of billions of parameters required pipeline parallelism. However, pipeline parallelism introduces "bubbles"—idle processor time while waiting for previous layers to compute—and demands heavy memory bandwidth to reload weights.
PaLM breaks this paradigm. By leveraging Google's Pathways execution engine, they managed to scale training to 6,144 TPU v4 chips across two geographically separated physical pods using only data and model parallelism. This demonstrates that we can continue scaling LLMs efficiently without the complexity and overhead of sequential pipelining.
The Big Idea
The core contribution of PaLM is twofold: a demonstration of highly efficient distributed scaling and the discovery of discontinuous capability jumps. While scaling from 8B to 62B parameters often follows predictable power-law scaling curves, moving to 540B parameters unlocks emergent reasoning capabilities—such as joke explanation, chain-of-thought mathematical deduction, and complex code repair—that are virtually non-existent at smaller scales.
How It Works
To achieve high training throughput, the authors combined a new distributed execution system (Pathways) with key architectural simplifications to the standard Transformer block.
1. Pathways Infrastructure
Rather than using a single monolithic coordinator, Pathways uses a client-server architecture to dispatch computations to multiple TPU Pods. Each Pod computes gradients on its half of the batch, and then performs a high-speed, bursty cross-pod gradient transfer over the data center network (DCN) at an aggregate burst rate of 81 Tbps.
Interactive: Standard vs. Parallel Transformer Block
PaLM uses a parallel layer formulation instead of the standard serialized one. This allows fusing the matrix multiplications of the Attention and MLP layers, accelerating training speed by ~15%.
Mathematical Formulation
Why it boosts efficiency:
Interactive: Multi-Query Attention (MQA) Memory Calc
Standard Multi-Head Attention stores unique Key (K) and Value (V) tensors for each attention head. In contrast, Multi-Query Attention shares a single Key and Value projection across all heads, dramatically reducing the KV Cache size during decoding.
Hyperparameters
Projection Sharing Scheme
KV Cache Memory Footprint
Key Results & Scaling Behavior
PaLM 540B achieves state-of-the-art few-shot performance on 28 out of 29 widely evaluated English NLP benchmarks. But the most interesting discovery is the nature of the scaling curves.
While many tasks improve log-linearly with scale (following standard scaling laws), approximately 25% of tasks in BIG-bench exhibit "discontinuous improvements". On these tasks, performance remains near random chance at 8B and 62B parameters, only to jump drastically at 540B.
Interactive Scaling Curve Explorer
Explore the difference between predictable log-linear scaling and emergent, discontinuous scaling across PaLM's three model sizes (8B, 62B, and 540B).
Loading chart...
Limitations & Training Spikes
Despite its success, training a 540B parameter model is fraught with challenges. The authors noted several key limitations and issues:
- Training Instability: During training of the 540B model, the loss spiked roughly 20 times despite gradient clipping. Because of the massive cost of training, they could not determine a principled mitigation strategy. Instead, they relied on a simple heuristic: rolling back to a checkpoint ~100 steps prior, skipping 200–500 batches of data, and resuming.
- Data Depletion: High-quality text data is finite. At 780B tokens, some subcorpora began to repeat. The authors' ablations showed that training longer on repeated data offers no downstream performance benefits, highlighting the impending bottleneck of high-quality data availability.
- Toxicity & Social Bias: While larger scale improves accuracy on pronoun resolution tasks (like Winogender), co-occurrence analysis reveals that PaLM continues to reflect and amplify societal stereotypes present in its training corpus.
Glossary
Model FLOPs Utilization (MFU)
A clean, implementation-independent metric proposed to replace Hardware FLOPs Utilization (HFU). MFU measures the ratio of the observed training throughput (tokens per second) to the theoretical maximum throughput of the system, excluding operations like rematerialization.
SwiGLU Activation
A variation of the Gated Linear Unit (GLU) activation function using Swish. Formulated as \(\text{Swish}(xW) \cdot xV\), it requires three matrix multiplications instead of two but consistently yields higher model quality.
Rotary Position Embeddings (RoPE)
A method of incorporating positional information by multiplying representation vectors by a rotation matrix. RoPE naturally generalizes to longer sequence lengths during inference.
Pipeline Parallelism
A distributed training strategy where the layers of a neural network are partitioned sequentially across devices. This creates execution "bubbles" where devices are idle while waiting for activations or gradients from other devices.
Cite this work
@article{chowdhery2022palm,
title={PaLM: Scaling Language Modeling with Pathways},
author={Chowdhery, Aakanksha and Narang, Sharan and Devlin, Jacob and Bosma, Maarten and Mishra, Gaurav and Roberts, Adam and Barham, Paul and Chung, Hyung Won and Sutton, Charles and Gehrmann, Sebastian and others},
journal={arXiv preprint arXiv:2204.02311},
year={2022}
}