VerifiedarXiv:2204.0231132 min
Large language models · Training systems

PaLM: Scaling Language Modeling with Pathways

A 540-billion-parameter language model trained on 6144 TPU chips at 46% of their peak throughput, then tested on hundreds of benchmarks.

PaLM is a decoder-only Transformer with a handful of changes that make it faster to train and to serve. Google trained it on 780 billion tokens across two TPU v4 pods with no pipeline parallelism, defined a new efficiency metric to report how well that went, and compared 8B, 62B and 540B versions to see which abilities grew with size.

Explaining the paperPaLM: Scaling Language Modeling with PathwaysChowdhery, Narang, Devlin, et al. · JMLR 2023 · arXiv:2204.02311 ↗

Posted on April 5, 2022, one week after the Chinchilla paper argued that models this large were being trained on too little text. Section 13 of the PaLM paper takes up that argument directly.

PaLM, the Pathways Language Model, continues the line of GPT-3: a Transformer trained only to predict the next token of text, then used by writing a prompt rather than by fine-tuning. Google trained three sizes, 8, 62 and 540 billion parameters, on the same data with the same recipe. The largest is dense, meaning every parameter takes part in processing every token, unlike mixture-of-experts models such as Switch Transformers that route each token through a small subset. At 540B it was slightly larger than the previous biggest dense model, Megatron-Turing NLG at 530B.

Four parts of the report explain the model, and this page covers them in order:

  1. the architecture: a standard Transformer with six changes, and where its 540 billion parameters sit;
  2. the training recipe: the data mixture, the Adafactor optimizer and its schedules, and the loss spikes;
  3. the system: how 6144 chips divided the work without pipeline parallelism, and MFU, the efficiency metric the paper introduced;
  4. the results: what improved smoothly with size, what jumped between 62B and 540B, and whether 780 billion tokens was enough data for a model this size.

What a decoder-only language model computes

Text first becomes a sequence of integer tokens. PaLM uses a SentencePiece vocabulary of 256,000 tokens built from its own training data. It is lossless: whitespace is kept exactly (which matters for code), any Unicode character missing from the vocabulary falls back to its UTF-8 bytes, each of which has a token, and numbers are always split into single digits, so 123.5 becomes the five tokens 1 2 3 . 5.

Training documents are concatenated, separated by a special [eod] token, and cut into sequences of exactly 2048 tokens, so no position is wasted on padding and a document can be split across two sequences. For each position tt the network outputs a probability for each of the 256,000 possible next tokens, computed from tokens 1 through tt only; a causal mask in attention blocks every later position. Training minimizes the average negative log-probability of the token that came next:

L(θ)=−1T∑t=1Tlog⁡pθ ⁣(xt∣x<t)\mathcal{L}(\theta) = -\frac{1}{T}\sum_{t=1}^{T}\log p_\theta\!\left(x_t \mid x_{<t}\right)
(1)

If the model gives the correct next token probability 0.25, that position contributes −ln⁡0.25≈1.39-\ln 0.25 \approx 1.39 nats (natural-log units) to the average. PaLM adds a small auxiliary term to (1), covered with the rest of the recipe below, and uses no label smoothing and no dropout.

A model trained this way is used few-shot. The prompt holds a short task description and kk worked examples, then the new question, and the model continues the text. "1-shot" means one example in the prompt, "5-shot" five, "0-shot" none. No weights change. The alternative in 2022 was the BERT and T5 route of fine-tuning on many labeled examples per task, which gives a separate copy of the model for each one; PaLM reports fine-tuning results too, but its headline numbers are few-shot.

How PaLM changes the Transformer block

The model is a stack of identical layers (118 at 540B), each reading and writing a vector of width dmodeld_\text{model} per token, called the residual stream. A standard pre-norm layer, as in GPT-2 and GPT-3, runs attention and then an MLP in series, each on a normalized copy of the stream, and adds each result back:

h=x+Attn(LN(x)),y=h+MLP(LN(h))h = x + \mathrm{Attn}(\mathrm{LN}(x)), \qquad y = h + \mathrm{MLP}(\mathrm{LN}(h))h=x+Attn(LN(x)),y=h+MLP(LN(h))\begin{gathered} h = x + \mathrm{Attn}(\mathrm{LN}(x)), \\ y = h + \mathrm{MLP}(\mathrm{LN}(h)) \end{gathered}
(2)

Substituting hh gives yy as xx plus the attention output plus an MLP that reads the stream after attention has written to it. (The paper prints this serial form with the attention term missing; the Provenance panel below has the details.) PaLM replaces it with the parallel formulation from GPT-J (Wang and Komatsuzaki, 2021), where both branches read the same normalized input:

y=x+Attn(LN(x))+MLP(LN(x))y = x + \mathrm{Attn}(\mathrm{LN}(x)) + \mathrm{MLP}(\mathrm{LN}(x))
(3)

In (2) the MLP cannot start until attention has finished. In (3) the first matrix multiplications of both branches take the same input, so they can be concatenated into one larger multiplication, and large multiplications keep a TPU closer to its peak rate. The paper reports roughly 15% faster training at large scale. In exchange, a layer's MLP no longer sees what that layer's attention just computed (the next layer does). In ablations this cost a little quality at 8B and none at 62B, and the authors extrapolated that it would be neutral at 540B without testing that size.

The other five changes, in the order the paper lists them:

SwiGLU MLP. The original Transformer MLP is ReLU(uW1)W2\mathrm{ReLU}(uW_1)W_2, two matrices. SwiGLU (Shazeer, 2020) multiplies two projections of the input elementwise, one passed through the Swish function:

MLP(u)=(Swish(uW1)⊙uW3) W2,Swish(z)=z σ(z)\mathrm{MLP}(u) = \big(\mathrm{Swish}(uW_1)\odot uW_3\big)\,W_2, \qquad \mathrm{Swish}(z) = z\,\sigma(z)MLP(u)=(Swish(uW1)⊙uW3) W2,Swish(z)=z σ(z)\begin{gathered} \mathrm{MLP}(u) = \big(\mathrm{Swish}(uW_1)\odot uW_3\big)\,W_2, \\ \mathrm{Swish}(z) = z\,\sigma(z) \end{gathered}
(4)

where σ\sigma is the logistic sigmoid and ⊙\odot is elementwise multiplication. There are now three matrices. Shazeer kept parameter count and compute equal to the two-matrix baseline by shrinking the hidden width to two thirds (from 3072 to 2048 in his experiments). PaLM does not shrink it: the hidden width is 4 dmodel4\,d_\text{model}, 73,728 at 540B, so each layer's MLP holds 3 × 18,432 × 73,728, about 4.08 billion parameters, 50% more than a two-matrix MLP of that width. The paper cites Shazeer's compute-matched comparison as the reason to use SwiGLU at all.

Multi-query attention. Standard multi-head attention projects each token into kk heads of query, key and value vectors, each of size hh. PaLM keeps the kk query heads but projects a single key and a single value of size hh that every head shares (Shazeer, 2019; see the multi-query attention explainer). The paper found this neutral for quality and training speed. The gain comes at generation time: a decoder generates one token at a time, and each new token attends to the keys and values of all earlier tokens, which are kept in a cache and read from memory at every step. For PaLM 540B (k=48k = 48, h=256h = 256, 118 layers) the cache per token shrinks from 2×48×2562 \times 48 \times 256 to 2×2562 \times 256 numbers per layer. Stored in bfloat16 that is 120,832 bytes per token instead of 5.8 MB, or 247 MB instead of 11.9 GB for one 2048-token sequence.

RoPE. Instead of adding a position vector to each token, rotary position embeddings rotate each query and key by angles proportional to the token's position, so the attention score between two tokens depends on their positions only through the distance between them (RoPE explainer); the paper chose it for its reported performance on long sequences.

Shared input and output embeddings. One 256,000 × dmodeld_\text{model} matrix both maps token ids to vectors at the input and maps the final vectors to 256,000 scores (logits) at the output. At 540B the matrix holds 4.72 billion parameters, stored once.

No biases. No dense layer and no layer norm has an additive bias term; layer norms keep only their scale. The paper reports that this made large models train more stably and gives no further mechanism.

Put together, one 540B layer computes the following, with shapes for a batch of sequences:

# One PaLM 540B layer, parallel form. x: [batch, seq, 18432]
u = layer_norm(x)                    # scale only: no bias anywhere in PaLM
q = (u @ Wq).reshape(b, s, 48, 256)  # 48 query heads of size 256
k = u @ Wk                           # [b, s, 256]: ONE key head, shared
v = u @ Wv                           # [b, s, 256]: ONE value head, shared
q, k = rope(q), rope(k)              # rotate by position
a = causal_attention(q, k, v)        # all 48 heads read the same k and v
attn_out = a.reshape(b, s, 12288) @ Wo           # back to width 18432
mlp_out = (swish(u @ W1) * (u @ W3)) @ W2        # SwiGLU, hidden 73728
y = x + attn_out + mlp_out           # one residual add for both branches
# Wq, Wk, Wv, W1 and W3 all multiply u, so they run as one fused matmul
# with 12288 + 256 + 256 + 73728 + 73728 = 160,256 output columns.

The query heads have total width 48 × 256 = 12,288, two thirds of dmodeld_\text{model} = 18,432; the output projection Wo maps back up to 18,432. Every new token in generation runs exactly this code once per layer, reading the cached k and v of all earlier tokens.

Counting the 540 billion parameters

The paper's Table 1 gives layers, heads and dmodeld_\text{model} for each size, but not the head size or the MLP width. A later Google paper on serving PaLM (Pope et al., 2022) lists them: head size h=256h = 256 and MLP width f=4 dmodelf = 4\,d_\text{model}. With kk query heads, one layer holds

Player=2 d (kh)+2 d h+3 d fP_\text{layer} = 2\,d\,(k h) + 2\,d\,h + 3\,d\,f
(5)

with d=dmodeld = d_\text{model}. The first term is the query and output projections, the second the one shared key and value projection, the third the three SwiGLU matrices. For 540B: 2 × 18,432 × 12,288 = 452,984,832, plus 2 × 18,432 × 256 = 9,437,184, plus 3 × 18,432 × 73,728 = 4,076,863,488, which is 4,539,285,504 per layer. Times 118 layers, plus the 256,000 × 18,432 embedding, gives 540,354,281,472, the paper's 540.35 billion. The same formula gives 8.63B and 62.50B for the smaller models. Layer-norm scales add about 2 million more, below the rounding.

Figure 1 shows the breakdown. The MLP holds 89% of PaLM 540B's parameters. Switch to multi-head attention and the model grows by 52 billion parameters, while the key/value cache per token grows 48-fold; switch to a compute-matched SwiGLU width and it shrinks to 380 billion.

Figure 1 · where the parameters are
Parameter count from (5) for each PaLM size. Only PaLM's own choices (multi-query, SwiGLU at width 4·dmodeld_\text{model}) reproduce Table 1. The key/value segment is the part multi-query attention removes, and the readout below the bar gives the per-token cache that generation reads.

Since every parameter is used once per token in the forward pass as a multiply and an add, a forward pass costs about 2N2N floating-point operations (FLOPs) per token for a model with NN parameters, about 1.08 trillion for PaLM 540B. Section 2.1 of the paper says "the number of FLOPs per token is approximately equal to the number of parameters", which holds if a multiply-add counts as one operation; Appendix B counts it as two, as this page does.

Data, optimizer and schedules

Data. The 780 billion training tokens come from the corpora used for LaMDA and GLaM: 50% multilingual social media conversations, 27% filtered multilingual webpages, 13% English books, 5% GitHub code, 4% multilingual Wikipedia and 1% English news. Webpages were scored by a quality classifier trained to recognize pages like known high-quality collections, then sampled in proportion to that score, so low-scoring pages appear less often but are not removed. The code is 196 GB of source in 24 languages from GitHub repositories whose licenses are not copyleft, deduplicated by edit distance. About 22% of the data is not English. All three model sizes see exactly one epoch of the same shuffled data, and the mixture proportions were chosen so no source repeats within it.

Initialization. Each weight matrix is drawn with standard deviation 1/nin1/\sqrt{n_\text{in}}, where ninn_\text{in} is its input width: about 0.0074 for the 540B model's 18,432-wide inputs. The embedding is drawn with standard deviation 1, because no layer norm is applied to it, and since the same matrix produces the output logits, those logits are multiplied by 1/dmodel1/\sqrt{d_\text{model}} to bring them back to a small scale.

Optimizer. PaLM trains with Adafactor (Shazeer and Stern, 2018) without its factorization, which makes it close to Adam with one difference, parameter scaling: each matrix's step is multiplied by the root-mean-square (RMS) size of that matrix's own entries. At initialization a hidden-layer matrix has RMS about 0.0074 and the embedding about 1, so with the same nominal learning rate the embedding steps 135 times further in absolute terms, and in relative terms both change by the same fraction of their size. The paper compares this to GPT-3's practice of lowering Adam's learning rate by hand for larger models, with the difference that the embeddings and layer-norm scales, which sit at a different size, are not scaled down with the rest.

The step size and the second-moment decay both follow Adafactor's proposed defaults, as functions of the step number kk:

ρk=min⁡ ⁣(10−2, 1k),β2,k=1−k−0.8\rho_k = \min\!\left(10^{-2},\ \tfrac{1}{\sqrt{k}}\right), \qquad \beta_{2,k} = 1 - k^{-0.8}ρk=min⁡ ⁣(10−2, 1k),β2,k=1−k−0.8\begin{gathered} \rho_k = \min\!\left(10^{-2},\ \tfrac{1}{\sqrt{k}}\right), \\ \beta_{2,k} = 1 - k^{-0.8} \end{gathered}
(6)

The paper's description, "10−2 for the first 10,000 steps, then decayed at a rate of 1/k1/\sqrt{k}", is the same schedule: 1/10,0001/\sqrt{10{,}000} is exactly 0.01, so the two pieces meet at step 10,000, and by the final step 255,000 the step size is 0.00198. The paper describes no warmup; the step size is flat from step 1. The decay rate β2\beta_2 sets how many recent steps the second-moment average covers, about 1/(1−β2)1/(1-\beta_2), which under (6) equals k0.8k^{0.8}. At step 1, β2=0\beta_2 = 0 and the estimate is the current squared gradient alone, so no bias correction is needed; by step 10 the average spans about 6 steps, by step 10,000 about 1,600, and at the end about 21,000. The paper prefers this to a fixed β2=0.99\beta_2 = 0.99 (a 100-step window) because rare tokens' embedding rows receive gradients in few batches, and a short window estimates their second moment badly.

The full update for one matrix, with the remaining settings, fits in a few lines:

# Adafactor as PaLM ran it (unfactored, with momentum), for one weight matrix W
# at step k. g is the gradient after global-norm clipping to 1.0.
rho   = min(1e-2, 1 / sqrt(k))           # "learning rate": a RELATIVE step size
beta2 = 1 - k ** -0.8                    # 0 at k = 1, 0.99995 at k = 255,000
v     = beta2 * v + (1 - beta2) * g * g  # second moment, one entry per weight
m     = 0.9 * m + 0.1 * g / sqrt(v)      # momentum on the normalized gradient
scale = max(1e-3, rms(W))                # parameter scaling
W     = W - rho * scale * m              # step is proportional to W's own size
W     = W - rho ** 2 * W                 # "dynamic weight decay of lr^2"

Gradients are clipped to a global norm of 1.0 before the update, and weight decay is "dynamic": the coefficient is the square of the current step size, 10−410^{-4} early on and 3.9×10−63.9 \times 10^{-6} at the end. (Where the paper is silent, such as where momentum enters and how the decay is applied, the listing follows the Adafactor paper's algorithm with momentum added; treat those lines as a reconstruction.)

z-loss. The softmax turns the 256,000 logits ziz_i into probabilities by dividing ezie^{z_i} by the normalizer ZZ. PaLM adds to (1) a penalty on the size of log⁡Z\log Z:

Lz=10−4 log⁡2Z,Z=∑i=1256,000ezi\mathcal{L}_z = 10^{-4}\,\log^2 Z, \qquad Z = \sum_{i=1}^{256{,}000} e^{z_i}Lz=10−4 log⁡2Z,Z=∑i=1256,000ezi\begin{gathered} \mathcal{L}_z = 10^{-4}\,\log^2 Z, \\ Z = \sum_{i=1}^{256{,}000} e^{z_i} \end{gathered}
(7)

Adding the same constant to every logit leaves the probabilities unchanged, so the main loss (1) puts no limit on how far the logits drift as a group. The penalty pins that free direction near log⁡Z=0\log Z = 0: at log⁡Z=5\log Z = 5 it costs 10−4 × 25 = 0.0025 nats, small next to the main loss, but its gradient grows with log⁡Z\log Z and keeps pulling the logits back; the paper reports that the term made training more stable.

Batch size. The 540B model trained with 512 sequences per batch (1M tokens) until step 50,000, 1024 (2M tokens) until step 115,000, and 2048 (4M tokens) until step 255,000. Small batches make more progress per token early in training; later, when the useful signal in each gradient is smaller relative to its noise, larger batches give better gradient estimates, and they also make larger matrix multiplications that run more efficiently on the TPUs. Summing the schedule gives 775.9 billion tokens, the paper's "780 billion".

Figure 2 plays all of these schedules against the step number. The step axis is logarithmic, so the first 10,000 steps, where the step size is flat and β2\beta_2 climbs fastest, take up most of the width, and the 140,000 steps at batch 2048 are squeezed into the right edge.

Figure 2 · the training schedules
505
PaLM 540B's step size (6), the window of steps the second-moment average covers, and the batch size, against the step on a log axis. Press Play or drag the slider; the readouts give the exact values and the tokens seen so far.

Loss spikes. The 540B run's loss jumped sharply about 20 times, at irregular steps, some late in training, despite gradient clipping; the smaller models had none. Without a known cause, the team restarted from a checkpoint about 100 steps before each spike and skipped about 200 to 500 batches around it (at 4M tokens a batch, up to about 2 billion tokens, a quarter of a percent of the data), and the loss did not spike again at that point. When they fed the same skipped batches to the model from a different, earlier checkpoint, no spike occurred either, so the authors attribute spikes to the combination of particular batches with a particular parameter state, not to bad data alone. The restarts are exactly reproducible because training is bitwise deterministic: the contents of every batch are a function of the step number, so a rerun from a checkpoint reproduces the original run exactly.

How PaLM was trained on 6144 chips without pipeline parallelism

PaLM 540B needs about 1.08 TB just to hold its weights in a 2-byte format, far more than one chip's memory, and its training needs thousands of chips working at once. There are three standard ways to divide the work:

Pipelines need little bandwidth between stages, so they are the usual choice across slow links: Gopher pipelined across four TPU v3 pods, and Megatron-Turing NLG combined pipeline, model and data parallelism on 2240 A100 GPUs. They cost two things the paper names. While the pipeline fills at the start of each step and drains at the end, most stages sit idle (the bubble), and every micro-batch reloads the stage's weights from memory.

Inside one pod. A TPU v4 pod here is 3072 chips attached to 768 host machines, joined by the pod's fast chip-to-chip network. Every weight tensor is divided over all 3072 chips in a 12 × 256 grid: 12-way model parallelism along one axis and 256-way fully sharded data parallelism along the other. Plain data parallelism would give each of the 256 chips that share a model-parallel slice a full copy of that slice; fully sharded data parallelism splits it 256 ways instead, the idea behind ZeRO. So each chip permanently holds 1/3072 of every matrix, about 176 million of the 540 billion parameters. In the forward pass, just before a layer runs, its weights are all-gathered along the 256-way axis, so each chip briefly has the full 1/12 model-parallel slice it needs, and they are discarded again after use.

Activations, the intermediate values the backward pass needs, are kept sharded too. From each layer the forward pass saves one fully sharded activation tensor; everything else is recomputed during the backward pass, a technique called rematerialization or activation recomputation. Recomputing costs FLOPs, but the freed memory allows a larger batch, and the paper found this gave higher throughput than storing more.

Across two pods. The two pods are connected only by the ordinary datacenter network, much slower than the links inside a pod. Pathways, Google's orchestration system, runs a single Python client that sends half of each batch to each pod and runs one step as three pieces of work:

  1. component A on both pods at once: forward and backward pass on that pod's half, including the gradient reductions within the pod;
  2. a transfer: each pod sends its gradients to the other;
  3. component B on both pods: add local and remote gradients and apply the optimizer update, so both pods end the step with bitwise-identical parameters.

Because both pods shard the weights the same way, each host only needs the gradient shard held by its counterpart in the other pod, so the 768 host pairs exchange data one-to-one, about 1.3 GB per pair per step. They all send at the same moment, the end of the backward pass, which makes an aggregate burst of 81 terabits per second; the Pathways network stack splits each transfer into small chunks routed over many links to avoid congestion. Two pods delivered 1.95 times the throughput of one, 97% of the ideal 2×, at twice the batch size. The paper attributes the missing 3% to the transfer starting only after the backward pass ends instead of overlapping with it.

How the paper measures training efficiency (MFU)

Earlier reports measured efficiency as "hardware FLOPs utilization" (HFU): FLOPs the chips executed divided by their peak. The paper argues this number is not comparable across systems. A run that recomputes more activations executes more FLOPs and reports a higher HFU without producing any more training, and the count depends on how a given system tallies its operations.

The paper's "model FLOPs utilization" (MFU) counts only the FLOPs the model itself requires, the same for any implementation, and divides the observed token throughput by the throughput the hardware could reach if it ran those FLOPs at peak. For a model with NN parameters, LL layers, HH heads of size QQ and sequence length TT, on hardware with total peak PP FLOP/s:

MFU=tokens/s×(6N+12LHQT)P\text{MFU} = \frac{\text{tokens/s} \times (6N + 12LHQT)}{P}
(8)

Here a multiply and an add count as two FLOPs. The 6N6N is the matrix-multiplication cost per token: each parameter does one multiply-add in the forward pass (2N2N), and the backward pass needs two matrix multiplications for each forward one, one for the gradient with respect to the layer's input and one for the gradient with respect to its weights (4N4N). The 12LHQT12LHQT covers the two attention products, QK⊤QK^\top and the weighted sum of values, which hold no parameters. For PaLM 540B the attention term is 12 × 118 × 48 × 256 × 2048 ≈ 3.56 × 1010, 1.1% of 6N6N, which is 3.24 × 1012.

The worked numbers: PaLM 540B averaged 238,300 tokens per second at batch 2048. Times 6N6N with N=540N = 540 billion, that is 7.72×10177.72 \times 10^{17} FLOP/s. The peak is 6144 chips × 275 TFLOP/s = 1.69×10181.69 \times 10^{18} FLOP/s. The ratio is 45.7%, and with the attention term 46.2%. A batch of 2048 × 2048 = 4.19 million tokens at that rate takes 17.6 seconds per step.

HFU for the same run is 57.8%. The extra comes from rematerialization: 75% of the non-attention forward pass is recomputed, adding 1.5N1.5N (75% of 2N2N) to the 6N6N, a factor of 1.25, and all of the attention forward pass is recomputed too. The paper's Table 22 lists 4.10 versus 3.28 TFLOPs per token for the two counts, and 46.2% × 1.25 = 57.8%.

In Figure 3, drag the throughput and the recompute fraction. Recomputation moves only the HFU bar; MFU moves only with tokens per second. Switch to Megatron-Turing NLG to reproduce its 29.7% (30.2% with attention) from 65,430 tokens per second on 2240 A100s at 312 TFLOP/s each.

Figure 3 · MFU and HFU
238.3K
75%
MFU from (8) and HFU, which also credits recomputed activations, for one training run. Defaults are each paper's reported throughput; at the right end of the throughput slider MFU without attention reaches 100% and HFU passes the hardware peak, which no real run can do.

Table 3 of the paper puts PaLM's 46.2% beside GPT-3 at 21.3% on V100 GPUs, Gopher at 32.5% on 4096 TPU v3 chips and Megatron-Turing NLG at 30.2%. The throughputs behind the GPT-3 and Gopher figures come from secondary reports (Gopher's from personal communication with its authors). The paper credits its margin to the pipeline-free parallelism, XLA compiler optimizations and the parallel layers of (3).

What 540 billion parameters changed on benchmarks

The evaluation runs to hundreds of tasks. The main results, each against the best prior result the paper could find:

English NLP. On the 29 benchmarks GPT-3 and GLaM used (question answering, cloze, Winograd-style, common sense, reading comprehension, SuperGLUE, natural language inference), PaLM 540B beat the best prior pretrained model on 28 of 29 few-shot and 24 of 29 one-shot. Averaged 1-shot, it scored 63.9 on the generation tasks and 74.7 on the understanding tasks, against GLaM's 58.4 and 68.7 and GPT-3's 52.9 and 65.4. PaLM 62B, at 57.7 and 67.3, already beats GPT-3 175B. PaLM 540B also beat Megatron-Turing NLG, a model of nearly the same size, on every task, which the paper takes as evidence that the data and training recipe matter, not only size.

MMLU, 57 multiple-choice subjects from mathematics to law, 5-shot: 69.3%, against Chinchilla's 67.5%. PaLM 62B scores 53.7% and PaLM 8B 25.3%, which is chance for four options.

Fine-tuned SuperGLUE. Fine-tuned, PaLM 540B averages 92.6 on the validation sets, between T5-11B (89.9) and the mixture-of-experts ST-MoE-32B (93.2), both encoder-decoder models, the architecture that tends to win classification fine-tuning at equal cost. Few-shot PaLM is well below its own fine-tuned scores (WiC 64.6 against 78.8, for example).

BIG-bench, a crowd-sourced suite of more than 150 tasks designed to be hard for language models. On the 58 tasks also run by GPT-3, Gopher and Chinchilla, 5-shot PaLM 540B scores above all three and above the average of the human raters, who were allowed search engines and calculators. It beats the prior best on 44 of the 58. Averaged human performance still beats it on 35% of the individual tasks.

Multi-step reasoning. Combined with chain-of-thought prompting, where the 8 examples in the prompt include written-out intermediate steps, PaLM 540B set the state of the art on GSM8K, a set of grade-school math word problems. The paper's Table 10:

Model and method (GSM8K)Accuracy
PaLM 540B, chain of thought, external calculator58%
PaLM 540B, chain of thought54%
PaLM 540B, no chain of thought17%
PaLM 62B, chain of thought33%
GPT-3 fine-tuned, chain of thought, calculator34%
GPT-3 fine-tuned, chain of thought, calculator, verifier (prior best)55%

The calculator is applied after the fact to the arithmetic in the generated steps and added at most 5 points on any dataset. Neither scale nor chain of thought is enough alone: 62B with chain of thought gets 33%, 540B without it 17%. Across seven reasoning datasets PaLM 540B with chain of thought reached the state of the art on four and came close on the other three, against prior systems that used task-specific fine-tuning, architectures or verifiers.

Code. About 5% of the training data, 39 billion tokens, is code, and only about 2.7 billion of that is Python. Even so, few-shot PaLM 540B matches the 12B Codex model on HumanEval (writing a Python function from its docstring), which was trained on about 100 billion Python tokens, roughly 50 times more. Fine-tuned on 7.75 billion further tokens, mostly Python, the resulting PaLM-Coder 540B solves 88.4% of HumanEval problems with 100 samples per problem (pass@100) and 80.8% of MBPP with 80. On DeepFix, repairing student C programs that fail to compile, it reaches an 82.1% compile rate against 71.7% for the prior best.

Discontinuous improvements between 62B and 540B

With three sizes spaced roughly 8× apart, the paper can ask whether the step from 62B to 540B gives about the same gain as the step from 8B to 62B. If a score rose by the same amount for every 8× increase in size, a straight line against the logarithm of size, the two gains would match; the paper cites the power-law trends of Kaplan et al. as the reason to expect that. Its definition of a discontinuity: project the 540B score as one more step of the same size, a62+(a62−a8)a_{62} + (a_{62} - a_8), and subtract that projection from the actual 540B score.

On BIG-bench's logical_sequence task (ordering events, months or numbers) the three models score 13%, 25% and 87%. The projection is 25 + 12 = 37%, so the discontinuity is 87 − 37 = +50 points. Over all 150 tasks, 25% had a discontinuity above +10 and 15% above +20. On english_proverbs (pick the proverb that fits a short story) 62B scores 25% and 540B 87%. On the two-sentence version of cause_and_effect the 8B model is near chance and 540B above 90%. Scale did not help everywhere: on navigate and mathematical_induction 540B only modestly beats 62B, and both stay far below the best human score.

Three things limit what the definition shows. The model sizes are not exactly evenly spaced on the log scale (62B to 540B is a factor of 8.6, 8B to 62B a factor of 7.2), so a strictly log-linear projection for logical_sequence is 38.1%, not 37%; that changes no conclusion. With three points there is no way to tell whether the curve between 62B and 540B rises sharply or gradually. And the measured quantity is usually accuracy, which gives no credit for a nearly right answer. Schaeffer, Miranda and Koyejo (2023) showed that such all-or-nothing metrics can turn a smooth improvement into an apparent jump; the paper does not test whether a continuous metric would show the same curve on these tasks.

Figure 4 has both cases. In the first mode, drag the three points (or use the sliders) and the readout applies the paper's rule. In the second, a made-up per-step accuracy improves in a straight line with log size, from 60% at 8B to 75% at 62B and 91% at 540B, and a task counts as solved only if all LL steps are right, so its accuracy is the per-step accuracy to the power LL. At L=1L = 1 the rule reports +1.3; at L=8L = 8 the same per-step skill registers as +30.

Figure 4 · the discontinuity rule
13%
25%
87%
Accuracy at the three PaLM sizes on a log size axis, the dashed one-more-step projection, and the gap the paper calls a discontinuity. "All-or-nothing scoring" is an illustration with synthetic numbers: a per-step accuracy that rises in a straight line, raised to the power LL.

PaLM and Chinchilla: model size versus training tokens

One week before PaLM appeared, Chinchilla (Hoffmann et al., 2022) reported that at a fixed training budget, a smaller model trained on more tokens reaches a lower loss: their 70B model on 1.4 trillion tokens beat the 280B Gopher on 300 billion at similar cost. The paper estimates training compute the same way, as

C≈6 N DC \approx 6\,N\,D
(9)

for NN parameters and DD training tokens (the 6N6N per token of (8)). For PaLM 540B that is 6 × 540 billion × 780 billion ≈ 2.53 × 1024 FLOPs, at 1.44 tokens per parameter. Chinchilla's results are often summarized as about 20 tokens per parameter; at this budget that rule points to roughly 145 billion parameters on 2.9 trillion tokens.

Section 13 of the paper poses the question in its own terms: how would 62B on 7 trillion tokens, 120B on 3.6 trillion, or 240B on 1.8 trillion compare with 540B on 780 billion? All four cost the same to train by (9), and a smaller model would cost proportionally less to serve, since each generated token costs about 2N2N FLOPs. The authors did not run these, and give three reasons. The experiment is as expensive as PaLM itself. Training a smaller model fast on the same 6144 chips would need a much larger batch than the 4M tokens already used, which may hurt sample efficiency. And with their mixture, some sources start repeating after about 780 billion tokens.

Figure 5 draws that budget as a line. Every point on the teal line costs the same as PaLM 540B; drag along it to read the token count each model size would need, and where it crosses the point at which PaLM's data runs out.

Figure 5 · one training budget, many model sizes
540B
Parameters against training tokens, both on log axes. The teal line is every model costing 2.53×10242.53 \times 10^{24} FLOPs by (9); the hollow circles are the paper's three alternatives, the cross is the 20-tokens-per-parameter point, and the dashed amber line is where PaLM's data mixture starts to repeat.

The paper compares where it can. On nine English tasks that both papers report, PaLM 540B averages 70.3, Chinchilla 65.2, PaLM 62B 64.8 and Gopher 59.5. Plotted against training FLOPs, Chinchilla sits a little above the curve through the PaLM models on 58 BIG-bench tasks and a little below it on the nine English tasks, while Gopher is well below both, which agrees with Chinchilla's claim that Gopher was undertrained. Because PaLM used a different corpus, the comparison cannot separate size, tokens and data quality.

Appendix F adds one direct test. PaLM 62B, resumed at 600 billion tokens and trained to 1.325 trillion on a refreshed mixture (new data instead of repeats, plus arXiv and math webpages), improves on every benchmark reported: MMLU from 53.7% to 62.8%, BIG-bench from 32.4 to 40.8, GSM8K from 33.0% to 48.7%. Its nine-task English average rises to 66.5, above Chinchilla's 65.2, while continuing on the original mixture, with its repeated sources, gave much worse few-shot results.

Memorization, bias and limits

Memorization. Prompted with the first 50 tokens of a random training passage and decoded greedily, PaLM 540B reproduced the next 50 tokens exactly for 2.4% of passages, PaLM 8B for 1.6%. The rate on held-out text from the same distribution is above zero because some held-out text, such as an open-source license with only the year changed, nearly duplicates training text. Repetition drives most of it: for the 540B model, 100-token spans that occur once in the training data were memorized 0.75% of the time, spans seen more than 500 times over 40%.

Contamination. Ten of the 29 English benchmarks are partly on the open web, either wholesale (SQuADv2, Winograd) or because their examples were extracted from web text (LAMBADA, ReCoRD, Web Questions). For those the paper split each evaluation set into a clean subset, where under 70% of an example's 8-grams appear in the training data, and the rest. Scores on the clean subsets went up about as often as down, and no more often down for 540B than for 8B. If the larger model had memorized the test sets, its clean-subset scores would have dropped more. For BIG-bench it checked that the benchmark's canary string never appears in the training data.

Bias and toxicity. Accuracy on Winogender, a test of gender bias in resolving pronouns, improves with scale, and PaLM 540B set a new 1-shot and few-shot best. Prompt-continuation analysis found the model associating Muslims with terrorism, extremism and violence at every size. The toxicity of PaLM's continuations tracks the toxicity of the prompt much more closely than human-written continuations do, and the 62B and 540B models are slightly more toxic overall than 8B. The analyses are in English only, which the paper lists as a limitation.

The broader limits are the paper's own. It varies model size and holds data fixed, so it cannot say how much of PaLM's gain over earlier models comes from the 540 billion parameters and how much from the cleaner corpus and recipe; the 62B results suggest the second is large. Neither weights nor code were released, so the results cannot be reproduced outside Google. The loss spikes were worked around by restarting and skipping data, and the paper leaves a principled fix to future work.

Provenance Verified against primary literatureHow we verify
Chowdhery et al. (2022), Sec. 2 and Table 1Six changes to a decoder-only Transformer (SwiGLU, parallel layers, multi-query attention, RoPE, shared input/output embeddings, no biases) and a 256k SentencePiece vocabulary. Layers, heads and d_model for 8B, 62B and 540B; batch sizes 256 to 2048.
Pope et al. (2022), Table D.1The paper does not print the head size or the MLP width. Pope et al. list PaLM 540B with d_ff = 73728 and d_head = 256. With a 256,000-token vocabulary these reproduce Table 1 exactly (8.63, 62.50, 540.35 billion) when only the weight matrices are counted; this page computed that check.
Shazeer (2020), GLU Variants Improve TransformerFFN_SwiGLU(x) = (Swish(xW) * xV) W2. To hold parameters and compute fixed, the paper shrinks d_ff from 3072 to 2048. PaLM keeps d_ff = 4 d_model.
Shazeer & Stern (2018), Algorithm 4 and 6Adafactor: alpha_t = max(eps2, RMS(X_{t-1})) rho_t, with proposed rho_t = min(1e-2, 1/sqrt t), beta2_t = 1 - t^-0.8, eps2 = 1e-3. PaLM's learning-rate and beta2 schedules are these defaults; PaLM adds momentum 0.9 and turns off factorization.
Chowdhery et al. (2022), Sec. 4 and App. BWithin a pod, 12-way model parallelism times 256-way fully sharded data parallelism over 3072 chips; across pods, Pathways two-way data parallelism, 1.95x the throughput of one pod. MFU = tokens/s x (6N + 12LHQT) / peak FLOP/s; 238.3K tokens/s on 6144 chips at 275 TFLOP/s gives 45.7% (46.2% with attention); HFU with rematerialization 57.8%.
Chowdhery et al. (2022), Sec. 6.2 and Sec. 13, App. FDiscontinuity defined as the 540B score minus a62 + (a62 - a8); 25% of 150 BIG-bench tasks exceed +10. Training FLOPs as 6 x params x tokens; PaLM 62B trained to 1.325T tokens on refreshed data.
Schaeffer, Miranda & Koyejo (2023)Argue that nonlinear or discontinuous metrics, such as exact match over a multi-token answer, produce apparent emergent abilities that a continuous metric would show as smooth. Figure 4's second mode illustrates the argument with synthetic numbers.
CodeNo official PaLM code or weights were released. The parallel block is the GPT-J design: mesh-transformer-jax, mesh_transformer/layers.py, TransformerLayerShard returns attn_out + dense_out computed from one norm(x). The serial reference is openai/gpt-2 src/model.py, block(): x = x + a; x = x + m.
correctionSection 2 prints the standard ("serialized") block as y = x + MLP(LayerNorm(x + Attention(LayerNorm(x))), which is missing a closing parenthesis and, more importantly, the attention residual: in it the attention output reaches y only through the MLP. The standard pre-norm block, as in GPT-2's block(), adds both branches to the residual stream: y = x + Attention(LayerNorm(x)) + MLP(LayerNorm(x + Attention(LayerNorm(x)))). This page teaches that form as (2). The paper's parallel form (3) is correct as printed. Two smaller notes. Section 5 writes the initialization as W ~ N(0, 1/sqrt(n_in)), where the second argument is a standard deviation (fan-in variance scaling means variance 1/n_in); read as a variance it would be wrong. And the claim that parallel layers are quality-neutral at 540B is an extrapolation from 8B (small degradation) and 62B (none); it was not measured at 540B.

Questions you might still have

?

Were PaLM's weights or code released?
No. The paper describes the model and training setup but released neither weights nor training code. The head size (256) and MLP width (73,728) used on this page come from a later Google paper on serving PaLM (Pope et al., 2022), and they reproduce the paper's parameter counts exactly.

?

Why did PaLM avoid pipeline parallelism?
A pipeline splits the layers into stages on different devices and feeds micro-batches through them, so stages sit idle while the pipeline fills and drains (the bubble), and each micro-batch reloads the weights. A TPU v4 pod is large enough (3072 chips on a fast interconnect) to hold the whole model with model and data parallelism alone, and the two pods only exchange gradients once per step, which tolerates the slower datacenter network between them.

?

Is an MFU of 46% low?
Not by the standards of its time. By the same formula the paper computes GPT-3 at 21.3%, Gopher at 32.5% and Megatron-Turing NLG at 30.2%. The remaining time goes to work MFU does not credit: communication between chips, recomputing activations (PaLM's hardware utilization, which counts recomputation, is 57.8%), and operations such as layer norms and softmax that move many bytes per FLOP.

?

Why use Adafactor instead of Adam?
Run without factorization, Adafactor stores the same per-weight second moment as Adam, so the memory saving it was designed for is not used. What PaLM takes from it is parameter scaling, which makes each matrix's step proportional to the matrix's own root-mean-square size, and the default schedules: step size min(0.01, 1/sqrt(k)) and beta2 = 1 - k^-0.8.

?

Do the discontinuous jumps prove that abilities emerge at 540B?
They show that on about a quarter of BIG-bench tasks the 62B-to-540B gain was much larger than the 8B-to-62B gain. With only three model sizes the shape between them is unknown, and Schaeffer, Miranda and Koyejo (2023) showed that all-or-nothing metrics such as exact match can turn a smooth improvement into an apparent jump. The chain-of-thought explainer on this site discusses the same question for reasoning tasks.

?

If multi-query attention does not speed up training, why use it?
Because it speeds up generation. Each new token attends to the cached keys and values of every earlier token, so decoding reads that cache once per token. With one shared key and value head instead of 48, PaLM 540B's cache is 48 times smaller: about 121 kB per token in bfloat16 instead of 5.8 MB. The multi-query attention explainer covers the quality and speed tradeoff.

?

How long and how much energy did training take?
The paper reports 1200 hours on 6144 TPU v4 chips plus 336 hours on 3072 chips, including downtime and repeated steps, about 2.56 x 10^24 FLOPs in total. At a measured 378.5 W per chip in a datacenter running on 89% carbon-free energy, it estimates 271.43 tonnes of CO2-equivalent.

?

What is the 62B model good for in this paper?
It is the comparison point that makes scale measurable: same data, same recipe, roughly a ninth of the parameters. It already beats GPT-3 175B on the 1-shot English averages, and trained longer on fresh data (1.325 trillion tokens) it beats Chinchilla on nine English tasks, which the paper uses to argue that data quality and token count matter alongside size.

Footnotes & further reading

  1. The paper: Chowdhery et al., PaLM: Scaling Language Modeling with Pathways (arXiv April 2022, v5 October 2022; JMLR 2023).
  2. The system: Barham et al., Pathways: Asynchronous Distributed Dataflow for ML (MLSys 2022). The within-pod sharding is from Xu et al., GSPMD (2021).
  3. Head size and MLP width: Pope et al., Efficiently Scaling Transformer Inference (MLSys 2023), Table D.1.
  4. Architecture sources: Shazeer, GLU Variants Improve Transformer (2020); Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (2019), covered in the multi-query attention explainer; the parallel block of GPT-J (mesh-transformer-jax).
  5. The optimizer: Shazeer & Stern, Adafactor: Adaptive Learning Rates with Sublinear Memory Cost (ICML 2018), Algorithms 4 to 6.
  6. On data versus size: Hoffmann et al., Training Compute-Optimal Large Language Models (2022), explained in the Chinchilla explainer. On emergence and metrics: Schaeffer, Miranda & Koyejo, Are Emergent Abilities of Large Language Models a Mirage? (NeurIPS 2023).