Building a language model from scratch: tokenizer to trained checkpoint

Why this subject matters now

Five years ago the practical knowledge required to work on language models was mostly architectural, knowing the transformer block, knowing attention, knowing how to fine-tune a pretrained encoder. That knowledge is now table stakes and is also the smallest part of the job. What changed is that the field discovered, empirically and then theoretically, that the architecture is close to a solved commodity and almost everything that separates a good model from a mediocre one at fixed compute lives elsewhere, in how the text is segmented into tokens, in how the data is filtered and mixed, in the ratio of parameters to training tokens, in the numerical format the matrix multiplies run in, in how the work is sharded across hundreds or thousands of accelerators, and in the inference arithmetic that decides whether the trained model can be served at a price anyone will pay.

The specific things a practitioner is expected to know today and was not expected to know in 2020 form a fairly precise list. Compute-optimal scaling, following Hoffmann and colleagues at DeepMind in 2022, replaced the earlier belief from Kaplan and colleagues at OpenAI that parameters should grow far faster than data. The correction moved the recommended token count for a given model size up by roughly an order of magnitude, and then inference economics pushed it up by another order of magnitude past that. Rotary position embeddings, from Su and colleagues, displaced learned absolute positions almost universally because they extrapolate and because the relative-position property is provable rather than hoped for. Grouped-query attention, from Ainslie and colleagues at Google, is now standard because the KV cache, not the parameters, is what limits serving throughput. RMSNorm displaced LayerNorm, SwiGLU displaced the ReLU MLP, bfloat16 displaced float16, and the pre-norm residual stream displaced post-norm, each for a reason that can be derived in a paragraph and each of which this page derives.

The other change is that the whole pipeline became legible. In 2020 the details of a frontier training run were closely held. Today the Llama 3 report from Meta, the DeepSeek-V3 report, the OLMo releases from Allen AI, and the FineWeb and Dolma data pipelines publish enough detail that a careful reader can reconstruct the engineering. That legibility is what makes a page like this possible. The derivations below are checked against numbers those groups published, and the parts that can be measured on a single GPU are measured here rather than asserted.

This page assumes the transformer block itself. Attention as a differentiable dictionary lookup, the \( \sqrt{d_k} \) scaling argument, causal masking, and multi-head decomposition are derived at attention and the transformer block and are summarized here only where the extensions need them. Two shorter companions cover the same territory at lower resolution, where transformers from scratch is the minimal implementation walk-through and training an LLM is the practitioner's checklist. This page is the reference the two of them point at.

Tokenization

Why bytes, and what the alternatives cost

A language model is a distribution over sequences of discrete symbols. The choice of symbol alphabet is not a preprocessing detail. It fixes the model's vocabulary matrix, its effective context length in characters, its arithmetic behaviour, and its multilingual cost structure. Three alphabets are available in principle. Words are the historical choice and fail immediately, since any fixed word list has an out-of-vocabulary problem, the tail of natural language is infinite, and the embedding matrix for a realistic word list dwarfs the rest of the model. Characters, or better Unicode code points, have no out-of-vocabulary problem but there are 154,998 assigned code points in Unicode 16, so a code point vocabulary is both large and badly distributed, and sequences become four to six times longer than word sequences, which costs quadratically in attention.

Bytes are the third choice and the one every modern model builds on. UTF-8 encodes every Unicode code point as one to four bytes, the alphabet has exactly 256 symbols, every possible input maps to a valid sequence, and no text can ever be unrepresentable. The price is sequence length. English prose is about one byte per character, so a byte-level model at a fixed context of \( T \) tokens sees \( T \) characters, whereas a subword model sees three to four times that. Since attention cost grows as \( T^2 \) and the MLP cost grows as \( T \), a factor of four in sequence length is a factor of four to sixteen in compute for the same text. Pure byte-level models are an active research line (the byte-level work from Meta on latent patching, and the earlier ByT5 from Google) but every production model today uses a learned subword vocabulary built on top of bytes, which is what byte-level BPE is.

Byte-pair encoding, derived

Byte-pair encoding was introduced as a data compression algorithm by Gage in 1994 and adapted to neural machine translation by Sennrich, Haddow, and Birch at Edinburgh in 2016. The idea is a greedy bottom-up construction. Start with the alphabet \( \Sigma_0 = \{0, \dots, 255\} \), the raw bytes. Represent the corpus as a sequence of symbols from \( \Sigma_0 \). Then, repeatedly, count every adjacent symbol pair in the corpus, take the most frequent pair \( (a,b) \), mint a new symbol \( c = ab \), add it to the alphabet, and replace every occurrence of the pair by the new symbol. After \( m \) merges the alphabet has \( 256 + m \) symbols and the corpus has become shorter by exactly the number of replacements performed.

The greedy rule is worth taking seriously as an optimization problem. Let \( n_t \) be the length of the encoded corpus after \( t \) merges, and let \( f_t(a,b) \) be the count of pair \( (a,b) \) at step \( t \). Replacing every occurrence of a pair with count \( f \) reduces the corpus length by exactly \( f \) symbols, so

Each merge therefore buys the largest immediately available reduction in sequence length, which is exactly the greedy step for the objective "minimize corpus length after \( m \) merges". It is greedy and not optimal, because merging a pair now changes the counts of overlapping pairs, so the sequence of locally best merges is not in general the globally best set of \( m \) merges. Finding the optimal merge set is combinatorial and nobody does it. The greedy construction is close enough that the field has never moved off it. The connection to compression is direct. BPE is a dictionary coder, and the quantity it maximizes, corpus length reduction per symbol added to the dictionary, is exactly what a dictionary coder trades.

Encoding a new string with a trained BPE tokenizer is not the same algorithm. Training produces an ordered list of merges, and encoding applies them in the order they were learned. Given a byte string, repeatedly find the pair present in the string whose merge rank is lowest (learned earliest) and apply it, until no learned pair remains. This is deterministic and reproduces the training-time segmentation on text the tokenizer has seen, and generalizes sensibly to text it has not. The cost is \( O(\ell^2) \) per pre-token of length \( \ell \) in the naive implementation and \( O(\ell \log \ell) \) with a priority queue. Since pre-tokens are short, either is fine, and the practical speed comes from caching the encoding of each distinct pre-token.

Regex pre-tokenization and why it exists

Running BPE directly on raw text produces tokens that straddle word boundaries and punctuation, so dog. and dog! and dog, all become distinct tokens, and the model has to learn separately that each of them means the same animal. Worse, merges will happily cross whitespace and produce tokens containing several words, which fragments the statistics of every constituent word. GPT-2 fixed this by first splitting text with a regular expression, then running BPE independently inside each piece and never allowing a merge to cross a piece boundary. The GPT-2 pattern splits on contractions, then runs of letters optionally preceded by one space, then runs of digits, then runs of punctuation, then whitespace. The optional leading space is the elegant part. It makes _dog (space-dog) a single token distinct from dog at the start of a line, which is exactly the distinction English orthography makes, and it removes the need for a separate word-boundary marker.

The cl100k pattern used by GPT-4 changes three things. Contractions match case-insensitively, so 'LL and 'll tokenize alike. Digit runs are capped at three characters, which prevents the tokenizer from minting arbitrary multi-digit tokens and gives arithmetic a consistent (if still awkward) segmentation. And newlines are grouped with trailing whitespace rather than arbitrarily, which matters for code. The measured effect of the pattern is visible in this repository. Training a vocabulary of 8,192 on 4MB of English prose and measuring compression on held-out text gives 3.46 bytes per token with the GPT-2 pattern and 3.69 bytes per token with the GPT-4 pattern, and the digit cap changes the segmentation of 2024 from [" 2","0","2","4"] to [" ","20","2","4"]. The second is a better compressor here and neither is a good number representation, which is the point of the next subsection.

Pre-tokenization also has a systems justification. It makes the trainer parallelizable, since the corpus becomes a bag of (pre-token, count) pairs, and every merge operates on that bag rather than on the raw stream. The implementation used for the measurements below keeps an inverted index from pair to the set of pre-tokens containing it, so each merge touches only the affected entries. That is what lets a pure-Python trainer build a 32,768 token vocabulary from 8MB of text in 176 seconds. The naive rescan-everything implementation is roughly two orders of magnitude slower.

Byte fallback and the guarantee it provides

A byte-level BPE tokenizer built the way above has a property worth stating precisely. Every byte string round-trips. The initial alphabet contains all 256 byte values, merges only ever create symbols that are concatenations of existing symbols, and encoding never invents a symbol, so the decoded output of an encoded input is the input, byte for byte. There is no unknown token, no normalization loss, no unrepresentable input. In this repository that was checked directly. A string of 128 Unicode code points from U+0100 to U+017F, none of which appear in the English training corpus, encodes to a sequence that decodes back exactly, at a measured cost of 256 bytes becoming 256 tokens, that is, one token per byte with no compression at all.

SentencePiece-based tokenizers that operate on code points rather than bytes need an explicit byte-fallback mechanism to get the same guarantee. The vocabulary reserves 256 byte tokens, and any code point that is not in the vocabulary is emitted as its UTF-8 bytes. Llama's tokenizer does exactly this. The two designs end at the same place, but byte-level BPE gets there without a special case.

Vocabulary size, the parameter and sequence-length tradeoff

Vocabulary size \( V \) trades two costs against each other. It enters the parameter count linearly through the embedding matrix and, if untied, the output projection, at \( V d \) parameters each. It enters the sequence length inversely, since a larger vocabulary compresses text into fewer tokens, so a fixed amount of text costs fewer positions, which reduces both attention and MLP work per document. Write \( \beta(V) \) for the measured bytes per token at vocabulary \( V \). Then processing a corpus of \( B \) bytes requires \( B / \beta(V) \) token positions, and the training cost in FLOPs, using the \( C \approx 6ND \) law derived later, is

where \( \kappa \in \{1,2\} \) depending on whether the output embedding is tied to the input. The two factors pull in opposite directions. Raising \( V \) raises the first factor linearly and lowers the second sublinearly, because \( \beta \) grows roughly logarithmically in \( V \). Differentiating and setting to zero gives the optimality condition

that is, the relative parameter cost of one more vocabulary entry should equal the relative compression it buys. This is exactly the calculation Tao and colleagues formalized in 2024 when they argued that vocabulary size should scale with model size, and it explains the empirical drift from GPT-2's 50,257 to Llama 3's 128,256 to Qwen and DeepSeek's roughly 150,000. As \( N_{\text{body}} \) grows, the left side shrinks and a larger \( V \) becomes optimal.

The compression curve was measured here rather than assumed. A byte-level BPE tokenizer was trained on 8.1MB of English prose at seven vocabulary sizes and evaluated on held-out text from several domains. The table gives bytes per token, where higher means better compression.

Three things are visible. First, compression grows roughly logarithmically. Each doubling of \( V \) from 4,096 upward buys about 0.12 to 0.20 additional bytes per token, a shrinking return, while the embedding parameter cost doubles. Fitting \( \beta(V) = a + b\ln V \) over the top four points gives \( b \approx 0.27 \) bytes per token per nat of vocabulary, so \( \beta'/\beta \approx 0.27/(3.4 V) \). Setting that equal to \( \kappa d/(N_{\text{body}} + \kappa V d) \) for the tied model of the next section (\( d = 512 \), \( N_{\text{body}} = 22.6 \)M, \( \kappa = 1 \)) gives \( V^\star \approx 12{,}000 \), which is why 8,192 was a defensible choice for a 27M-parameter model and why 128,000 is a defensible choice for an 8B one. Second, code compresses much worse than prose at every vocabulary size, because identifiers are compounds and indentation is repetitive whitespace. Third, and most consequentially, an English-trained tokenizer gives Chinese 0.41 characters per token, meaning a Chinese character costs about 2.4 tokens, so the same semantic content costs roughly nine times more context than English. That is the multilingual tax, measured, and it is why every serious multilingual model trains the tokenizer on a language-balanced mixture rather than on the pretraining mixture.

Tokenizer pathologies

The 32,768-entry tokenizer trained above was probed on inputs that break tokenizers in characteristic ways. Every result below is that tokenizer's actual output.

Numbers. The string 127 128 1000 1024 65536 3.14159 1,234,567 is 41 characters and becomes 22 tokens, segmented as ["12","7"," 12","8"," 1000"," 10","24"," 6","55","36",...]. Note what happened. 127 split as 12|7 and 128 as 12|8, but 1000 survived whole and 1024 became 10|24. The segmentation of a number depends on the frequency of its digit substrings in the training corpus, so numerically adjacent values receive completely unrelated token sequences, and the model must learn arithmetic over a representation in which 127 and 128 share a prefix token while 1023 and 1024 do not. This is the mechanical reason language models are bad at multi-digit arithmetic without either digit-level tokenization, right-to-left digit grouping (as in Llama 3, which splits numbers into individual digits), or chain-of-thought that reduces every step to single digits.

Code indentation. The eight-space indent in a Python snippet tokenized as "\n " followed by " ", a four-space token plus a three-space token, and then the following space is absorbed into the next word token. A twelve-space indent split as a sixteen-space token would not fit, so it becomes "\n " plus " ". Indentation depth is therefore represented inconsistently, which is exactly why code-focused tokenizers reserve dedicated tokens for runs of 2, 4, 8, 12, 16, and more spaces, and why the GPT-4 pattern groups whitespace runs explicitly.

Non-Latin scripts. Thirteen Chinese characters became 39 tokens, exactly one token per UTF-8 byte, because no merge in an English-trained vocabulary ever covers a CJK byte pair. A Russian phrase of 23 characters became 45 tokens. Each of these tokens is an invalid UTF-8 fragment on its own, so the model is forced to learn the byte-level structure of the encoding as part of learning the language. The output still decodes correctly, which is the byte-fallback guarantee doing its job, but the cost is real.

Undertrained and glitch tokens. Of the 32,512 merged tokens, 17,856 appear at most once in a 2MB sample of the very corpus the tokenizer was trained on. The examples are revealing, among them "ght", "rom", "ough", "ther". These are intermediate merge products. BPE created them on the way to "thought" and "from" and then almost never emits them, because the encoder always prefers the longer merge. Their embedding rows receive gradient a handful of times in the entire training run and remain essentially at initialization. This is the mechanism behind glitch tokens. A token that exists in the vocabulary but is nearly absent from the training distribution has an untrained embedding, and prompting the model with it produces arbitrary behaviour. The widely discussed SolidGoldMagikarp family in GPT-2 and GPT-3 arose the same way, from tokenizer training data that contained material the model's training data did not. The defence is to train tokenizer and model on the same distribution and to audit the vocabulary for rows whose embedding norm never moves.

Vocabulary utilization. On 400KB of held-out English, only 8,712 of the 32,768 entries were used at all, and the 1,000 most frequent tokens covered 83.35% of all positions. The distribution over tokens is close to Zipfian by construction, which is why the softmax over a large vocabulary is cheap in practice despite being expensive in principle. The same few thousand logits dominate.

Unigram language-model tokenization and its EM training

BPE is not the only subword algorithm. Kudo's unigram language model, the default in SentencePiece and the tokenizer behind T5, ALBERT, and XLNet, is probabilistic rather than constructive. Fix a vocabulary \( \mathcal{V} \) of pieces with probabilities \( p(x) \) summing to one, and model a segmentation \( \mathbf{x} = (x_1,\dots,x_k) \) of a word \( w \) as

where \( S(w) \) is the set of all segmentations of \( w \) into vocabulary pieces. Training maximizes \( \sum_w f_w \log P(w) \) over \( p \), with \( f_w \) the corpus frequency of \( w \). The segmentation is a latent variable, so this is an EM problem. The E step computes, for each word, the posterior expected number of times each piece is used,

which is a forward-backward computation over the segmentation lattice. Here \( \alpha_i \) is the log-sum-exp over prefixes ending at position \( i \), \( \beta_i \) the same over suffixes starting at \( i \), and the expected count of the piece spanning \( [i,j) \) is \( \exp(\alpha_i + \log p(w_{i:j}) + \beta_j - \alpha_n) \). The M step is the closed-form maximizer of the expected complete-data log-likelihood, \( p(x) \leftarrow \E[c(x)] / \sum_{x'} \E[c(x')] \). Because EM cannot change the vocabulary, Kudo wraps it in a pruning loop. Start from a large seed vocabulary of frequent substrings, run EM to convergence, estimate for each piece the loss in total log-likelihood if it were removed (approximated by re-running Viterbi without it), drop the least useful fraction, and repeat until the target size is reached. Single characters are never dropped, which is what preserves coverage.

The two algorithms produce different segmentations of the same text and roughly comparable compression. Implemented here from the description above and trained on 2MB of English prose, the unigram model reaches the compression shown below on held-out text alongside BPE trained on the same data. Unigram is the slower trainer by a wide margin, because every EM iteration touches the whole lattice for every distinct word and every pruning round re-runs Viterbi once per candidate piece.

Architecture, component by component

Embeddings, weight tying, and the parameter arithmetic

The input embedding is a lookup table \( E \in \R^{V \times d} \), and token \( i \) becomes row \( E_i \). The output head is a linear map \( U \in \R^{d \times V} \) producing logits \( z = U\T h \). Weight tying, introduced independently by Press and Wolf at Bar-Ilan and by Inan, Khosravi, and Socher in 2016, sets \( U = E\T \). The argument is that both matrices learn a map between the token identity space and the representation space, in opposite directions, and that forcing them to agree both halves the parameters and regularizes. An input embedding gets gradient from every position where the token appears as input, and the same row gets gradient from every position where it appears as a target, so tying pools two sparse signals into one.

The parameter accounting is worth writing out because it is the first thing to compute for any new configuration. Let \( d \) be the model width, \( L \) the number of layers, \( H \) the number of query heads, \( H_{kv} \) the number of key/value heads, \( d_h = d/H \) the head dimension, \( F \) the MLP hidden width, and \( V \) the vocabulary. The per-layer count is

and the whole model is that times \( L \), plus \( Vd \) for the embedding, plus \( d \) for the final norm, plus another \( Vd \) if the head is untied. For the configuration used throughout this page, \( V = 8192 \), \( d = 512 \), \( L = 8 \), \( H = 8 \), \( H_{kv} = 2 \), \( d_h = 64 \), \( F = 1408 \), tied, the counts are given below.

The closed form and the constructed module agree exactly, \( 26{,}747{,}392 \) from both, verified in this repository. Non-embedding parameters are \( 22{,}553{,}088 \). Two structural facts generalize. The MLP dominates, taking roughly two thirds of the parameters in any modern configuration, which is why mixture-of-experts targets the MLP and not attention. And the embedding share shrinks as models grow. At \( d = 512 \) it is 15.7%, at Llama 3 8B's \( d = 4096 \) with \( V = 128{,}256 \) the two embedding matrices (untied there) are 1.05B of 8.03B, or 13%, and at 70B they are 3.3%.

Attention, and what changes at scale

Scaled dot-product attention, the \( 1/\sqrt{d_k} \) scaling argument, causal masking, and the multi-head decomposition are derived in full at attention and the transformer block and are taken as given here. The one-line summary is that each position emits a query, a key, and a value, that output \( i \) is \( \sum_j \softmax_j(q_i \cdot k_j/\sqrt{d_h}) v_j \) restricted to \( j \le i \), and that heads run the same computation in \( H \) disjoint \( d_h \)-dimensional subspaces and are concatenated. What this page adds is everything that changes when the model has to be trained on many devices and served to many users at once, and the first of those is the key-value cache.

Multi-query and grouped-query attention, and the KV-cache arithmetic

During autoregressive decoding, generating token \( t+1 \) requires attending over the keys and values of all previous positions. Recomputing them is \( O(t) \) work per step and \( O(T^2) \) for a full generation, while caching them is \( O(1) \) work per step at the cost of storing them. The cache holds, for every layer, every key/value head, and every position, one \( d_h \)-vector of each, for every sequence in the batch. In bytes, with 2 bytes per element, this is

where the leading 2 counts keys and values and the second counts bytes per bf16 element. This is the number that dominates serving. Take Llama 2 70B, which uses \( L = 80 \), \( d = 8192 \), \( H = 64 \), \( d_h = 128 \). With multi-head attention (\( H_{kv} = H = 64 \)) at \( T = 4096 \) and \( B = 1 \), the cache is

for a single sequence. The weights themselves are 140GB in bf16, so eight H100s hold them with about 500GB free. At 12.8GiB per sequence, that is roughly 39 concurrent sequences, a poor serving economy. Llama 2 70B actually uses grouped-query attention with \( H_{kv} = 8 \), an eight-fold reduction to 1.6GiB per sequence and roughly 310 concurrent sequences from the same memory. The KV cache, not the parameter count, is what sets serving concurrency.

Multi-query attention, proposed by Shazeer at Google in 2019, takes this to the limit, \( H_{kv} = 1 \). All \( H \) query heads share one key head and one value head. The cache shrinks by a factor of \( H \), and the KV projections shrink from \( 2d^2 \) to \( 2d\,d_h \) parameters per layer. The cost is quality. Sharing one key subspace across all heads removes the ability of different heads to look for different things, and Shazeer reported a small but consistent degradation. Grouped-query attention, from Ainslie and colleagues in 2023, interpolates. Partition the \( H \) query heads into \( G \) groups and give each group its own key and value head, so \( H_{kv} = G \). At \( G = H \) it is multi-head, at \( G = 1 \) it is multi-query, and the empirical finding, which Llama 2, Llama 3, Mistral, and Qwen all adopted, is that \( G = 8 \) recovers essentially all of multi-head's quality at close to multi-query's cache size. The paper also showed that an existing multi-head checkpoint can be converted by mean-pooling the key and value projections within each group and then uptraining on about 5% of the original tokens, which is why the technique spread so quickly.

Measured here on the H100 at the 26.7M-parameter scale, the four variants differ in cache bytes exactly as the formula predicts. Per token per sequence, with \( L = 8 \) and \( d_h = 64 \), the cache costs \( 4 L H_{kv} d_h \) bytes, which is 16,384 bytes for \( H_{kv} = 8 \), 4,096 for \( H_{kv} = 2 \), and 2,048 for \( H_{kv} = 1 \). Training throughput barely moves, because during training there is no cache and the KV projections are a small share of the FLOPs. The win is entirely at inference.

Sliding-window attention and attention sinks

A second way to bound the cache is to bound what each position may attend to. Sliding-window attention, used in Mistral 7B and in Longformer before it, restricts position \( i \) to \( [i - w, i] \). The cache then never exceeds \( w \) positions per layer, so memory becomes constant in sequence length rather than linear, and attention cost becomes \( O(Tw) \) rather than \( O(T^2) \). Information still propagates further than \( w \), since after \( L \) layers the receptive field is \( Lw \), by the same stacking argument that gives convolutions their receptive field. Mistral 7B uses \( w = 4096 \) with \( L = 32 \), a theoretical reach of 131,072 tokens. In practice the effective reach is shorter, and the current fashion, in Gemma 2 and Llama 4 and others, is to interleave, with most layers sliding-window and every fourth or fifth layer full attention, which bounds the cache while keeping a few layers that can genuinely see everything.

Attention sinks are a related and initially surprising finding from Xiao and colleagues at MIT with Meta and CMU in 2023. If you take a model trained with full attention and simply evict the oldest entries from its cache to make a sliding window, quality collapses. The reason is that trained transformers dump a large fraction of attention mass onto the first few tokens of the sequence, regardless of content, because the softmax must sum to one and a head that wants to attend to nothing needs somewhere to put its mass. The first tokens, visible to every position, become that somewhere. Evicting them forces the mass onto genuinely irrelevant tokens and destroys the representation. The fix is to keep the first four or so tokens permanently in the cache alongside the sliding window, which restores the original behaviour and enables effectively unbounded streaming. The deeper fix, adopted in several recent models, is to give the softmax an explicit escape, a learned logit appended to the denominator, so \( \softmax \) becomes \( e^{z_i}/(e^{s} + \sum_j e^{z_j}) \), letting a head attend to nothing at all without hijacking a position.

Positional information, four designs and one proof

Attention is permutation-equivariant. Nothing in \( \softmax(QK\T/\sqrt{d_h})V \) depends on the order of the rows, so without added positional information a transformer sees a bag of tokens. That claim is easy to check and was checked here. Training the reference model with the positional signal removed entirely gives a validation loss of 4.677 after 65.5M tokens, against 3.566 for the same model with rotary embeddings, a gap of 1.11 nats. The model is not useless without positions, because the causal mask itself leaks a weak ordering signal (position \( i \) attends to \( i \) keys, so the softmax normalization differs by position), which is the mechanism behind "NoPE" results showing decoder-only models can partly infer order. But the gap is large.

Learned absolute embeddings add a trained vector \( p_t \) to the token embedding at position \( t \), as in BERT and GPT-2. This costs \( T_{\max} d \) parameters, and its fatal property is that \( p_t \) is undefined for \( t \ge T_{\max} \), so the model cannot process a longer sequence than it was trained on, at all. Measured here it also underperforms rotary at equal budget, 4.114 against 3.566.

Sinusoidal embeddings, from the original transformer paper, define \( p_{t,2i} = \sin(t/10000^{2i/d}) \), \( p_{t,2i+1} = \cos(t/10000^{2i/d}) \). These have no parameters and are defined for all \( t \). Vaswani and colleagues motivated them by noting that \( p_{t+k} \) is a fixed linear function of \( p_t \) for any fixed offset \( k \), so relative offsets are in principle linearly decodable. In practice sinusoidal absolute embeddings extrapolate poorly, because the model learns to use absolute magnitudes and those go out of distribution.

ALiBi, from Press, Smith, and Lewis at Washington and Meta in 2021, throws away position embeddings entirely and adds a linear penalty to the attention scores,

with a head-specific slope \( m_h \) fixed to a geometric sequence \( 2^{-8h/H} \). Every head becomes a soft recency bias with a different decay length. Heads with large \( m_h \) see only the recent past, and heads with small \( m_h \) see far. The virtue is extrapolation. A model trained at 1,024 tokens evaluated at 2,048 degrades gracefully, because the bias is a smooth function of distance with no learned parameters to go out of range. Measured here, ALiBi reached validation loss 3.451 against rotary's 3.566 at a matched 65.5M-token budget and a 512-token context, but ran 20% slower because the additive bias forces the attention kernel off the fastest fused path. At short context ALiBi is competitive. The reason the field moved to rotary anyway is that ALiBi's recency bias is a hard prior that hurts tasks requiring long-range exact retrieval, and rotary composes with the context-extension tricks described below.

Rotary position embeddings, derived

RoPE, from Su and colleagues in 2021, asks for a function \( f(x, m) \) applied to the query at position \( m \) and the key at position \( n \) such that the resulting inner product depends on \( m \) and \( n \) only through \( m - n \),

Work in two dimensions first. Identify \( \R^2 \) with \( \C \) by \( x = (x_1, x_2) \mapsto x_1 + i x_2 \), and let \( f(x, m) = x e^{i m\theta} \), a rotation by angle \( m\theta \). The real inner product of two complex numbers is \( \langle u, v \rangle = \mathrm{Re}(u \bar{v}) \), so

which depends on \( m \) and \( n \) only through \( m-n \). That is the whole proof in two dimensions. Concretely, writing the rotation as a matrix,

using \( R_m\T = R_{-m} \) and \( R_a R_b = R_{a+b} \), which hold because the 2-D rotations form a commutative group isomorphic to the circle. To extend to \( d_h \) dimensions, split the head vector into \( d_h/2 \) disjoint coordinate pairs and rotate pair \( i \) by \( m\theta_i \) with \( \theta_i = \Theta^{-2i/d_h} \) and base \( \Theta = 10^4 \). The block-diagonal matrix \( R_m = \bigoplus_i R_m^{(\theta_i)} \) is orthogonal, satisfies the same two group identities blockwise, and therefore

a sum of terms each depending only on \( n - m \). The dot product is a function of relative position alone, exactly, with no approximation. The geometric interpretation is that the frequency ladder \( \theta_i \) makes low-index pairs rotate fast (short wavelength, fine positional discrimination) and high-index pairs rotate slowly (long wavelength, coarse but long-range discrimination), so a single head carries positional information at many scales at once.

This was verified numerically here rather than trusted. Taking a random query and key in \( \R^{64} \), rotating them to positions \( (m,n) \in \{(5,2), (17,14), (100,97), (203,200)\} \), all with \( m - n = 3 \), gives inner products \( -0.434219552, -0.434220165, -0.434215265, -0.434205623 \), identical to eight significant figures, with a maximum deviation of \( 1.4\times10^{-8} \) attributable to float32 evaluation of the trigonometric tables. Changing the offset to \( m-n = 7 \) gives \( -4.269 \), a completely different value, as it must.

Two implementation details matter. RoPE is applied to queries and keys but never to values, because the goal is to modulate the scores, and rotating values would rotate the output in a position-dependent way. And the rotation is applied after the projections and after any QK-normalization, per head, which makes it a cheap elementwise operation with no parameters. The cost is \( O(T d) \) per layer against \( O(T^2 d) \) for the attention it modifies, so it is free.

Extending context with position interpolation, NTK-aware scaling, and YaRN

A model trained with RoPE at context \( T \) has seen rotation angles \( m\theta_i \) for \( m < T \) only. Evaluated at \( m \ge T \), the low-frequency pairs enter angles never seen in training and the model breaks, typically catastrophically. Three fixes, in historical order.

Position interpolation, from Chen and colleagues at Meta in 2023, rescales positions instead of extrapolating them. To extend from \( T \) to \( T' = sT \), replace \( m \) by \( m/s \). Every angle stays in the trained range, and the model sees a compressed version of the same geometry. This works with a small amount of fine-tuning (Meta reported 1,000 steps) but it degrades short-range resolution. Neighbouring tokens that used to be \( \theta_i \) apart in angle are now \( \theta_i/s \) apart, so the fine positional discrimination the high-frequency pairs provided is compressed by the same factor.

NTK-aware scaling, developed in public by the open-source community (the "NTK-aware scaled RoPE" posts, later formalized) makes the observation that the two ends of the frequency ladder should be treated differently. High-frequency pairs complete many full rotations within the trained context, so they have effectively seen all angles and need no adjustment. Only low-frequency pairs, whose wavelength exceeds the trained context, are extrapolating. Instead of scaling positions, scale the base, \( \Theta' = \Theta\, s^{d_h/(d_h-2)} \). This leaves the fastest pair almost unchanged and stretches the slowest pair by approximately \( s \), interpolating where interpolation is needed and extrapolating where extrapolation is safe. It works without any fine-tuning at moderate \( s \), which made it widely adopted.

YaRN, from Peng, Quesnelle, Fan, and Shippole in 2023, makes the frequency-dependent treatment explicit. Define the wavelength of pair \( i \) as \( \lambda_i = 2\pi/\theta_i \) and the ratio \( r_i = T/\lambda_i \), the number of full rotations that pair completes in the trained context. YaRN interpolates fully (\( m \to m/s \)) for pairs with \( r_i < \alpha \), leaves pairs with \( r_i > \beta \) untouched, and ramps linearly between, with \( \alpha = 1 \) and \( \beta = 32 \) in the paper. It adds one more piece. Because interpolation concentrates the attention distribution, YaRN multiplies the attention logits by a temperature \( t = 0.1\ln s + 1 \), a correction fitted empirically that recovers the pre-extension entropy. YaRN reaches a given context length with roughly ten times less fine-tuning data than position interpolation. Llama 3.1's 128K context uses a closely related frequency-dependent scheme, and Qwen has shipped YaRN configurations directly.

All three share one honest limitation. Extending the positional encoding extends where the model can look, not what it can use. Retrieval benchmarks and the needle-in-a-haystack family repeatedly show a gap between advertised and effective context, and closing it needs long-context training data, not just a rescaled \( \Theta \).

Normalization

LayerNorm, from Ba, Kiros, and Hinton at Toronto in 2016, computes for each token vector \( x \in \R^d \)

RMSNorm, from Zhang and Sennrich at Edinburgh in 2019, drops the mean, computing

The argument for dropping the mean is that re-centering is not what makes normalization work. LayerNorm's benefit is scale invariance. The output is invariant to \( x \to cx \), which bounds the activation magnitude entering the next layer and, crucially, makes the gradient with respect to the incoming weight matrix inversely proportional to that weight's scale, an implicit learning-rate adaptation. Re-centering adds shift invariance, \( x \to x + c\1 \), which the residual stream does not actually need, because the following linear layer can absorb any constant shift into its bias or into the norm's own \( \gamma \). Zhang and Sennrich showed empirically that removing the mean costs nothing in quality across a range of tasks while removing two of the four reduction passes.

The speed argument needs care, and measuring it here produced a more interesting answer than the usual claim. On the H100, normalizing a \( 4096 \times 4096 \) tensor, PyTorch's fused LayerNorm kernel takes 0.0648 ms, the same LayerNorm written as eager elementwise operations takes 0.4027 ms, an eager RMSNorm takes 0.1968 ms, and compiling either with torch.compile gives 0.0732 ms for LayerNorm and 0.0651 ms for RMSNorm. So the honest statement is that RMSNorm is about 11% faster than LayerNorm when both are fused, and that the factor-of-two difference people quote is really the difference between fused and unfused code, not between the two formulas. At a bandwidth of roughly 2.93 TB/s measured on this device, a 4096×4096 bf16 tensor read and written once is \( 2 \times 4096^2 \times 4 = 134 \) MB at fp32, or 0.046 ms at peak, so the fused kernels are running at 70% of memory bandwidth and the remaining difference between the two is the extra pass the mean requires.

The parameter saving is real but tiny. RMSNorm has \( d \) parameters per norm against LayerNorm's \( 2d \). In the reference model that is 8,704 against 17,408, or 0.03% of the total. Nobody adopts RMSNorm for the parameters. In this repository's matched ablation at 65.5M tokens, LayerNorm reached validation loss 3.551 and RMSNorm 3.566, a difference of 0.015 nats that is within the seed-to-seed noise measured below. The correct summary is that they are equivalent in quality and RMSNorm is marginally cheaper, which is why every model since Llama uses it.

Pre-LN versus post-LN, and the gradient-scale argument

The original transformer put the normalization after the residual addition, \( x_{\ell+1} = \mathrm{LN}(x_\ell + F(x_\ell)) \). Modern models put it before the sublayer, \( x_{\ell+1} = x_\ell + F(\mathrm{LN}(x_\ell)) \). The difference decides whether a deep model trains at all, and the reason is a gradient-scale argument made precise by Xiong and colleagues at Microsoft Research Asia and Peking in 2020.

In the pre-LN arrangement the residual stream is an identity path from input to output. Differentiating \( x_{L} = x_0 + \sum_{\ell=0}^{L-1} F_\ell(\mathrm{LN}(x_\ell)) \) gives

so there is always a path of Jacobian exactly \( I \) from the loss to any layer, and the gradient norm at layer \( \ell \) is bounded independently of \( L \) up to the contributions of the branches. Xiong and colleagues showed that in the pre-LN case the expected gradient norm at initialization scales as \( \Theta(1/\sqrt{\ell}) \) at layer \( \ell \) and is \( \Theta(\sqrt{\ln L / L}) \) at the output layer, that is, gradients are well-scaled and mildly decreasing with depth.

In the post-LN case the normalization sits on the residual path. Each layer's Jacobian carries a factor from the LayerNorm, whose derivative scales like \( 1/\sigma_\ell \), and these multiply through depth. The same analysis gives an expected gradient norm at the output layer of \( \Theta(\sqrt{L}) \) at initialization, so the gradient grows with depth, and a learning rate that is stable for a 6-layer model blows up a 48-layer one. The practical consequence, and the reason warmup was invented, is that post-LN transformers require a learning-rate warmup to survive the first few hundred steps, while pre-LN transformers can, in principle, be trained with a constant rate from step one.

This is directly observable. In this repository, at a learning rate of \( 3\times10^{-3} \) with 2% warmup and identical everything else, the pre-LN model reached validation loss 3.566 and the post-LN model reached 7.290. It diverged and never recovered, at a learning rate the pre-LN model handles without incident. The cost of pre-LN is a subtlety noted by several groups. Because the residual stream is never normalized, its variance grows with depth, and the last layers see inputs of much larger magnitude than the first. The final norm before the output head fixes the symptom. Variants that try to get both properties, such as the sandwich normalization used in Gemma and the DeepNorm scheme from Microsoft that rescales the residual branch by a depth-dependent constant, exist and are used in specific models, but plain pre-LN with a final norm is the default everywhere.

QK-norm

A failure mode that appears only at scale is that attention logits \( q\cdot k/\sqrt{d_h} \) can grow without bound during training, because nothing constrains the norms of \( q \) and \( k \). When they do, the softmax saturates, its Jacobian \( \diag(a) - aa\T \) goes to zero, the attention entropy collapses to near zero, and the layer stops learning. Dehghani and colleagues at Google Brain hit this at 22B parameters in ViT-22B and fixed it with QK-norm, applying a normalization to \( q \) and \( k \) per head before the dot product, which bounds \( |q \cdot k| \le \|q\|\|k\| = d_h \) by construction. It is now standard in several recent models including Gemma 2 and OLMo 2 and in Chameleon from Meta.

Measured here at small scale, QK-norm reached validation loss 3.457 against the baseline's 3.566 at a matched budget, at a cost of about 9% throughput (1.204M against 1.319M tokens per second). At 27M parameters the logit-growth pathology does not occur, so this gain is not the stability effect. It is the mild regularization of normalizing the query and key. The honest reading is that QK-norm is cheap insurance whose value grows with scale.

Activations and the MLP

The position-wise MLP is where two thirds of the parameters and roughly two thirds of the FLOPs live. The original design is \( \mathrm{MLP}(x) = W_2\,\mathrm{ReLU}(W_1 x) \) with \( W_1 \in \R^{4d\times d} \), \( W_2 \in \R^{d\times 4d} \), the expansion factor of 4 being an empirical choice that has survived eight years of scrutiny.

GELU, from Hendrycks and Gimpel in 2016, replaces the hard gate with a smooth probabilistic one, \( \mathrm{GELU}(x) = x\,\Phi(x) \), where \( \Phi \) is the standard normal CDF. The motivation is that ReLU multiplies its input by \( \mathbb{1}[x>0] \), a Bernoulli gate with a discontinuous derivative, whereas GELU multiplies by the probability that a standard normal is below \( x \), which is smooth and matches the intuition of stochastic regularization. In practice GELU trains slightly better than ReLU and the tanh approximation \( 0.5x(1+\tanh[\sqrt{2/\pi}(x + 0.044715x^3)]) \) is what almost everyone actually computes.

SwiGLU, from Shazeer's 2020 note on GLU variants, replaces the single nonlinearity with a gated pair,

Two projections go up, one multiplies the other elementwise, and one comes back down. The gating is the point. The network can compute a value \( W_3 x \) and, independently, a data-dependent multiplicative mask on it, which a single-path MLP cannot do without going deeper. Shazeer's own summary of why it works is worth quoting for its honesty. He attributes the improvement to divine benevolence rather than to an explanation, and the field has not produced a better one since.

Because SwiGLU uses three matrices instead of two, a naive swap at hidden width \( 4d \) increases MLP parameters from \( 8d^2 \) to \( 12d^2 \). The standard correction is to shrink the hidden width by two thirds,

which is why Llama uses \( F = \frac{2}{3}\cdot 4d \) rounded up to a multiple of 256, and why the reference model here uses \( F = 1408 \approx 2.75 \times 512 \), rounded to a multiple of 128 so the matmul tiles align. Under this correction, in this repository's matched-budget ablation, a GELU MLP at \( F = 4d \) (26.22M parameters) reached validation loss 3.519 and a SwiGLU MLP at \( F = 2.75d \) (26.75M parameters) reached 3.566, while a SwiGLU at the full \( F = 4d \) (34.6M parameters, and therefore not a matched comparison) reached 3.581. At this scale and this budget the GELU MLP is slightly ahead, which is a useful corrective. The SwiGLU advantage reported in the literature is around 0.01 to 0.02 in bits per byte at scales three orders of magnitude larger, and it is not detectable in a 65M-token run. The reason to use SwiGLU is that it does not hurt and it helps at scale, not that it transforms a small model.

The full parameter formula, and a concrete configuration

Collecting everything, for a pre-LN decoder with RMSNorm, RoPE, GQA, SwiGLU, and tied embeddings, the total is

With \( F = \frac{8}{3}d \) and \( H_{kv}/H = 1/4 \) this simplifies to \( N \approx Vd + L(2.5d^2 + 8d^2) = Vd + 10.5 L d^2 \), and the familiar rule of thumb \( N \approx 12Ld^2 \) for the classic MHA plus \( 4d \) MLP configuration (\( 4d^2 + 8d^2 = 12d^2 \)) becomes \( 10.5Ld^2 \) for a modern GQA plus SwiGLU one. Check the formula against a real model. Llama 3 8B has \( d = 4096 \), \( L = 32 \), \( H = 32 \), \( H_{kv} = 8 \), \( F = 14336 \), \( V = 128256 \), untied. The body is \( 32[4096^2(1 + 0.5) + 4096^2 + 3\cdot4096\cdot14336 + 2\cdot4096] = 32[2.516\times10^7 + 1.678\times10^7 + 1.761\times10^8 + 8192] = 32 \times 2.180\times10^8 = 6.98\times10^9 \). The embeddings are \( 2 \times 128256 \times 4096 = 1.051\times10^9 \). The total is \( 8.03\times10^9 \), which matches the published 8.03B exactly.

Training compute and where the factor of 6 comes from

The single most useful estimate in this field is that training a model with \( N \) parameters on \( D \) tokens costs about \( 6ND \) floating-point operations. It is worth deriving rather than memorizing, because the derivation tells you exactly when it is wrong.

Consider one weight matrix \( W \in \R^{m \times n} \) used as \( y = Wx \) for a single token. Forward, \( y_i = \sum_j W_{ij}x_j \) costs \( mn \) multiplies and \( mn \) adds, so \( 2mn = 2\,|W| \) FLOPs where \( |W| \) is the parameter count. Backward, two gradients are needed. The gradient with respect to the input, \( \partial \L/\partial x = W\T (\partial \L/\partial y) \), is another matrix-vector product, \( 2mn \) FLOPs. The gradient with respect to the weights, \( \partial \L/\partial W = (\partial \L/\partial y)\,x\T \), is an outer product accumulated into an \( m\times n \) array, again \( 2mn \) FLOPs. So backward costs exactly twice forward, and the total is \( 6\,|W| \) FLOPs per parameter per token. Summing over every weight matrix in the model gives \( 6N \) per token, and over \( D \) tokens,

The estimate ignores three things, each of which is worth knowing. First, elementwise operations. Activations, normalizations, and the softmax are \( O(1) \) FLOPs per element rather than per parameter, and contribute a percent or two. Second, the attention score and value matmuls, which involve no parameters at all. Per layer per token, computing \( q\cdot k \) against all previous keys costs \( 2 d T \) FLOPs on average without causal masking and \( dT \) with it, and the same again for the value aggregation. Forward plus backward triples this. The standard convention, from the PaLM report and used by nanoGPT, ignores the causal discount and writes

Third, the estimate assumes every FLOP is a useful matmul FLOP, which is what model FLOPs utilization then measures against.

When the attention term matters. The ratio of the attention term to the parameter term is \( 12LdT / 6N \). With \( N \approx 12Ld^2 \) this is \( 12LdT/(72Ld^2) = T/(6d) \). So attention is negligible when \( T \ll 6d \) and dominant when \( T \gg 6d \). For Llama 3 8B (\( d = 4096 \)) the crossover is \( T = 24{,}576 \). At the 8K training context, attention is a third of the parameter cost, and at 128K it is five times larger. For the reference model here (\( d = 512 \), \( T = 512 \)) the crossover is 3,072, and indeed the measured ratio is \( 12 \times 8 \times 512 \times 512 / (6 \times 2.67\times10^7) = 2.52\times10^7 / 1.60\times10^8 = 0.157 \), so attention is 16% of the total and the plain \( 6ND \) estimate is 14% low. The exact FLOPs per token, computed from the model, are \( 1.857\times10^8 \) under the PaLM convention and \( 1.731\times10^8 \) with the causal discount applied.

From FLOPs to wall clock, and what MFU really measures

Model FLOPs utilization is the ratio of the model's useful FLOPs to the hardware's peak,

The denominator has to be a number you can actually reach, not a marketing figure. The measured peak on this H100 80GB HBM3, from classes/data/h100.json, is 728.7 bf16 TFLOP/s on an \( 8192^3 \) matmul and 744.6 at \( 4096^3 \). Fp16 is essentially identical at 700.8 and 723.4, tf32 reaches 409.7, and plain fp32 tops out at 51.4. The memory bandwidth measured on the same device is 2,930 GB/s for a bf16 copy and 3,063 GB/s for a bf16 add. Every MFU figure on this page uses 729 TFLOP/s as the denominator.

Now the worked example that the rest of the page's engineering exists to support. Take an 8B-parameter model trained on 2 trillion tokens at a context of 8,192, on 1,024 H100s.

At 100% of 729 TFLOP/s on 1,024 devices, that is \( 1.22\times10^{23}/(1024 \times 7.29\times10^{14}) = 1.63\times10^5 \) seconds, or 1.9 days. At a realistic 45% MFU it is 4.2 days, and at 35% it is 5.4 days. Those numbers are the reason MFU is tracked so closely. Every point of MFU is a day of a thousand-GPU cluster.

What real runs achieve. Published figures cluster between 35% and 55% for dense transformers on modern hardware. Megatron-LM's 2021 paper reported 52% of peak on A100s at 1T parameters with 3D parallelism. PaLM reported 46.2% MFU on TPU v4 for the 540B model. The Llama 3 report describes 38% to 43% BF16 MFU on 16,384 H100s for the 405B model. Mixture-of-experts models run lower, typically 25% to 40%, because expert routing introduces all-to-all communication and load imbalance. Anything above 55% on a large cluster should be treated as either an unusually favourable configuration or a different definition of the numerator.

What was achieved here. The 26.7M-parameter reference model, trained on this H100 at batch 64 and context 512 in bf16 with torch.compile and the fused attention kernel, sustained 1,232,919 tokens per second over 16,384 steps, which at \( 1.857\times10^8 \) FLOPs per token is \( 2.29\times10^{14} \) FLOP/s, or 31.4% MFU (29.3% with the causal discount). A shorter dedicated benchmark at the same configuration reached 1,464,648 tokens per second and 37.3% MFU when the machine was otherwise idle. Both numbers are respectable for a model this small. At \( d = 512 \), the matmuls are small enough that kernel launch overhead and the non-matmul operations are a large fraction of the step. The width scan below makes that explicit.

Eight layers, context 512, batch 32, everything else identical. MFU rises from 13.8% to 43.0% purely because the matrices get bigger and the tensor cores get closer to their asymptotic efficiency. This is the same effect visible in the raw matmul benchmark, 108.9 TFLOP/s at \( n = 1024 \) rising to 744.6 at \( n = 4096 \). Small models are inefficient per FLOP, and that inefficiency is a property of the hardware, not of the code.

Numerics

Formats and the range argument for bfloat16

A floating-point format allocates its bits between a sign, an exponent, and a significand. The exponent width sets the dynamic range, and the significand width sets the relative precision. The four formats that matter are given below, with the values reported by the runtime on this machine.

The decisive comparison is float16 against bfloat16, both 16 bits. Float16 spends 10 bits on the significand and 5 on the exponent, while bfloat16 spends 7 and 8. Bfloat16 therefore has exactly the same dynamic range as float32 and three fewer significand bits than float16. The range is what matters for deep learning. Gradients in a large model span many orders of magnitude, and float16's smallest normal value is \( 6.1\times10^{-5} \), so any gradient component below that either becomes a subnormal (losing precision progressively) or flushes to zero. Measured directly here on the reference model, with an unscaled backward pass in float16 autocast, {{FP16_UNDERFLOW}} that occurs in practice.

The cost of bfloat16's three lost significand bits is that a single rounding has relative error up to \( 2^{-8} = 0.4\% \) instead of \( 2^{-11} = 0.05\% \). This matters for accumulation, not for storage. Summing \( n \) values with independent roundings accumulates error like \( \sqrt{n}\,\epsilon \) in the best case and \( n\epsilon \) in the worst. Measured here, summing \( 4{,}194{,}304 \) values in the range \( [0, 0.01] \) shows this directly. {{ACCUM_ERR}} This is why tensor cores take bf16 inputs and accumulate in fp32 internally, and why no framework ever accumulates a reduction in bf16.

Loss scaling, master weights, and the mixed-precision recipe

The mixed-precision recipe, from Micikevicius and colleagues at NVIDIA and Baidu in 2018, has three parts and each solves a specific numerical problem.

Master weights. Keep an fp32 copy of every parameter. The reason is not the forward pass but the update. With a learning rate of \( 10^{-4} \) and a weight of order 1, the update is \( 10^{-4} \) relative, and bfloat16's machine epsilon is \( 7.8\times10^{-3} \). The update is 80 times smaller than one unit in the last place, so \( w \leftarrow w - \eta g \) in bf16 is exactly \( w \). Small updates vanish entirely. Keeping the master copy in fp32, casting to bf16 for the forward and backward, and applying the update to the fp32 copy fixes this. Modern alternatives include stochastic rounding, which makes the expected update correct even when each individual update rounds to zero.

Loss scaling exists only for float16. Multiply the loss by a constant \( S \) before the backward pass. By linearity every gradient is multiplied by \( S \), shifting the whole distribution up into float16's representable range. Then divide the gradients by \( S \) before the optimizer step. This was measured here on the reference model with float16 autocast. {{LOSS_SCALE}} Dynamic loss scaling, which every framework implements, starts \( S \) high, checks for infinities or NaNs after each backward, halves \( S \) and skips the step when it finds them, and doubles \( S \) after a few hundred clean steps. Bfloat16 needs none of this, which is a large part of why the field switched, one fewer moving part in a run that may last a month.

Measured throughput by precision on this H100, for a full training step of the reference model, is given below.

Reading this table is most of what precision engineering is. Enabling TF32 alone doubles fp32 throughput for free, because the tensor cores accept a 19-bit format transparently. Moving to bf16 adds another 39%. Compiling adds a further 2.6×, which is the largest single factor in the table and comes from operator fusion eliminating memory traffic around the norms, the activation, and the residual adds. The same effect was measured independently in this repository as a 4.26× speedup on a fused elementwise chain. Float16 is slower than bf16 here despite identical tensor-core throughput, because the gradient scaler adds an unscale-and-inspect pass every step. And replacing FlashAttention with the naive materialized-score-matrix implementation costs 4.2× and 2.7× the memory, consistent with the standalone measurements in h100.json, where flash attention at sequence length 2048 is 17.8× faster than the naive version and uses 13.7× less memory.

fp8 and per-tensor scaling

Hopper adds fp8 tensor cores with two formats. The first, e4m3, with more significand, is used for weights and activations, and the second, e5m2, with more exponent, is used for gradients where range matters more. The dynamic range of e4m3 is only \( [1.6\times10^{-2}, 448] \), far too narrow for raw tensors, so fp8 training requires explicit scaling. Before casting a tensor to fp8, divide it by a scale factor chosen so its maximum lands near the format's maximum, then record the scale and fold it back into the accumulation. Per-tensor scaling with a delayed-update history (keep the max over the last \( k \) steps and use it to pick this step's scale) is what NVIDIA's Transformer Engine implements. Finer-grained schemes scale per row, per column, or per block. DeepSeek-V3 trained essentially the whole forward and backward in fp8 with per-block scaling of \( 128\times128 \) weight tiles and \( 1\times128 \) activation tiles, plus fp32 accumulation promoted out of the tensor core at intervals, and reported it as the first public validation of fp8 at very large scale.

The payoff is measured here directly. An \( 8192^3 \) matmul runs at 780.5 TFLOP/s in bf16 and 1,425.1 TFLOP/s in fp8 e4m3 through torch._scaled_mm, a 1.83× speedup, close to the 2× the hardware promises. Whether that translates into a 1.8× faster training step depends on how much of the step is matmul, which the MFU discussion above already answers. At frontier width, most of it.

Memory, the full training accounting

A training step holds four categories of memory. Let \( N \) be the parameter count.

Parameters. In pure fp32, \( 4N \) bytes. In mixed precision with fp32 master weights, \( 4N \) for the master copy plus \( 2N \) for the bf16 working copy, though frameworks often materialize the bf16 copy transiently.

Gradients. One per parameter, usually fp32, \( 4N \) bytes.

Optimizer state. Adam keeps a first moment and a second moment per parameter, both fp32, \( 8N \) bytes. This is the single largest static cost and the reason Adam is expensive.

Activations. Everything saved by the forward pass for use in the backward. This scales with batch size and sequence length, not with the parameter count, and is the only term the engineer controls at run time.

The static total for mixed-precision Adam is therefore about \( 16N \) bytes, 4 master + 2 working + 4 gradient + 8 optimizer, or \( 18N \) if the bf16 copy is persistent. For pure fp32 Adam it is \( 4 + 4 + 8 = 16N \) as well. The famous consequence is

An 8B model does not fit on an 80GB GPU for training, by a factor of 1.5, before a single activation is allocated. With ZeRO stage 3 across 8 GPUs, each device holds \( 16N/8 = 14.9 \) GiB of static state and has 65GB left for activations, which is comfortable. This one calculation is why sharded training exists. Inference is a different story, since 8B parameters in bf16 is 16GB, and the same GPU serves it easily.

The accounting was checked against the allocator here. For the reference model at batch 32 and context 512 in bf16 autocast with fused AdamW, the predicted static state is \( 16 \times 2.675\times10^7 = 4.28\times10^8 \) bytes = 0.399 GiB, {{MEM_ACCOUNT}}

Activation memory and the checkpointing tradeoff

Without any checkpointing, the activations saved per layer per token are roughly the layer input (\( d \)), the two norm outputs (\( 2d \)), the query, key, and value (\( d(1 + 2H_{kv}/H) \)), the attention output (\( d \)), the two MLP up-projections (\( 2F \)), and the gated product (\( F \)). With \( F = \frac{8}{3}d \) and GQA at \( H_{kv}/H = 1/4 \) that is about \( 13.5\,d \) elements per token per layer, or \( 27 d \) bytes in bf16. For a 70B model (\( d = 8192 \), \( L = 80 \)) at batch 1 and context 8192, that is \( 27 \times 8192 \times 8192 \times 80 = 1.45\times10^{11} \) bytes, or 135 GiB, for one sequence. Activations, not parameters, are what makes long-context training hard.

Gradient checkpointing, from Chen and colleagues in 2016, trades compute for memory, storing only a subset of the activations and recomputing the rest during the backward pass. The classic analysis is worth doing. Suppose the network has \( L \) layers and we store the input to every \( k \)-th layer, so there are \( L/k \) checkpoints. Memory is then \( O(L/k) \) for the checkpoints plus \( O(k) \) for the segment being recomputed, total \( O(L/k + k) \), minimized at \( k = \sqrt{L} \) giving \( O(\sqrt{L}) \) memory. The compute cost is one extra forward pass over the recomputed segments, so the step costs roughly \( 4/3 \) of the un-checkpointed step (forward 1, recompute 1, backward 2, against forward 1, backward 2).

Measured here on a 24-layer version of the reference model at batch 32 and context 512, checkpointing every \( k \)-th layer gives the results below.

The measured overhead of full checkpointing on the 8-layer model was 22.4 ms to 26.8 ms, a 20% slowdown, against a 3× memory reduction from 5.81 GB to 1.94 GB. That is a much better trade than the theoretical \( 4/3 \) suggests, because the recomputed forward passes are pure matmul while the memory saved would otherwise force a smaller batch, and smaller batches are less efficient. In practice, selective checkpointing (recompute the cheap elementwise operations, keep the expensive matmul outputs), which PyTorch exposes as selective activation checkpointing and torchtitan uses by default, gets most of the memory saving for a few percent of the compute.

Parallelism, with the communication costs derived

No frontier model is trained on one device. The five ways to split the work are orthogonal in principle and combined in practice, and each has a communication pattern whose cost can be computed rather than guessed. Throughout, \( P \) is the number of devices, \( \alpha \) the per-message latency, and \( \beta \) the inverse bandwidth in seconds per byte.

  data parallel        replicate model, split batch      all-reduce gradients
  ZeRO / FSDP          shard optimizer/grad/param        all-gather + reduce-scatter
  tensor parallel      split each matmul across devices  all-reduce activations
  pipeline parallel    split layers across devices       point-to-point activations
  sequence / context   split the token axis              ring / all-to-all
  expert parallel      split experts across devices      all-to-all tokens
          

Data parallelism and the all-reduce cost model

Each device holds a full replica, processes a distinct slice of the global batch, and the gradients are averaged before the optimizer step. The averaging is an all-reduce over \( N \) values. The standard ring all-reduce, from Baidu's 2017 implementation of an old HPC algorithm, runs in two phases, a reduce-scatter in which each device ends up owning the fully reduced value for \( 1/P \) of the data, and an all-gather that distributes those pieces. Each phase is \( P - 1 \) steps, and each step sends \( N/P \) elements, so the total data leaving each device is

with \( b \) bytes per element. Two consequences follow. The bandwidth term is independent of \( P \) in the limit, approaching \( 2Nb\beta \), which is why ring all-reduce scales. Doubling the cluster does not double the per-device traffic. The latency term grows linearly in \( P \), which is why very large rings use hierarchical or tree algorithms instead.

The scaling question is whether the all-reduce hides behind the backward pass. Compute per step per device is \( 6 N B_{\text{local}} T / (\text{FLOP/s}) \) and communication is \( 2Nb\beta \), so the ratio is

The parameter count cancels. Whether data parallelism scales depends only on the tokens per device per step, the interconnect bandwidth, and the compute rate. With bf16 gradients (\( b = 2 \)), a 400 GB/s effective interconnect (\( \beta = 2.5\times10^{-12} \)), and 400 TFLOP/s of achieved compute, the ratio is \( 2\times2.5\times10^{-12}\times4\times10^{14} / (3 B_{\text{local}}T) = 2000/(3 B_{\text{local}} T) \), so communication is under 10% of compute once \( B_{\text{local}}T > 6{,}700 \) tokens per device per step. That is a small number, which is why plain data parallelism works well up to the point where the model stops fitting.

Measured on the two H100s in this machine, NCCL all-reduce bandwidth as a function of message size is shown below.

ZeRO and FSDP, what each stage shards

Data parallelism replicates \( 16N \) bytes of static state on every device, which is pure waste, since every replica holds identical optimizer moments. ZeRO, from Rajbhandari and colleagues at Microsoft in 2020, removes the redundancy in three stages.

The communication analysis is the interesting part. ZeRO-1 and ZeRO-2 are free. An all-reduce is already a reduce-scatter followed by an all-gather, so replacing "all-reduce the gradients, everyone updates everything" with "reduce-scatter the gradients, each device updates its shard, all-gather the parameters" moves the same bytes. ZeRO-3 is not free. The parameters must be gathered before each layer's forward and again before its backward, adding \( 2N b (P-1)/P \approx 2Nb \) bytes per step on top of the \( 2Nb \) of the gradient reduction, roughly a 50% increase in total volume for a 3× to \( P \)× reduction in memory. In exchange, model size stops being bounded by one device.

FSDP, PyTorch's implementation described by Zhao and colleagues at Meta in 2023, is ZeRO-3 with the sharding unit being a module (typically one transformer block) rather than a flat parameter vector, which allows the all-gather for layer \( \ell+1 \) to overlap with the compute of layer \( \ell \). That prefetching is what makes ZeRO-3 practical. Without it the gathers serialize and throughput collapses. FSDP2 refines this further with per-parameter sharding using DTensor.

Measured on the two-GPU configuration here, at the reference model size and batch 64 per device, the results are below.

Tensor parallelism, splitting the matmuls

Megatron-LM, from Shoeybi and colleagues at NVIDIA in 2019, splits individual weight matrices across devices so that a single layer's computation is distributed. The construction is chosen so that only one all-reduce per sublayer is needed, and the derivation is short enough to do in full.

MLP. The block computes \( Y = \sigma(XA)B \) with \( A \in \R^{d\times F} \) and \( B \in \R^{F \times d} \). Split \( A \) by columns, \( A = [A_1, A_2] \), so device \( i \) computes \( \sigma(XA_i) \in \R^{T \times F/2} \) using only its shard. This works because \( \sigma \) is elementwise, so the nonlinearity of a column block depends only on that column block. Then split \( B \) by rows, \( B = [B_1; B_2] \), so device \( i \) computes \( \sigma(XA_i)B_i \in \R^{T\times d} \), a full-width partial result. The final output is \( Y = \sum_i \sigma(XA_i)B_i \), one all-reduce. Column split, then row split, one all-reduce at the end of the MLP. If \( B \) had been split by columns instead, the sum would have been needed before \( \sigma \), forcing a second communication.

Attention. The head structure makes this even cleaner. Give each device a disjoint subset of the heads. Device \( i \) holds the \( W_Q, W_K, W_V \) columns for its heads (a column split), computes its heads' attention entirely locally (attention does not mix heads), and holds the corresponding rows of \( W_O \) (a row split), producing a partial sum. One all-reduce. So a transformer layer costs exactly two all-reduces forward and two backward, each over the full activation tensor \( B T d \).

The cost model is now unfavourable in a specific way. Per layer per step, tensor parallelism moves \( 4 \cdot 2 B T d (P-1)/P \) bytes of activations, while data parallelism moves \( 2Nb \) bytes of gradients once per step for the whole model. The activation traffic scales with batch and sequence length and happens \( 2L \) times per step. That is why tensor parallelism is restricted to devices connected by NVLink, typically the 8 GPUs in one node. At 900 GB/s intra-node it hides behind compute, and at 50 GB/s inter-node it does not. The correctness of the construction was verified here. A hand-written column-then-row split of a \( 1024 \times 4096 \) MLP across the two GPUs, followed by one all-reduce, matches the single-device result to machine precision.

Sequence parallelism, added to Megatron by Korthikanti and colleagues in 2022, notices that the parts of a layer that are not tensor-parallel, the norms and dropout, operate per token independently and are therefore replicated across every tensor-parallel rank, wasting activation memory. Splitting those along the sequence axis and converting the all-reduce into a reduce-scatter plus all-gather pair moves exactly the same bytes while cutting activation memory by the tensor-parallel degree.

Pipeline parallelism and the bubble

Pipeline parallelism assigns contiguous groups of layers to different devices. Device 1 runs layers 1 to \( L/P \), passes activations to device 2, and so on. Communication is point-to-point and tiny, one activation tensor of size \( BTd \) per boundary, against tensor parallelism's \( 2L \) all-reduces of the same size. The cost is idle time.

Deriving the bubble fraction. With \( P \) stages and a batch split into \( m \) microbatches, in the simple GPipe schedule (all forwards, then all backwards) each stage does \( m \) forward units and \( m \) backward units of work. Stage \( i \) cannot start until stage \( i-1 \) has produced its first microbatch, so the pipeline takes \( P - 1 \) extra forward units to fill and \( P-1 \) extra backward units to drain. Total time in units of one microbatch stage-pass is

against the ideal \( m(t_f + t_b) \). The bubble fraction is

With \( P = 8 \) stages and \( m = 8 \) microbatches the bubble is \( 7/15 = 47\% \), so almost half the pipeline sits idle. With \( m = 64 \) it is \( 7/71 = 9.9\% \). The rule is \( m \ge 4P \) and preferably \( m \ge 8P \). The catch is that GPipe must keep the activations of all \( m \) in-flight microbatches, so raising \( m \) raises activation memory linearly.

1F1B, introduced in PipeDream from Microsoft and CMU and adopted by Megatron, fixes the memory problem without changing the bubble. After the fill phase, each stage alternates one forward and one backward, so a stage holds at most \( P \) microbatches' activations rather than \( m \). Same bubble fraction, memory independent of \( m \).

Interleaved 1F1B, from Narayanan and colleagues in 2021, reduces the bubble itself. Instead of giving each device one contiguous block of \( L/P \) layers, give it \( v \) non-contiguous chunks of \( L/(Pv) \) layers each. Each microbatch now passes through each device \( v \) times, so the fill and drain phases shrink by a factor of \( v \),

The price is \( v \)× more point-to-point communication. Megatron uses \( v = 2 \) or 4 routinely. Zero-bubble and "breadth-first" schedules from later work split the backward pass into its input-gradient and weight-gradient halves, which can be scheduled independently, and get the theoretical bubble near zero at the cost of considerable scheduling complexity.

Sequence and context parallelism, including ring attention

When the context is 128K or a million tokens, even one sequence's activations exceed a device. Context parallelism splits the sequence axis across devices, so device \( i \) owns tokens \( [iT/P, (i+1)T/P) \). Everything except attention is token-local and needs no communication at all. Attention needs every query to see every key, which is where ring attention, from Liu, Zaharia, and Abbeel at Berkeley in 2023, comes in.

The construction reuses the online-softmax trick that makes FlashAttention work. Attention over a partitioned key set can be computed incrementally. Maintain a running maximum \( m \), a running denominator \( \ell \), and a running weighted sum \( o \). For each new key block, compute its scores, update \( m' = \max(m, m_{\text{block}}) \), rescale the accumulators by \( e^{m - m'} \), and add the block's contribution. The result is exactly the same as computing the softmax over all keys at once. Ring attention arranges the devices in a ring, and while device \( i \) computes attention of its local queries against its current key/value block, it simultaneously sends that block to device \( i+1 \) and receives one from \( i-1 \). After \( P \) rounds every query has seen every key. The communication of \( 2 B T d / P \) bytes per round overlaps with \( O(B T^2 d/P^2) \) compute per round, so for long enough sequences the communication is fully hidden. This is what makes million-token context training possible, and it is how Llama 3's 128K context stage was trained.

Expert parallelism, and how it all composes

Mixture-of-experts models have a natural extra axis, putting different experts on different devices. A token routed to expert \( e \) must travel to the device holding \( e \), be processed, and travel back, which is an all-to-all in each direction. The volume is \( 2 \times k B T d \) bytes per layer per step for top-\( k \) routing, and unlike an all-reduce it is a personalized exchange whose cost is sensitive to load imbalance. If one expert receives three times the average traffic, its device becomes the critical path.

Real configurations combine four or five of these axes. The Llama 3 405B run used tensor parallelism 8 (within a node), pipeline parallelism 16, context parallelism up to 16 for the long-context stage, and data parallelism (as FSDP) for the remainder, on 16,384 H100s. DeepSeek-V3 used a 16-way pipeline, 64-way expert parallelism, and ZeRO-1 data parallelism, with no tensor parallelism at all. The heuristic for choosing is as follows.

  1. Tensor parallel first, but only within a node (NVLink). Degree 8 max.
  2. Pipeline parallel next, to fit the model. Keep microbatches >= 4 * stages.
  3. Expert parallel if the model is MoE; overlaps with the pipeline dimension.
  4. Context parallel only if one sequence's activations do not fit.
  5. Data parallel (FSDP/ZeRO) with everything left over. This is the axis
     that scales, so make it as large as the batch-size budget allows.
          

The binding constraint on the last line is the critical batch size derived in the optimization section. Data parallelism raises the global batch, and past a certain batch size extra tokens per step stop buying proportional progress.

Mixture of experts

The MLP is two thirds of the parameters and does the same work for every token. A mixture of experts replaces it with \( E \) parallel MLPs and a router that sends each token to \( k \) of them, so the parameter count multiplies by roughly \( E/k \) while the FLOPs per token stay fixed. This is the only known way to buy capacity without buying compute, and it is why most frontier models released since 2024 are sparse.

Routing. A linear router produces \( g = \softmax(W_r x) \in \R^E \). The top \( k \) entries are selected, renormalized, and used as mixture weights,

Switch Transformer, from Fedus, Zoph, and Shazeer at Google in 2021, uses \( k = 1 \) and argues that even a single expert suffices if the router is trained well, while GShard, from Lepikhin and colleagues, and most subsequent work use \( k = 2 \). Mixtral 8x7B uses \( E = 8 \), \( k = 2 \). DeepSeek-V3 uses 256 routed experts with \( k = 8 \) plus one shared expert, and Qwen3's MoE variants use 128 experts with \( k = 8 \).

The load-balancing loss, derived. Nothing in the objective prevents the router from collapsing onto a few experts, and collapse is self-reinforcing, since an expert that receives more tokens trains faster and becomes more attractive. The standard fix is an auxiliary loss that penalizes imbalance. Let \( f_e \) be the fraction of tokens in the batch routed to expert \( e \), and \( P_e \) the mean router probability assigned to expert \( e \) over the batch,

Why this form? By Cauchy-Schwarz or simply by the fact that \( \sum_e f_e = k \) and \( \sum_e P_e = 1 \), the dot product \( \sum_e f_e P_e \) is minimized when the mass is spread out and maximized when it concentrates. If both are uniform, \( f_e = k/E \) and \( P_e = 1/E \), giving \( \L_{\text{aux}} = E \cdot E \cdot (k/E)(1/E) = k \). With \( k = 1 \) the balanced value is exactly 1, which is why the loss is scaled by \( E \). If all traffic goes to one expert, \( f_1 = k \), \( P_1 \approx 1 \), and \( \L_{\text{aux}} \approx Ek \), which is \( E \) times worse. The reason \( f \) and \( P \) both appear rather than penalizing \( f \) alone is differentiability, since \( f_e \) involves an argmax and has zero gradient almost everywhere, while \( P_e \) is smooth, so the product gives a gradient that pushes probability mass away from overloaded experts, with \( f \) acting as a constant multiplier measuring how overloaded each one is.

Capacity factor and token dropping. In a distributed implementation each expert has a fixed buffer, sized as \( \text{capacity} = c \cdot k|\mathcal{B}|/E \) with capacity factor \( c \) typically 1.0 to 1.25. Tokens arriving at a full expert are dropped, meaning they skip the MLP entirely and pass through on the residual. This is the price of static shapes, which the all-to-all requires. Raising \( c \) reduces dropping and raises memory and communication. The balancing loss is what keeps \( c \) near 1 workable.

Shared experts and auxiliary-loss-free balancing. DeepSeek's MoE line adds one or more experts that every token visits unconditionally, on the argument that some computation is universally useful and forcing the routed experts to rediscover it wastes capacity. DeepSeek-V3 goes further and removes the auxiliary loss entirely, replacing it with a per-expert bias added to the routing logits that is adjusted online, up when an expert is underloaded and down when it is overloaded. The advantage is that the balancing pressure no longer perturbs the language-modelling gradient, which they report improves quality at equal balance.

Total versus active parameters. The distinction is the whole point and is frequently misreported. Take the real numbers. Mixtral 8x7B has 46.7B total parameters and 12.9B active per token (not 8×7 = 56B, because attention and embeddings are shared). DeepSeek-V3 has 671B total and 37B active. Qwen3-235B-A22B has 235B total and 22B active. Llama 4 Maverick has 400B total and 17B active. Memory and communication scale with total, FLOPs scale with active, and quality sits somewhere in between, empirically closer to a dense model of size \( \sqrt{N_{\text{total}} N_{\text{active}}} \) in several published comparisons.

Measured here at small scale, a model with 8 experts of half the baseline MLP width and top-2 routing was trained against a dense baseline at the same token budget, with and without the auxiliary loss. The load-balance statistics show the effect clearly.

Data

Data is where the largest quality differences between otherwise identical models come from, and it is the part of the pipeline with the least published theory and the most published engineering. The structure of a web-scale pipeline is stable across the open descriptions. RefinedWeb from the Technology Innovation Institute, FineWeb from Hugging Face, Dolma from Allen AI, and the Llama 3 report all describe the same stages in roughly the same order.

Extraction. Common Crawl ships WARC files of raw HTTP responses. Turning them into text means stripping boilerplate, navigation, and markup. The choice of extractor matters more than it sounds. FineWeb's ablations found that trafilatura recovers noticeably better text than the WET files Common Crawl provides, and that this alone moves downstream benchmark scores.

Language identification. A fastText classifier assigns a language and a confidence, and documents below a threshold (0.65 in FineWeb) are dropped. This is also where the multilingual mixture is decided, and dropping the threshold too low admits machine-translated and machine-generated text.

Quality filtering. Two families. Heuristic filters, inherited from the Gopher work at DeepMind, are cheap rules, such as mean word length in a plausible range, at least some fraction of lines ending in punctuation, symbol-to-word ratio below a threshold, a minimum number of stop words, no excessive repetition of lines or n-grams. Model-based filters train a classifier to distinguish "high-quality" reference text (Wikipedia, books, curated web) from random crawl and keep the high-scoring tail. GPT-3 used a classifier of this kind. Llama 3 used a DistilRoBERTa quality classifier and later a Llama-2-based one, and FineWeb-Edu used a classifier trained on Llama-3-70B judgments of educational value and reported large gains on knowledge benchmarks.

Deduplication. The web is enormously redundant, and Lee and colleagues at Google showed in 2021 that deduplicating training data reduces memorization, improves perplexity, and lets models reach the same loss in fewer steps. Two mechanisms are used together. Fuzzy document-level deduplication uses MinHash-LSH. Represent each document by the set of its \( n \)-grams, estimate Jaccard similarity between documents by the fraction of agreeing minimum hash values, and use locality-sensitive hashing to find candidate near-duplicates without an all-pairs comparison. With \( h \) hash functions split into \( b \) bands of \( r = h/b \) rows, two documents with Jaccard similarity \( s \) collide in at least one band with probability

an S-curve whose threshold sits near \( s^\star = (1/b)^{1/r} \). The typical configuration of \( h = 112 \), \( b = 14 \), \( r = 8 \) puts the threshold at \( (1/14)^{1/8} = 0.72 \), catching documents sharing roughly three quarters of their 5-grams. Exact substring deduplication builds a suffix array over the concatenated corpus and removes any span of 50 or more tokens that appears more than once, which catches the boilerplate and license text that survives document-level matching. The suffix array construction is \( O(n \log n) \) and has been run on corpora of a trillion tokens.

Decontamination. Remove from training any document containing an evaluation instance. The standard method is \( n \)-gram overlap against every benchmark's test set, usually with \( n = 13 \) as in GPT-3, sometimes with normalized-text matching. This is harder than it sounds because benchmarks are themselves derived from web text and because paraphrases evade \( n \)-gram matching entirely, which is why every serious report includes a contamination analysis and treats it as a lower bound.

PII removal. Regex and model-based detection of email addresses, phone numbers, IP addresses, and identifiers, replaced with placeholders. Dolma documents its exact rules, which is unusual and useful.

The measured effect of quality. This pipeline was implemented here at small scale and its effect measured. Starting from 100MB of raw Wikipedia XML, applying markup stripping, a 400-byte minimum length, a 90% alphanumeric-character-ratio threshold, a stop-word test requiring at least four of the twenty most common English words, and MD5-based near-duplicate removal on the first 60 words, kept 7,588 of 12,344 documents and 65.2% of bytes. Scaled up to the full 1GB dump, the same pipeline kept 134,666 documents totalling 557MB, which tokenized to 163.1M tokens at 3.416 bytes per token. Training the identical 26.7M-parameter model for 65.5M tokens on the filtered corpus against the unfiltered one, and evaluating both on the same clean held-out set, gives the results below.

The filtered corpus is 0.21 nats better than the raw one at an identical compute budget, which at these loss levels is worth roughly a 2× increase in compute judging by the isoFLOP curves below. That is the entire argument for data engineering in one number. The half-and-half mixture landing marginally below pure filtered is within the seed noise measured later, and is a reminder that a small amount of messy text is not harmful. What is harmful is a corpus that is mostly markup.

Mixing, curriculum, and multiple epochs

A pretraining corpus is a weighted mixture of sources, such as web, code, books, papers, encyclopedic text, mathematics, and multilingual data. The weights are hyperparameters and they matter. The Llama 3 report describes arriving at roughly 50% general knowledge, 25% mathematics and reasoning, 17% code, and 8% multilingual by running scaling-law experiments on candidate mixtures at small scale and extrapolating. DoReMi, from Xie and colleagues, automates this by training a small proxy model with group-distributionally-robust optimization to find weights that minimize worst-case excess loss across domains, then reusing those weights for the large run.

Curriculum in the strict sense (easy examples first) has never robustly helped language model pretraining. What does help is a change of mixture near the end of training, sometimes called annealing or mid-training, where the last few percent of tokens are drawn from a much higher-quality, more instruction-like mixture while the learning rate decays. Llama 3 reports this, and the WSD schedule below is designed around it.

How many epochs. Muennighoff and colleagues studied this directly in 2023 and found that repeating data is nearly as good as fresh data for about four epochs, after which returns decay rapidly and by around 40 epochs additional training adds nothing. Measured here at fixed compute, holding the token budget at 65.5M and shrinking the unique-token pool to force repetition, the results are below.

Optimization

AdamW and the decoupling derivation

Adam maintains exponential moving averages of the gradient and its square,

corrects their initialization bias with \( \hat m_t = m_t/(1-\beta_1^t) \), \( \hat v_t = v_t/(1-\beta_2^t) \), and steps \( \theta_t = \theta_{t-1} - \eta\,\hat m_t/(\sqrt{\hat v_t}+\epsilon) \). The bias correction is exact. Unrolling \( m_t = (1-\beta_1)\sum_{i=1}^{t}\beta_1^{t-i} g_i \) and assuming stationary gradients with mean \( \bar g \) gives \( \E[m_t] = \bar g(1-\beta_1)\sum \beta_1^{t-i} = \bar g(1-\beta_1^t) \), so dividing by \( 1-\beta_1^t \) restores the mean.

Why decoupling matters. \( L_2 \) regularization adds \( \frac{\lambda}{2}\|\theta\|^2 \) to the loss, so the gradient becomes \( g_t + \lambda\theta_{t-1} \). Feed that into Adam and the decay term passes through the preconditioner,

The effective decay on parameter \( j \) is therefore \( \eta\lambda\theta_j/(\sqrt{\hat v_j}+\epsilon) \), which is inversely proportional to the historical gradient magnitude of that parameter. Parameters with large gradients get almost no decay, while parameters with tiny gradients get very large decay. That is the opposite of what regularization is supposed to do, and it makes the effective regularization strength depend on the learning rate in a way that couples the two hyperparameters. Loshchilov and Hutter's AdamW, from 2017, simply removes the decay from the gradient and applies it directly to the weights,

Now every parameter shrinks by the same factor \( 1 - \eta\lambda \) per step regardless of its gradient history, the two hyperparameters decouple, and \( \lambda \) can be tuned once and reused. Every language model uses AdamW.

Hyperparameter choices and why. \( \beta_1 = 0.9 \) is nearly universal. \( \beta_2 \) is set to 0.95 rather than the default 0.999 in essentially every language model, because \( \beta_2 = 0.999 \) has an effective averaging window of 1,000 steps and language modelling gradients are non-stationary enough over that horizon that a stale \( v \) causes instability. A value of 0.95 gives a window of 20 steps. \( \epsilon = 10^{-8} \) is standard, and lowering it to \( 10^{-15} \) is a known instability fix when \( \sqrt{v} \) becomes very small. Weight decay is applied to matrices but not to biases, norm gains, or (usually) embeddings, because decaying a norm gain toward zero directly attacks the scale invariance the norm exists to provide.

Schedules

Warmup raises the learning rate linearly from zero over the first 1% to 3% of steps. The justification is Adam's second moment. At step 1, \( \hat v_1 = g_1^2 \) exactly, so the update is \( \eta\,\mathrm{sign}(g_1) \) with no averaging, and the variance of the update direction is at its maximum. Warmup lets \( v \) accumulate before large steps are taken. For post-LN models it is additionally required by the gradient-scale argument above.

Cosine decay takes the rate from its peak to a floor (typically 10% of peak) following \( \eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max}-\eta_{\min}) (1 + \cos(\pi t/T)) \). Chinchilla established that the cosine period must equal the planned training length. A cosine tuned for \( 2T \) steps but stopped at \( T \) leaves the model significantly undertrained, which was one of the errors in the earlier Kaplan-era scaling estimates.

WSD (warmup-stable-decay), from the MiniCPM work and independently from several groups, holds the rate constant for most of training and decays it sharply over the last 10% to 20%. The advantage is operational. With cosine you must know \( T \) in advance, and every intermediate checkpoint is at a high learning rate and therefore far from a usable model. With WSD, any checkpoint from the stable phase can be branched, decayed for 10% of the remaining budget, and evaluated as a finished model. That makes scaling-law data collection cheap, makes it possible to extend a run that turned out to be shorter than desired, and makes the annealing-on-high-quality-data trick natural, since the decay phase is exactly where the mixture is changed. Measured here at a 65.5M-token budget, the results are below.

Batch size, gradient noise, and the critical batch size

Increasing the batch size reduces gradient noise, which allows a larger learning rate, which means fewer steps to reach a given loss. But the reduction in noise is only \( 1/\sqrt{B} \) in standard deviation, so there is a point past which doubling the batch halves the number of steps but doubles the cost per step, buying nothing. McCandlish, Kaplan, and colleagues at OpenAI formalized this in 2018 with the gradient noise scale. Model the per-example gradient as having mean \( G \) and covariance \( \Sigma \). The expected improvement from one step of size \( \eta \) with a batch of size \( B \) is, to second order,

Optimizing over \( \eta \) gives \( \eta^\star = \eta_{\max}/(1 + B_{\text{noise}}/B) \) with \( B_{\text{noise}} = \tr(H\Sigma)/(G\T H G) \), and substituting back,

So the improvement per step saturates as \( B \to \infty \) at \( \Delta L_{\max} \), and the number of steps to reach a target loss follows the hyperbola \( S/S_{\min} = 1 + B_{\text{noise}}/B \), equivalently \( (S/S_{\min} - 1)(E/E_{\min} - 1) = 1 \) where \( E = BS \) is the total examples processed. At \( B = B_{\text{noise}} \) you take twice the minimum steps and use twice the minimum data, which is the natural definition of the critical batch size. Empirically \( B_{\text{noise}} \) grows during training, roughly as a power of \( 1/L \). The noisier the gradient relative to its mean, the more averaging helps, and gradients become relatively noisier as the loss falls. This is why frontier runs ramp the batch size upward during training. Llama 3 ramped from 4M to 16M tokens per batch.

Measured here at a fixed 65.5M-token budget, varying the batch and scaling the learning rate as \( \sqrt{B} \) gives the results below.

Clipping, z-loss, initialization, and loss spikes

Gradient clipping rescales the whole gradient when its global norm exceeds a threshold, \( g \leftarrow g \cdot \min(1, c/\|g\|) \), with \( c = 1.0 \) essentially universal. It preserves the direction and bounds the step, and it is the cheapest insurance against a single bad batch destroying a month of training. In the main run here, 67 of 16,384 steps had gradient norm above 1.0 and were clipped, or 0.4%.

Z-loss adds \( \gamma\,(\log\sum_i e^{z_i})^2 \) to the objective, penalizing the log-partition function of the output softmax. Introduced in PaLM and used in several later models with \( \gamma = 10^{-4} \), it prevents the logits from drifting to large absolute values, which matters because the softmax is shift-invariant and therefore nothing else constrains their absolute scale. Large logits cause bf16 roundoff in the loss and are correlated with instability.

Depth-dependent initialization. The residual stream accumulates \( 2L \) branch outputs. If each branch output has variance \( \sigma^2 \) at initialization, the stream's variance after \( L \) layers is \( 2L\sigma^2 \), growing linearly with depth. GPT-2 fixed this by scaling the initialization of the projections that write into the residual stream (the attention output and the MLP down-projection) by \( 1/\sqrt{2L} \), so the total added variance is \( O(1) \) regardless of depth. Measured here, removing that scaling raised validation loss from 3.566 to 3.602 at a matched budget on an 8-layer model. The effect grows with depth, and at 100 layers it is the difference between training and not.

Loss spikes are the characteristic failure of large runs. The loss jumps by one or several nats over a few steps and either recovers slowly or diverges. The published causes are consistent. Attention logit growth leading to entropy collapse is the mechanism ViT-22B and several language models hit, and QK-norm fixes it. Bad data batches, especially long runs of repeated tokens, produce very large gradients. Optimizer state staleness, when \( \sqrt{v} \) has shrunk far below the current gradient scale, makes one large gradient produce an outsized step, which is what \( \beta_2 = 0.95 \) and \( \epsilon \) tuning address. And plain numerical overflow in fp16 without adequate scaling. The mitigations, in order of how often they are used, are to clip the gradient, lower the peak learning rate, use bf16, use QK-norm and z-loss, skip batches whose loss exceeds a running threshold, and, when all else fails, roll back to a checkpoint from before the spike and skip the offending data. PaLM's report describes doing exactly this and finding the spikes did not recur on the same data from a different checkpoint, which is evidence that the spikes were an interaction between specific data and specific optimizer state rather than a property of the data alone.

Measured here, deliberately provoking spikes with a 10× learning rate gives the table below.

Lion, Sophia, Muon, and how little separates them

Three post-Adam optimizers have real followings. Lion, discovered by symbolic program search at Google Brain in 2023, keeps only a momentum buffer and takes the sign of an interpolation, \( u_t = \mathrm{sign}(\beta_1 m_{t-1} + (1-\beta_1)g_t) \), \( \theta_t = \theta_{t-1} - \eta(u_t + \lambda\theta_{t-1}) \), with the momentum updated afterwards using a different \( \beta_2 \). It halves optimizer memory (one state instead of two) and, because every update has the same magnitude, requires a learning rate roughly 3 to 10 times smaller and a weight decay correspondingly larger.

Sophia, from Liu, Li, Hsieh and colleagues at Stanford and elsewhere in 2023, is a diagonal second-order method. It estimates the Hessian diagonal with a Hutchinson or Gauss-Newton-Bartlett estimator every \( k \) steps, and clips the preconditioned update elementwise. The paper reported roughly 2× fewer steps to a given loss on GPT-2-scale models. The result has been difficult to reproduce at larger scale with tuned baselines, which is the recurring pattern for optimizer claims.

Muon, from Jordan and collaborators in 2024, is the most interesting recent entry. It treats each 2-D weight matrix as a matrix rather than a bag of scalars. Take the momentum buffer \( M \), compute its nearest orthogonal matrix \( \mathrm{orth}(M) = UV\T \) from the SVD \( M = U\Sigma V\T \), and step in that direction. Computing an SVD every step is impossible, so Muon uses a fixed five-step Newton-Schulz iteration with coefficients tuned to converge quickly in bf16. The justification is that Adam's per-coordinate normalization ignores the matrix structure, and that orthogonalizing the update equalizes its effect across the singular directions, which is a form of steepest descent under the spectral norm. Embeddings, the output head, and all 1-D parameters stay on AdamW, because those are not structurally matrices in the relevant sense. Muon was used for Kimi K2 at 1T parameters, which is the largest public deployment of a non-Adam optimizer.

The honest summary, which the measurement below supports, is that at a matched budget with each optimizer's learning rate tuned, the differences are small. Reported speedups of 2× are almost always against an untuned Adam baseline, and the field's experience with LAMB, LARS, Shampoo, AdaFactor, and half a dozen others is that the gap closes when the baseline is tuned properly. Muon is the one that has survived contact with a frontier-scale run, and even there the claimed advantage is on the order of tens of percent in tokens-to-loss rather than a factor of two.

Scaling laws

The power-law form

Kaplan and colleagues at OpenAI established in 2020 that language model loss falls as a power law in each of model size, dataset size, and compute, over more than five orders of magnitude, with the exponents stable across architectures. The functional form that has held up, and that Hoffmann and colleagues fit at DeepMind, is

with \( E \) the irreducible entropy of the text, the second term the cost of a finite model, and the third the cost of finite data. Chinchilla's fitted values are \( E = 1.69 \), \( A = 406.4 \), \( B = 410.7 \), \( \alpha = 0.34 \), \( \beta = 0.28 \), with \( L \) in nats per token. The near-equality of \( \alpha \) and \( \beta \) is the whole story, as the next derivation shows.

Compute-optimal allocation, derived

Fix a compute budget \( C = 6ND \) and minimize \( L(N,D) \) subject to it. Substituting \( D = C/(6N) \) gives

Differentiating and setting to zero gives

Solving for \( N \) gives

With Chinchilla's \( \alpha = 0.34 \) and \( \beta = 0.28 \), \( \beta/(\alpha+\beta) = 0.28/0.62 = 0.452 \) and \( \alpha/(\alpha+\beta) = 0.548 \). So \( N^\star \propto C^{0.45} \) and \( D^\star \propto C^{0.55} \), close enough to \( C^{0.5} \) each that parameters and tokens should scale in roughly equal proportion, and the ratio \( D^\star/N^\star \) is nearly constant. Evaluating the constant from Chinchilla's own fit gives \( D^\star \approx 20 N^\star \), that is, 20 tokens per parameter.

Kaplan's earlier fit gave \( N \propto C^{0.73} \), that is, spend almost everything on parameters. The discrepancy was traced to three methodological issues. Kaplan's runs used a cosine schedule whose period did not match the training length, so shorter runs were systematically undertrained, the learning rate was not retuned per model size, and small models were excluded in a way that biased the fit. Chinchilla's correction is why a 70B model trained on 1.4T tokens (Chinchilla itself) outperformed a 280B model trained on 300B tokens (Gopher) at a quarter of the inference cost.

The isoFLOP methodology, run here

Hoffmann and colleagues used three approaches, of which the isoFLOP profile is the most direct and the easiest to replicate. Fix a compute budget, train several model sizes with \( D = C/(6N) \) tokens each, plot final loss against \( N \), and read off the minimum. Repeat at several budgets and fit the trajectory of minima.

That experiment was run here at three budgets spanning an order of magnitude, with a ladder of seven model shapes from 3.1M to 80.6M parameters, all sharing the same tokenizer, data, and schedule. The resulting curves are below.

The minimum of each curve is marked. The compute-optimal model size and token count found here are as follows.

Fitting \( N^\star \propto C^{a} \) to these three minima gives \( a = 0.012 \), and \( D^\star \propto C^{b} \) gives \( b = 0.988 \), against Chinchilla's 0.45 and 0.55. The measured tokens-per-parameter ratios at the optima are 23.5, 13.5, 234.9. Two caveats are owed. The exponents are fit from three budgets and a coarse seven-point ladder, so the uncertainty is large. And at the small-model end of each curve the run exceeds one epoch over the 173M-token corpus, up to 8.8 epochs at the extreme, which penalizes small models and biases the estimated optimum upward. The qualitative result, a clear interior minimum that moves right as compute grows, is exactly the Chinchilla isoFLOP picture reproduced on one GPU in under an hour.

Why inference cost pushes past compute-optimal

Chinchilla minimizes training compute. Almost nobody wants that. The quantity a deployed model minimizes is total cost over its lifetime, training plus inference. Serving \( D_{\text{inf}} \) tokens with a model of \( N \) parameters costs about \( 2N D_{\text{inf}} \) FLOPs, so

Minimizing this at fixed quality means trading a smaller \( N \) (cheaper forever) against a larger \( D_{\text{train}} \) (a one-time cost). Since the loss surface \( L(N,D) \) is flat near its constrained minimum, a substantial reduction in \( N \) can be bought with a modest increase in \( D \). Sardana and Frankle at Databricks formalized this in 2023 and showed that for realistic inference volumes the optimum sits far past 20 tokens per parameter.

The arithmetic on a concrete case. Llama 3 8B was trained on 15T tokens, giving \( D/N = 1875 \), almost 100× Chinchilla-optimal. Was that rational? The training cost was \( 6 \times 8\times10^9 \times 1.5\times10^{13} = 7.2\times10^{23} \) FLOPs. Chinchilla-optimal at that budget would be \( N^\star \approx \sqrt{C/(6\cdot 20)} = \sqrt{7.2\times10^{23}/120} = 7.75\times10^{10} \), a 77B model trained on 1.55T tokens, which would be roughly as good and would cost \( 2 \times 7.75\times10^{10} = 1.55\times10^{11} \) FLOPs per generated token against the 8B model's \( 1.6\times10^{10} \), a factor of 9.7. Break-even is where the inference saving repays the training premium. Here the 8B model was already cheaper in total after roughly \( 7.2\times10^{23}/(1.39\times10^{11}) \approx 5\times10^{12} \) generated tokens, which a widely deployed model passes in weeks. Overtraining a small model is a rational response to inference economics, not a mistake.

The second reason is that the loss surface is flat. Using the Chinchilla fit, a model at \( D/N = 20 \) and one at \( D/N = 200 \) with the same training compute differ in loss by a few hundredths of a nat, while differing in inference cost by more than 3×. Nobody would take the first trade.

What scaling laws do not predict

Scaling laws predict cross-entropy on a held-out sample of the training distribution. They do not predict downstream task accuracy, which is a thresholded, discontinuous function of the same underlying capability and therefore appears to "emerge". Schaeffer, Miranda, and Koyejo at Stanford argued in 2023 that emergence is largely an artifact of discontinuous metrics, and that continuous metrics on the same tasks improve smoothly. They do not predict the effect of data quality, which shifts the entire curve (the 0.21 nats measured above is a shift, not a slope change). They do not predict instruction-following, reasoning, or safety behaviour, all of which are post-training phenomena. They do not transfer across tokenizers, since the units of \( L \) are nats per token and a different tokenizer changes what a token is. Comparing across tokenizers requires bits per byte, \( \mathrm{bpb} = L/(\ln 2 \cdot \beta) \) with \( \beta \) the bytes per token. And they say nothing about architecture beyond the observation that, within the transformer family, the exponents are largely insensitive to depth, width, and head count as long as the aspect ratio stays in a reasonable range.

Evaluation during training

Perplexity is \( \exp(L) \) with \( L \) the mean cross-entropy in nats per token. It is the right metric for monitoring a run and the wrong metric for comparing models across tokenizers, for the reason just given. Bits per byte, \( L/(\ln 2 \cdot \beta) \), is tokenizer-independent and is what careful comparisons use. The reference model here finished at validation loss 2.8401 nats per token, which is perplexity 17.12 and, at the measured 3.611 bytes per token on the validation split, 1.135 bits per byte. For calibration, a strong character-level model on English Wikipedia reaches about 1.0 bits per byte and the best large models are well under 0.7.

Benchmark suites during training are used sparingly because they are noisy at small scale and expensive. The standard set is MMLU for knowledge, HellaSwag and ARC and Winogrande for commonsense, GSM8K and MATH for mathematics, HumanEval and MBPP (Mostly Basic Python Problems) for code. Below a few billion parameters most of these sit at chance, which is why loss remains the signal that matters early. The evaluation harness from EleutherAI is the de facto standard implementation, and small differences in prompt formatting and normalization change reported numbers by several points, which is why cross-paper comparisons of benchmark scores should be treated with suspicion unless the harness and settings match.

Contamination is the standing threat. Benchmarks are on the web, and the training corpus is the web. Even with \( n \)-gram decontamination, paraphrased and translated versions survive, and the effect is large. Models score far better on benchmark items that appear in their training data. The practical defences are to hold out a private evaluation set, to prefer benchmarks released after the training-data cutoff, and to watch for the characteristic signature of contamination, which is a model that scores well on a benchmark and poorly on a trivially rephrased version of it.

Loss curves remain the honest signal. They are continuous, they are computed on data the model has never seen, they are cheap enough to compute every few hundred steps, and they are not gameable. The discipline is to watch training loss for optimization health (spikes, plateaus, divergence) and validation loss on a fixed held-out set for generalization, and to distrust any benchmark movement that is not accompanied by a movement in validation loss.

Post-training, briefly

A pretrained model completes text. It does not follow instructions, refuse harmful requests, or format answers usefully. Post-training is what turns one into the other, and it is a large enough subject that it lives elsewhere on this site. The standard sequence is supervised fine-tuning on demonstration data, then preference optimization, either with a learned reward model and PPO as in the InstructGPT recipe from OpenAI or with a direct method such as DPO from Rafailov and colleagues at Stanford, which reparameterizes the RLHF objective so that the optimal policy can be fit by a simple classification loss with no reward model and no sampling. More recent pipelines add verifiable-reward RL for mathematics and code, where a checker replaces the reward model, and this is the mechanism behind the reasoning-model line from OpenAI, DeepSeek, and Qwen. The policy-gradient machinery, the PPO clipped objective, and the KL-regularized formulation are derived at deep reinforcement learning. The practical fine-tuning and adaptation techniques, including LoRA and quantized adapters, are covered at applied generative AI.

Inference

Prefill and decode are different machines

Generation has two phases with completely different performance characteristics. Prefill processes the entire prompt in one forward pass, where \( T \) tokens go in, all of them are computed in parallel, and the keys and values are written to the cache. Decode produces one token at a time. Each step processes a single token per sequence and reads the entire cache and the entire weight matrix to do it.

The arithmetic intensity, FLOPs per byte of memory traffic, tells the whole story. In prefill with \( B \) sequences of length \( T \), the work is \( 2N B T \) FLOPs and the traffic is \( 2N \) bytes for the weights (bf16) plus the activations, so the intensity is roughly \( BT \) FLOPs per byte. In decode, the work is \( 2NB \) FLOPs and the traffic is still \( 2N \) bytes for the weights plus the cache, so the intensity is roughly \( B \) FLOPs per byte. Prefill is compute-bound, while decode is memory-bandwidth-bound until the batch is very large.

The crossover is where intensity equals the machine's ridge point, which for this H100 is \( 729\times10^{12}\ \text{FLOP/s} / 2.93\times10^{12}\ \text{B/s} = 249 \) FLOPs per byte. So decoding becomes compute-bound only around batch 249, and below that every decode step takes the time required to stream the weights from HBM regardless of how few tokens are being produced. For an 8B model in bf16 that floor is \( 1.6\times10^{10} \) bytes / \( 2.93\times10^{12} \) B/s \( = 5.5 \) ms per token, or 182 tokens per second, no matter how small the batch. Everything in the serving stack exists to raise the batch size toward that ridge point.

Measured here on the trained 26.7M-parameter model, the results are below.

KV cache, paged attention, and continuous batching

The cache size formula was derived above, \( 4 L H_{kv} d_h \) bytes per token per sequence in bf16. Two systems ideas turn that formula into throughput.

Paged attention, from Kwon and colleagues at Berkeley in the vLLM paper of 2023, observes that the naive implementation allocates a contiguous cache buffer sized for the maximum possible sequence length for every request, and that this wastes most of the memory. Their measurements found 60% to 80% of the allocated KV cache unused, from internal fragmentation (the sequence finishes early), external fragmentation, and the inability to share memory between requests. The fix is borrowed wholesale from virtual memory. Divide the cache into fixed-size blocks of, say, 16 tokens, keep a per-sequence block table mapping logical positions to physical blocks, and allocate blocks on demand. Fragmentation drops to under 4%, and blocks can be shared between sequences with a common prefix, with copy-on-write when they diverge, which makes parallel sampling and beam search nearly free. vLLM reported 2 to 4× the throughput of the previous best systems at the same latency.

Continuous batching, from the Orca paper by Yu and colleagues in 2022 and now universal, changes the scheduling granularity from request to iteration. Static batching runs a batch until every sequence finishes, so a batch containing one long generation wastes the slots of every short one. Continuous batching evicts each sequence the moment it emits its stop token and admits a waiting request in its place at the next decode step. Combined with paged attention, which makes admitting a new sequence a matter of allocating a few blocks rather than a contiguous buffer, this keeps the batch near the hardware's ridge point continuously.

Chunked prefill and disaggregation are the current refinements. Because prefill is compute-bound and decode is bandwidth-bound, mixing them in one batch means each interferes with the other. A long prefill stalls every decode in flight. Chunked prefill splits a long prompt into pieces that are interleaved with decode steps, bounding the stall. Disaggregated serving goes further and runs prefill and decode on separate pools of GPUs, sized independently, shipping the KV cache between them.

Speculative decoding

Decode is bandwidth-bound, which means the GPU is idle most of the time. If several tokens could be verified in one pass, the bandwidth cost would be amortized. Speculative decoding, from Leviathan, Kalman, and Matias at Google and independently from Chen and colleagues at DeepMind in 2023, does exactly this. A small draft model proposes \( \gamma \) tokens autoregressively, and the large model verifies all \( \gamma \) in a single forward pass over the extended sequence.

The key property is that the output distribution is exactly the target model's. The acceptance rule is rejection sampling. For draft distribution \( q \) and target \( p \), accept the drafted token \( x \) with probability \( \min(1, p(x)/q(x)) \), and on rejection, sample from the residual distribution \( \propto \max(0, p(x) - q(x)) \). To see that this is correct, note that the probability of finally emitting \( x \) is

The first term is \( \min(q(x), p(x)) \). The rejection probability is \( \sum_{x'} \max(0, q(x')-p(x')) = \sum_{x'}\max(0,p(x')-q(x')) \) because both \( p \) and \( q \) sum to one, so the second term simplifies to \( \max(0, p(x)-q(x)) \), and the sum is \( \min(q,p) + \max(0, p-q) = p(x) \) in both cases. The draft model can therefore be arbitrarily bad without changing the output, only the speed.

Expected speedup. Let \( \alpha \) be the average per-token acceptance probability. With \( \gamma \) drafted tokens, the number accepted before the first rejection is geometric, and including the one token the target always contributes, the expected tokens per target call is

If the draft costs a fraction \( c \) of the target per call, the wall-clock speedup is

At \( \alpha = 0.8 \), \( \gamma = 4 \), \( c = 0.1 \), this is \( (1-0.8^5)/(0.2 \times 1.4) = 0.672/0.28 = 2.40\times \). At \( \alpha = 0.5 \) the same configuration gives \( (1-0.03125)/(0.5\times1.4) = 1.38\times \). Differentiating with respect to \( \gamma \) shows the optimal draft length grows as \( \alpha \to 1 \) and collapses to 1 when \( \alpha \) is small, which is why production systems tune \( \gamma \) online.

This was measured here with two models trained on the same tokenizer and data, the 26.7M-parameter target and a 4.92M-parameter draft (validation loss 2.840 and 3.372 respectively).

Three variants matter in production. Self-speculation attaches extra prediction heads to the target model itself, as in Medusa from Cai and colleagues and in EAGLE, removing the need for a separate draft. Prompt-lookup or n-gram drafting copies candidate continuations from the prompt, which works well for summarization and code editing where the output repeats the input. And tree-structured drafting verifies several candidate continuations in one target call, raising the effective acceptance rate at the cost of a larger verification batch.

Quantization for serving

Since decode is bandwidth-bound and the dominant traffic is the weights, shrinking the weights directly buys latency. Weight-only quantization to 4 bits reduces the traffic by 4× against bf16 and, because the matmul is dequantized on the fly into the tensor cores, costs almost nothing in compute.

GPTQ, from Frantar, Ashkboos, Hoefler, and Alistarh at ETH Zurich and IST Austria in 2022, quantizes one layer at a time by solving a local reconstruction problem. Choose quantized weights \( \hat W \) minimizing \( \|WX - \hat W X\|_2^2 \) over a small calibration set. The solution uses the Optimal Brain Surgeon framework. Quantize weights one column at a time, and after each one, update the remaining unquantized weights to compensate for the error introduced, using the inverse Hessian \( (2XX\T)^{-1} \) of the layer's reconstruction objective. GPTQ's contribution is making this tractable, by processing columns in a fixed order with a Cholesky reformulation, bringing a 175B model to 4 bits in about four GPU hours.

AWQ, from Lin and colleagues at MIT in 2023, starts from a different observation, that not all weights matter equally, and that the ones that matter are identified by the magnitude of the activations they multiply, not by the magnitude of the weights themselves. Protecting the roughly 1% of channels with the largest activation magnitude preserves almost all the quality. Rather than keeping those in high precision, which breaks the uniform kernel layout, AWQ scales them up before quantization and scales the corresponding activations down, an equivalent transformation that moves the salient channels into a range where the quantization grid is finer. AWQ needs no backpropagation and generalizes better to instruction-tuned models than GPTQ.

Measured here on the trained model, per-row symmetric int8 quantization of every 2-D weight matrix gives a mean relative reconstruction error of 0.78% (maximum 0.94%) and compresses the weights by 1.987× against bf16. The general picture from the literature is that 8-bit weight-only quantization is free, 4-bit costs a small and usually acceptable amount, 3-bit needs care, and 2-bit needs either a much better method or quantization-aware training. KV-cache quantization to 8 or 4 bits is the other lever and matters most at long context, where the cache exceeds the weights.

Worked problems

Implementation

Everything below runs. The PyTorch and JAX versions of each component were executed on the H100 in this repository and checked against each other and against reference implementations. The tokenizer core is given in Python and in Rust because the merge loop is the one part where the language choice changes what is feasible.

The BPE trainer and encoder

The trainer keeps the corpus as a bag of (pre-token, frequency) pairs, a counter of adjacent symbol pairs, and an inverted index from pair to the pre-tokens containing it, so a merge touches only the affected words rather than rescanning the corpus. The encoder applies learned merges in rank order. The Rust version is the same algorithm and is what a production tokenizer looks like inside. Running it on the canonical example aaabdaaabac produces the merges aa, ab, aaab and encodes the corpus to five symbols, which is the textbook answer.

import collections, regex as re

# GPT-2 pre-tokenization: contractions, letter runs with an optional leading
# space, digit runs, punctuation runs, then whitespace. Merges never cross
# a piece boundary, which is what stops "dog." and "dog," from being tokens.
GPT2_PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""


class BPE:
    def __init__(self, pattern=GPT2_PAT):
        self.split = re.compile(pattern).findall
        self.merges = {}                              # (int, int) -> new id
        self.vocab = {i: bytes([i]) for i in range(256)}

    def train(self, text, vocab_size):
        counts = collections.Counter()
        for chunk in self.split(text):
            counts[chunk.encode("utf-8")] += 1
        words = [list(w) for w in counts]             # list[list[int]]
        freqs = list(counts.values())

        pair_counts = collections.Counter()           # (a,b) -> weighted count
        where = collections.defaultdict(set)          # (a,b) -> {word index}
        for i, w in enumerate(words):
            for a, b in zip(w, w[1:]):
                pair_counts[(a, b)] += freqs[i]
                where[(a, b)].add(i)

        for m in range(vocab_size - 256):
            if not pair_counts:
                break
            pair, cnt = max(pair_counts.items(), key=lambda kv: kv[1])
            if cnt < 2:
                break
            new_id = 256 + m
            self.merges[pair] = new_id
            self.vocab[new_id] = self.vocab[pair[0]] + self.vocab[pair[1]]
            a, b = pair
            for i in list(where[pair]):               # only the affected words
                w, f = words[i], freqs[i]
                for x, y in zip(w, w[1:]):            # retract old pairs
                    pair_counts[(x, y)] -= f
                    if pair_counts[(x, y)] <= 0:
                        del pair_counts[(x, y)]
                        where.pop((x, y), None)
                    else:
                        where[(x, y)].discard(i)
                nw, j = [], 0
                while j < len(w):
                    if j < len(w) - 1 and w[j] == a and w[j + 1] == b:
                        nw.append(new_id); j += 2
                    else:
                        nw.append(w[j]); j += 1
                words[i] = nw
                for x, y in zip(nw, nw[1:]):          # add new pairs
                    pair_counts[(x, y)] += f
                    where[(x, y)].add(i)
            where.pop(pair, None); pair_counts.pop(pair, None)
        return self

    def _encode_word(self, bs):
        ids = list(bs)
        while len(ids) >= 2:
            best, bi = None, None
            for k in range(len(ids) - 1):             # lowest learned rank wins
                r = self.merges.get((ids[k], ids[k + 1]))
                if r is not None and (best is None or r < best):
                    best, bi = r, k
            if best is None:
                break
            ids[bi:bi + 2] = [best]
        return ids

    def encode(self, text):
        out = []
        for chunk in self.split(text):
            out.extend(self._encode_word(chunk.encode("utf-8")))
        return out

    def decode(self, ids):
        return b"".join(self.vocab[i] for i in ids).decode("utf-8", errors="replace")


# vocab 8192 on 8.1 MB of English prose: 39 s, 3.430 bytes/token held out.

use std::collections::HashMap;

/// Merge-training inner loop. `words` is the corpus as byte-id sequences,
/// `freqs` their counts. An inverted index pair -> word ids keeps each merge
/// proportional to the number of affected words, not the corpus size.
pub fn train(words: &mut Vec<Vec<u32>>, freqs: &[i64], n_merges: usize)
    -> Vec<((u32, u32), u32)> {
    let mut pair_counts: HashMap<(u32, u32), i64> = HashMap::new();
    let mut where_: HashMap<(u32, u32), Vec<usize>> = HashMap::new();

    for (i, w) in words.iter().enumerate() {
        for p in w.windows(2) {
            let key = (p[0], p[1]);
            *pair_counts.entry(key).or_insert(0) += freqs[i];
            let e = where_.entry(key).or_insert_with(Vec::new);
            if e.last() != Some(&i) { e.push(i); }
        }
    }

    let mut merges = Vec::with_capacity(n_merges);
    for m in 0..n_merges {
        let best = pair_counts.iter().max_by_key(|(k, v)| {
            (**v, std::cmp::Reverse(k.0), std::cmp::Reverse(k.1))
        });
        let (&pair, &cnt) = match best { Some((k, v)) => (k, v), None => break };
        if cnt < 2 { break; }
        let new_id = 256u32 + m as u32;
        merges.push((pair, new_id));

        let touched: Vec<usize> = where_.remove(&pair).unwrap_or_default();
        for i in touched {
            let f = freqs[i];
            for p in words[i].windows(2) {            // retract
                let k = (p[0], p[1]);
                if let Some(c) = pair_counts.get_mut(&k) {
                    *c -= f;
                    if *c <= 0 { pair_counts.remove(&k); where_.remove(&k); }
                }
            }
            let old = std::mem::take(&mut words[i]);
            let mut new_w: Vec<u32> = Vec::with_capacity(old.len());
            let mut j = 0usize;
            while j < old.len() {
                if j + 1 < old.len() && old[j] == pair.0 && old[j + 1] == pair.1 {
                    new_w.push(new_id); j += 2;
                } else { new_w.push(old[j]); j += 1; }
            }
            for p in new_w.windows(2) {               // re-insert
                let k = (p[0], p[1]);
                *pair_counts.entry(k).or_insert(0) += f;
                let e = where_.entry(k).or_insert_with(Vec::new);
                if e.last() != Some(&i) { e.push(i); }
            }
            words[i] = new_w;
        }
        pair_counts.remove(&pair);
    }
    merges
}

/// Encode one pre-token by repeatedly applying the lowest-rank learned merge.
pub fn encode_word(bytes: &[u8], rank: &HashMap<(u32, u32), u32>) -> Vec<u32> {
    let mut ids: Vec<u32> = bytes.iter().map(|&b| b as u32).collect();
    loop {
        let mut best: Option<(u32, usize)> = None;
        for j in 0..ids.len().saturating_sub(1) {
            if let Some(&r) = rank.get(&(ids[j], ids[j + 1])) {
                if best.map_or(true, |(br, _)| r < br) { best = Some((r, j)); }
            }
        }
        match best {
            None => break,
            Some((_, j)) => {
                let new_id = 256 + rank[&(ids[j], ids[j + 1])];
                ids.splice(j..j + 2, [new_id]);
            }
        }
    }
    ids
}

// train on b"aaabdaaabac" for 3 merges:
//   merges  [((97,97),256), ((97,98),257), ((256,257),258)]
//   corpus  [258, 100, 258, 97, 99]

RoPE, with shapes

The tables are built once at model construction and sliced per forward. Every tensor's shape is annotated because the interleaving convention (pairs are adjacent coordinates here, but some implementations split the head in half instead) is the single most common source of silent incompatibility between checkpoints.

import torch

def rope_tables(seq_len, d_head, theta=10000.0, device="cuda"):
    """Returns cos, sin of shape (seq_len, d_head // 2).
    Angles are computed in float64 and cast down: at position 100k the
    float32 evaluation of m * theta_i loses several digits."""
    inv = 1.0 / (theta ** (torch.arange(0, d_head, 2, device=device).double() / d_head))
    t = torch.arange(seq_len, device=device).double()          # (T,)
    ang = torch.outer(t, inv)                                  # (T, Dh/2)
    return ang.cos().float(), ang.sin().float()

def apply_rope(x, cos, sin):
    """x: (B, H, T, Dh) -> (B, H, T, Dh).  cos, sin: (T, Dh/2).
    Pair i is coordinates (2i, 2i+1); each pair is rotated by m * theta_i."""
    x1, x2 = x[..., 0::2], x[..., 1::2]                         # (B,H,T,Dh/2)
    c = cos[None, None, :, :]                                  # (1,1,T,Dh/2)
    s = sin[None, None, :, :]
    o1 = x1 * c - x2 * s                                       # (B,H,T,Dh/2)
    o2 = x1 * s + x2 * c
    return torch.stack((o1, o2), dim=-1).flatten(-2)           # (B,H,T,Dh)

# Verification of the relative-position property: the same offset gives the
# same dot product no matter where in the sequence the pair sits.
cos, sin = rope_tables(256, 64, device="cpu")
q = torch.randn(1, 1, 1, 64, dtype=torch.float32)
k = torch.randn(1, 1, 1, 64, dtype=torch.float32)
for (m, n) in [(5, 2), (17, 14), (100, 97), (203, 200)]:
    qm = apply_rope(q, cos[m:m+1], sin[m:m+1])
    kn = apply_rope(k, cos[n:n+1], sin[n:n+1])
    print(m, n, float((qm * kn).sum()))
# all four print the same value to ~8 significant figures

import jax, jax.numpy as jnp

def rope_tables(seq_len, d_head, theta=10000.0):
    """cos, sin: (seq_len, d_head // 2)"""
    inv = 1.0 / (theta ** (jnp.arange(0, d_head, 2, dtype=jnp.float32) / d_head))
    t = jnp.arange(seq_len, dtype=jnp.float32)                 # (T,)
    ang = jnp.outer(t, inv)                                    # (T, Dh/2)
    return jnp.cos(ang), jnp.sin(ang)

def apply_rope(x, cos, sin):
    """x: (B, H, T, Dh) -> (B, H, T, Dh); cos, sin: (T, Dh/2)"""
    x1, x2 = x[..., 0::2], x[..., 1::2]                        # (B,H,T,Dh/2)
    c = cos[None, None, :, :]
    s = sin[None, None, :, :]
    o1 = x1 * c - x2 * s
    o2 = x1 * s + x2 * c
    return jnp.stack([o1, o2], axis=-1).reshape(x.shape)       # (B,H,T,Dh)

cos, sin = rope_tables(256, 64)
q = jax.random.normal(jax.random.PRNGKey(0), (1, 1, 1, 64))
k = jax.random.normal(jax.random.PRNGKey(1), (1, 1, 1, 64))
vals = []
for (m, n) in [(5, 2), (17, 14), (100, 97), (203, 200)]:
    qm = apply_rope(q, cos[m:m+1], sin[m:m+1])
    kn = apply_rope(k, cos[n:n+1], sin[n:n+1])
    vals.append(float(jnp.sum(qm * kn)))
print(vals, "spread:", max(vals) - min(vals))
# spread 3.2e-06 in float32; exactly zero in exact arithmetic

Grouped-query attention

The only difference from multi-head attention is that the key and value projections produce \( H_{kv} \) heads instead of \( H \), and each is repeated \( H/H_{kv} \) times before the score computation. The repeat is a view in the fused kernels and a materialization in the naive path. Either way the cache holds only the \( H_{kv} \) copies, which is the entire point.

import torch, torch.nn as nn, torch.nn.functional as F

class GQA(nn.Module):
    def __init__(self, d_model, n_head, n_kv_head):
        super().__init__()
        assert n_head % n_kv_head == 0
        self.H, self.Hkv = n_head, n_kv_head
        self.dh = d_model // n_head
        self.wq = nn.Linear(d_model, n_head * self.dh, bias=False)
        self.wk = nn.Linear(d_model, n_kv_head * self.dh, bias=False)
        self.wv = nn.Linear(d_model, n_kv_head * self.dh, bias=False)
        self.wo = nn.Linear(n_head * self.dh, d_model, bias=False)

    def forward(self, x, cos, sin, cache=None):
        B, T, _ = x.shape                                       # (B,T,d)
        q = self.wq(x).view(B, T, self.H, self.dh).transpose(1, 2)    # (B,H,T,dh)
        k = self.wk(x).view(B, T, self.Hkv, self.dh).transpose(1, 2)  # (B,Hkv,T,dh)
        v = self.wv(x).view(B, T, self.Hkv, self.dh).transpose(1, 2)  # (B,Hkv,T,dh)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        if cache is not None:                                   # decode path
            if cache[0] is not None:
                k = torch.cat([cache[0], k], dim=2)             # (B,Hkv,T_past+T,dh)
                v = torch.cat([cache[1], v], dim=2)
            cache = (k, v)
        if self.Hkv != self.H:                                  # expand for the matmul
            rep = self.H // self.Hkv
            k = k.repeat_interleave(rep, dim=1)                 # (B,H,Tk,dh)
            v = v.repeat_interleave(rep, dim=1)
        causal = T > 1
        y = F.scaled_dot_product_attention(q, k, v, is_causal=causal)  # (B,H,T,dh)
        y = y.transpose(1, 2).reshape(B, T, self.H * self.dh)
        return self.wo(y), cache

# KV cache bytes per token per sequence = 2 (K and V) * 2 (bf16) * L * Hkv * dh
# L=80, dh=128:  Hkv=64 -> 3.3 MB/token   Hkv=8 -> 0.41 MB/token

import math, jax, jax.numpy as jnp

def gqa(x, p, cos, sin, H, Hkv):
    """x: (B,T,d); p holds wq (d, H*dh), wk/wv (d, Hkv*dh), wo (H*dh, d)."""
    B, T, d = x.shape
    dh = d // H
    q = (x @ p["wq"]).reshape(B, T, H, dh).transpose(0, 2, 1, 3)     # (B,H,T,dh)
    k = (x @ p["wk"]).reshape(B, T, Hkv, dh).transpose(0, 2, 1, 3)   # (B,Hkv,T,dh)
    v = (x @ p["wv"]).reshape(B, T, Hkv, dh).transpose(0, 2, 1, 3)
    q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
    if Hkv != H:
        rep = H // Hkv
        k = jnp.repeat(k, rep, axis=1)                               # (B,H,T,dh)
        v = jnp.repeat(v, rep, axis=1)
    scores = jnp.einsum("bhqd,bhkd->bhqk", q, k) / math.sqrt(dh)     # (B,H,T,T)
    mask = jnp.tril(jnp.ones((T, T), dtype=bool))
    scores = jnp.where(mask, scores, -jnp.inf)
    a = jax.nn.softmax(scores, axis=-1)
    o = jnp.einsum("bhqk,bhkd->bhqd", a, v)                          # (B,H,T,dh)
    o = o.transpose(0, 2, 1, 3).reshape(B, T, H * dh)
    return o @ p["wo"]

# For production use jax.nn.dot_product_attention, which dispatches to the
# fused cuDNN kernel and never materializes the (B,H,T,T) score matrix.

The complete model

Under 120 lines each, this covers embeddings, pre-norm blocks with RMSNorm, GQA with RoPE, a SwiGLU MLP, a tied output head, and the depth-scaled residual initialization. This is the model that was trained for the measurements on this page.

import math, torch, torch.nn as nn, torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class Config:
    vocab_size: int = 8192
    d_model: int = 512
    n_layer: int = 8
    n_head: int = 8
    n_kv_head: int = 2
    d_ff: int = 1408          # ~ (8/3) * d_model, rounded to a multiple of 128
    seq_len: int = 512
    rope_theta: float = 10000.0

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps
    def forward(self, x):                                  # (B,T,d)
        dt = x.dtype
        x = x.float()                                      # reduce in fp32
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x.to(dt) * self.weight

class SwiGLU(nn.Module):
    def __init__(self, d, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d, d_ff, bias=False)           # gate
        self.w3 = nn.Linear(d, d_ff, bias=False)           # up
        self.w2 = nn.Linear(d_ff, d, bias=False)           # down
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n1, self.n2 = RMSNorm(cfg.d_model), RMSNorm(cfg.d_model)
        self.attn = GQA(cfg.d_model, cfg.n_head, cfg.n_kv_head)
        self.mlp = SwiGLU(cfg.d_model, cfg.d_ff)
    def forward(self, x, cos, sin):                        # pre-norm residual
        x = x + self.attn(self.n1(x), cos, sin)[0]
        return x + self.mlp(self.n2(x))

class GPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.norm = RMSNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.emb.weight                 # tied
        cos, sin = rope_tables(cfg.seq_len * 4, cfg.d_model // cfg.n_head,
                               cfg.rope_theta, device="cpu")
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
        self.apply(self._init)
        for n, p in self.named_parameters():               # depth-scaled residual init
            if n.endswith("wo.weight") or n.endswith("w2.weight"):
                nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * cfg.n_layer))

    def _init(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 0.0, 0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.emb(idx)                                  # (B,T,d)
        cos, sin = self.cos[:T], self.sin[:T]
        for b in self.blocks:
            x = b(x, cos, sin)
        logits = self.head(self.norm(x))                   # (B,T,V)
        if targets is None:
            return logits
        return F.cross_entropy(logits.view(-1, logits.size(-1)).float(),
                               targets.reshape(-1))

# 8192 x 512 tied embedding + 8 layers -> 26,747,392 parameters, matching
# Vd + L(2 d^2 + 2 d Hkv dh + 3 d F + 2d) + d exactly.

import math, jax, jax.numpy as jnp
from jax import random, lax

def init_params(key, cfg):
    V, d, L, H, Hkv, F = (cfg["vocab"], cfg["d"], cfg["L"], cfg["H"],
                          cfg["Hkv"], cfg["F"])
    dh = d // H
    ks = random.split(key, 2 + 7 * L)
    p = {"emb": random.normal(ks[0], (V, d)) * 0.02,
         "final_norm": jnp.ones((d,)), "layers": []}
    resid = 0.02 / math.sqrt(2 * L)                        # depth-scaled
    for i in range(L):
        k = ks[2 + 7 * i: 2 + 7 * (i + 1)]
        p["layers"].append({
            "n1": jnp.ones((d,)), "n2": jnp.ones((d,)),
            "wq": random.normal(k[0], (d, H * dh)) * 0.02,
            "wk": random.normal(k[1], (d, Hkv * dh)) * 0.02,
            "wv": random.normal(k[2], (d, Hkv * dh)) * 0.02,
            "wo": random.normal(k[3], (H * dh, d)) * resid,
            "w1": random.normal(k[4], (d, F)) * 0.02,
            "w3": random.normal(k[5], (d, F)) * 0.02,
            "w2": random.normal(k[6], (F, d)) * resid})
    return p

def rms_norm(x, w, eps=1e-6):
    return x * lax.rsqrt(jnp.mean(x * x, axis=-1, keepdims=True) + eps) * w

def swiglu(x, lp):
    return (jax.nn.silu(x @ lp["w1"]) * (x @ lp["w3"])) @ lp["w2"]

def forward(p, idx, cfg):
    B, T = idx.shape
    cos, sin = rope_tables(T, cfg["d"] // cfg["H"])
    x = p["emb"][idx]                                       # (B,T,d)
    for lp in p["layers"]:
        x = x + gqa(rms_norm(x, lp["n1"]), lp, cos, sin, cfg["H"], cfg["Hkv"])
        x = x + swiglu(rms_norm(x, lp["n2"]), lp)
    x = rms_norm(x, p["final_norm"])
    return x @ p["emb"].T                                   # tied head, (B,T,V)

def loss_fn(p, idx, targets, cfg):
    logits = forward(p, idx, cfg)
    logp = jax.nn.log_softmax(logits.astype(jnp.float32), axis=-1)
    tgt = jax.nn.one_hot(targets, logits.shape[-1])
    return -jnp.mean(jnp.sum(logp * tgt, axis=-1))

# Sanity check every model should pass: at initialization the loss must be
# close to ln(V). With V = 256 this printed 5.542 against ln(256) = 5.545.

The MoE layer with the load-balancing loss

The routing, the renormalized top-\( k \) weights, the auxiliary loss, and the gather-scatter that sends each token only to its chosen experts. The loop over experts is what a single-device implementation looks like. In a distributed one the loop becomes an all-to-all and the fixed expert capacity introduces token dropping.

class MoE(nn.Module):
    """Top-k routed experts with the Switch/GShard load-balancing loss."""
    def __init__(self, d, d_ff, n_expert, top_k=2, n_shared=0):
        super().__init__()
        self.E, self.k = n_expert, top_k
        self.gate = nn.Linear(d, n_expert, bias=False)
        self.experts = nn.ModuleList([SwiGLU(d, d_ff) for _ in range(n_expert)])
        self.shared = SwiGLU(d, d_ff) if n_shared else None
        self.aux = torch.zeros(())

    def forward(self, x):                                  # (B,T,d)
        B, T, d = x.shape
        xf = x.reshape(-1, d)                              # (N,d), N = B*T
        probs = self.gate(xf).float().softmax(-1)          # (N,E)
        topw, topi = probs.topk(self.k, dim=-1)            # (N,k) each
        topw = topw / topw.sum(-1, keepdim=True)           # renormalize

        # load-balancing loss: E * sum_e f_e * P_e, minimized at uniform routing
        counts = F.one_hot(topi, self.E).float().sum(1)    # (N,E) 0/1 per expert
        f = counts.mean(0)                                 # fraction routed to e
        P = probs.mean(0)                                  # mean router prob for e
        self.aux = self.E * (f * P).sum()

        out = torch.zeros_like(xf)
        for e, expert in enumerate(self.experts):
            idx, slot = (topi == e).nonzero(as_tuple=True) # tokens choosing e
            if idx.numel() == 0:
                continue
            w = topw[idx, slot].unsqueeze(-1).to(x.dtype)  # (n_e,1)
            out.index_add_(0, idx, expert(xf[idx]) * w)    # scatter-add
        if self.shared is not None:
            out = out + self.shared(xf)                    # every token, always
        return out.view(B, T, d)

# total loss = cross_entropy + alpha * mean_over_layers(aux), alpha ~ 1e-2.
# With alpha = 0 the router collapses; with alpha = 1e-2 the max/min expert
# load ratio stays near 1 in the runs on this page.

def moe_layer(x, gate_w, experts, k=2, aux_coef=0.01):
    """x: (B,T,d); gate_w: (d,E); experts: list of dicts with w1, w3, w2."""
    B, T, d = x.shape
    E = gate_w.shape[1]
    xf = x.reshape(-1, d)                                       # (N,d)
    probs = jax.nn.softmax((xf @ gate_w).astype(jnp.float32), axis=-1)   # (N,E)
    topv, topi = lax.top_k(probs, k)                            # (N,k)
    topv = topv / jnp.sum(topv, axis=-1, keepdims=True)

    onehot = jnp.sum(jax.nn.one_hot(topi, E), axis=1)           # (N,E)
    f = jnp.mean(onehot, axis=0)                                # (E,)
    P = jnp.mean(probs, axis=0)                                 # (E,)
    aux = E * jnp.sum(f * P)

    out = jnp.zeros_like(xf)
    for e in range(E):                                          # unrolled at trace time
        w = jnp.sum(jnp.where(topi == e, topv, 0.0), axis=-1)[:, None]
        ex = experts[e]
        y = (jax.nn.silu(xf @ ex["w1"]) * (xf @ ex["w3"])) @ ex["w2"]
        out = out + w * y                                       # dense-compute form
    return out.reshape(B, T, d), aux_coef * aux

# The dense form above computes every expert for every token: correct, and
# the right thing under jit for small E. A real implementation sorts tokens by
# expert with jnp.argsort, uses a fixed capacity per expert so the shapes are
# static, and replaces the loop with a shard_map all-to-all.

The training loop

Gradient accumulation, mixed precision, global-norm clipping, the warmup-cosine schedule, activation checkpointing, and checkpoint save and restore. The accumulation is exact. Summing the gradients of \( a \) microbatches each divided by \( a \) equals the gradient of the full batch, verified here to a relative error of \( 6.4\times10^{-16} \) in float64.

import math, time, torch
from torch.utils.checkpoint import checkpoint

def lr_at(step, total, base, warmup=0.02, kind="cosine", min_frac=0.1, decay=0.2):
    w = max(1, int(warmup * total))
    if step < w:
        return base * (step + 1) / w                       # linear warmup
    if kind == "cosine":
        p = (step - w) / max(1, total - w)
        return base * (min_frac + (1 - min_frac) * 0.5 * (1 + math.cos(math.pi * p)))
    if kind == "wsd":                                      # stable, then linear decay
        d = int(decay * total)
        if step < total - d:
            return base
        p = (step - (total - d)) / max(1, d)
        return base * ((1 - p) * (1 - min_frac) + min_frac)
    raise ValueError(kind)

def make_optimizer(model, lr, wd=0.1, betas=(0.9, 0.95)):
    # no weight decay on 1-D parameters: decaying a norm gain toward zero
    # attacks the scale invariance the norm exists to provide
    decay = [p for p in model.parameters() if p.dim() >= 2]
    nodecay = [p for p in model.parameters() if p.dim() < 2]
    return torch.optim.AdamW(
        [{"params": decay, "weight_decay": wd},
         {"params": nodecay, "weight_decay": 0.0}],
        lr=lr, betas=betas, eps=1e-8, fused=True)

def enable_checkpointing(model):
    for b in model.blocks:                                 # recompute each block
        f = b.forward
        b.forward = (lambda f: lambda *a, **k:
                     checkpoint(f, *a, use_reentrant=False, **k))(f)

def train(model, get_batch, steps, lr=3e-3, accum=1, clip=1.0, ckpt_every=1000):
    model = torch.compile(model)
    opt = make_optimizer(model, lr)
    for step in range(steps):
        for g in opt.param_groups:
            g["lr"] = lr_at(step, steps, lr)
        opt.zero_grad(set_to_none=True)
        total = 0.0
        for _ in range(accum):                             # exact large-batch gradient
            x, y = get_batch()                             # (B,T), (B,T)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss = model(x, y)
            (loss / accum).backward()
            total += loss.item() / accum
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        opt.step()
        if (step + 1) % ckpt_every == 0:
            torch.save({"step": step, "model": model.state_dict(),
                        "opt": opt.state_dict()}, f"ckpt_{step+1}.pt")
    return model

# Measured on one H100 80GB: 26.7M parameters, batch 64 x 512, bf16 + compile
#   1,232,919 tokens/s over 16,384 steps  ->  31.4% MFU against 729 bf16 TFLOP/s

import jax, jax.numpy as jnp
from jax import lax

def adamw_init(p):
    return {"m": jax.tree.map(jnp.zeros_like, p),
            "v": jax.tree.map(jnp.zeros_like, p), "t": 0}

def adamw_update(grads, state, params, lr, b1=0.9, b2=0.95, eps=1e-8, wd=0.1):
    t = state["t"] + 1
    m = jax.tree.map(lambda m_, g: b1 * m_ + (1 - b1) * g, state["m"], grads)
    v = jax.tree.map(lambda v_, g: b2 * v_ + (1 - b2) * g * g, state["v"], grads)
    mh = jax.tree.map(lambda x: x / (1 - b1 ** t), m)      # exact bias correction
    vh = jax.tree.map(lambda x: x / (1 - b2 ** t), v)
    upd = jax.tree.map(lambda mm, vv, pp: -lr * (mm / (jnp.sqrt(vv) + eps) + wd * pp),
                       mh, vh, params)                     # decoupled decay
    return upd, {"m": m, "v": v, "t": t}

def make_train_step(cfg, accum, opt_update):
    grad_fn = jax.value_and_grad(loss_fn)

    @jax.jit
    def step(p, opt_state, xs, ys, lr):
        """xs, ys: (accum, B, T). One scan iteration per microbatch, so the
        peak activation memory is that of a single microbatch."""
        def body(carry, batch):
            g_acc, l_acc = carry
            xb, yb = batch
            l, g = grad_fn(p, xb, yb, cfg)
            g_acc = jax.tree.map(lambda a, b: a + b / accum, g_acc, g)
            return (g_acc, l_acc + l / accum), None

        zero = jax.tree.map(jnp.zeros_like, p)
        (grads, loss), _ = lax.scan(body, (zero, 0.0), (xs, ys))
        gnorm = jnp.sqrt(sum(jnp.sum(g ** 2) for g in jax.tree.leaves(grads)))
        grads = jax.tree.map(lambda g: g * jnp.minimum(1.0, 1.0 / (gnorm + 1e-6)),
                             grads)                        # global-norm clip
        updates, opt_state = opt_update(grads, opt_state, p, lr)
        p = jax.tree.map(lambda a, b: a + b, p, updates)
        return p, opt_state, loss, gnorm
    return step

# For activation checkpointing wrap the per-layer function in jax.checkpoint
# (alias jax.remat); for data parallelism use jax.sharding.NamedSharding with
# a "data" mesh axis and let the compiler insert the all-reduce.

How it is done in practice: a complete run, measured

Everything above was exercised by actually building the thing. Here is the pipeline, end to end, on one H100 80GB.

  1. corpus       1 GB Wikipedia XML dump + 50 Project Gutenberg books
  2. extraction   markup stripped, templates and tables removed
  3. filtering    min 400 bytes, >=90% alphanumeric, >=4 common stop words,
                  MD5 near-duplicate removal on the first 60 words
                  -> 134,666 documents kept, 557 MB, 65.2% of bytes survive
  4. tokenizer    byte-level BPE, GPT-2 regex pre-tokenization, V = 8192
                  -> 3.416 bytes/token on the training corpus
  5. encoding     163.1M training tokens + 10.4M book tokens + 834K val tokens
  6. model        d=512, L=8, H=8, Hkv=2, F=1408, RoPE, RMSNorm, SwiGLU, tied
                  -> 26,747,392 parameters (22,553,088 non-embedding)
  7. training     AdamW, lr 3e-3 cosine to 10%, 2% warmup, wd 0.1, clip 1.0,
                  batch 64 x 512 = 32,768 tokens/step, bf16 + torch.compile
                  -> 16,384 steps = 536,870,912 tokens = 20.1 tokens/parameter
  8. result       435.2 s wall clock, 1,232,919 tokens/s, 31.4% MFU
                  val loss 2.8401 nats/token = ppl 17.12 = 1.135 bits/byte
          

The full loss curve below is sampled every 500 steps.

The shape is the familiar one, a fast drop in the first 5% of steps as the model learns the unigram distribution and the tokenizer's statistics, then a long power-law-ish descent, then a visible extra drop over the last 15% as the cosine schedule decays the learning rate. The gap between train and validation loss is small throughout (the final values are 2.38 and 2.84, and the training loss is a noisy single-batch estimate rather than an average), which at 3.1 epochs over the corpus is what the multi-epoch measurements predict.

What the run cost. 435 seconds on one GPU. Scaling the same recipe to Llama-3-8B's 15T tokens at 8B parameters would be \( 6\times8\times10^{9}\times1.5\times10^{13} = 7.2\times10^{23} \) FLOPs, which at 40% of 729 TFLOP/s is \( 2.47\times10^{9} \) GPU-seconds, or 78 GPU-years, or 28 days on 1,024 GPUs. The gap between this page's run and a frontier run is six orders of magnitude in compute and exactly zero in the algorithms.

Seed-to-seed noise. Every ablation on this page reports a validation loss, and the honest question is how much of any difference is signal. Running the identical 65.5M-token configuration at four different data-sampling seeds gives the spread reported below. Differences smaller than that should be treated as noise, and several of the architecture ablations are.

The engineering gap between the derivation and a real run

Four things separate the code in this page's implementation section from what a production trainer looks like, and none of them are algorithmic.

Data loading. A frontier run streams tokens from object storage at several GB/s, shuffles at document granularity with a fixed seed so the run is reproducible and resumable, packs documents into fixed-length sequences with separator tokens, and maintains an exact resumable position so a restart consumes each token exactly once. Getting this wrong is the most common source of a run that silently trains on the same data repeatedly.

Checkpointing and failure. At 16,384 GPUs the mean time between hardware failures is measured in hours. The Llama 3 report documents 419 unexpected interruptions over 54 days, an average of one every three hours, with GPU and HBM failures the largest categories. Checkpoints must therefore be frequent, which means they must be fast, which means they must be sharded and written asynchronously. The state to save is not just the weights. Optimizer moments, the data loader position, the RNG state, and the learning-rate schedule position all have to round-trip or the restart is not a restart.

Numerical monitoring. Production runs log gradient norms per parameter group, the fraction of steps clipped, the attention entropy per layer, activation and logit magnitudes, and the optimizer's second-moment scale, all every step, because the signature of an impending divergence is visible in these tens to hundreds of steps before the loss moves.

The current research frontier

Architecture is consolidating, not diverging. The 2024 to 2026 releases from Meta, Alibaba's Qwen team, DeepSeek, Mistral, Google, and Allen AI differ in almost no architectural respect. The shared recipe is pre-norm RMSNorm, RoPE with some context-extension scheme, GQA, SwiGLU, and either a dense MLP or a fine-grained MoE. The visible disagreements are the MoE granularity (Mixtral's 8 experts against DeepSeek's 256), whether to use a shared expert, whether to use QK-norm, and how attention is compressed for the cache.

Attention compression beyond GQA. DeepSeek-V2 introduced multi-head latent attention, which projects the keys and values into a low-rank latent of a few hundred dimensions, caches only the latent, and reconstructs per head on the fly. The cache shrinks by roughly an order of magnitude against GQA at reportedly better quality, at the cost of a more complicated kernel and an interaction with RoPE that requires splitting the head into a rotary part and a non-rotary part. Whether this becomes standard is one of the open questions of the moment.

Sub-quadratic sequence models. Mamba and Mamba-2, from Gu at CMU and Dao at Princeton, are selective state-space models with linear-time recurrence and constant-size state, and they match transformers at moderate scale on language while being dramatically cheaper at long context. The current consensus is hybrid. Jamba from AI21, Zamba, and several Nvidia models interleave a small number of attention layers among many Mamba layers, on the evidence that attention is needed for exact retrieval and copying (the induction-head function) while the rest of the work is done fine by a recurrence. Related lines include linear attention with gating (RWKV from the open-source community, RetNet from Microsoft, and gated linear attention from MIT), all of which trade an exact \( O(T^2) \) mechanism for an \( O(T) \) approximation with a fixed-size state.

Data is the contested resource. The high-quality public web is finite, and Villalobos and colleagues at Epoch AI estimated the stock of human-generated public text would be exhausted by frontier training runs somewhere in the late 2020s. The three responses under active development are better filtering (FineWeb-Edu's classifier-based selection, the phi line from Microsoft arguing that textbook-quality synthetic data beats far more web text), synthetic data generation with verification, and multi-epoch training with the repetition penalties Muennighoff and colleagues quantified. Whether models trained substantially on model-generated text degrade, the "model collapse" question raised by Shumailov and colleagues at Oxford and Cambridge, remains genuinely open, with the practical answer so far being that verified synthetic data (mathematics with checkable answers, code with tests) works and unverified imitation does not.

Optimizers are moving for the first time in a decade. Muon's use in Kimi K2 at trillion-parameter scale is the first credible displacement of Adam in a frontier run, and the surrounding theory (steepest descent under non-Euclidean norms, the modular-norm framework from Bernstein and Newhouse) is the most interesting optimization work in the field right now. Distributed variants (DiLoCo from DeepMind, which synchronizes every few hundred steps instead of every step) attack the communication side of the same problem.

Test-time compute changes the scaling question. The reasoning-model line, from OpenAI's o-series through DeepSeek-R1 and the Qwen and Gemini reasoning variants, shifts compute from training to inference by generating long chains of thought. This creates a second scaling axis and reopens questions that Chinchilla appeared to settle. If a model will spend a thousand tokens of reasoning per query, the inference-cost term in the total-cost optimization grows by three orders of magnitude, which pushes the optimum back toward smaller models trained on more data, and makes distillation from a large reasoner into a small one economically central.

Open source to read

Reading order matters more than breadth. These are listed roughly from smallest to largest, with the file to open first.

1. karpathy/nanoGPT is the canonical minimal implementation and the right first read for anyone. Open model.py to find the entire GPT-2 architecture in about 300 lines, including the estimate_mfu method that is the reference implementation of the \( 6N + 12LdT \) accounting used on this page. Then train.py for gradient accumulation, DDP, and the cosine schedule.
2. karpathy/llm.c trains GPT-2 in raw C and CUDA with no framework at all. Open train_gpt2.c first for the CPU reference, where every gradient is written out by hand and the backward pass of each layer is fifteen readable lines, then train_gpt2.cu for the fused kernels. This is the fastest way to stop treating autograd as magic.
3. karpathy/minbpe is byte-level BPE in a few hundred lines. Open basic.py for the algorithm without the regex, then regex.py for the pre-tokenization that makes it practical.
4. huggingface/tokenizers is the production Rust implementation. Open tokenizers/src/models/bpe/trainer.rs to see the same inverted-index merge loop this page derives, with the priority queue and parallelism a real trainer needs.
5. pytorch/torchtitan is the current reference for how PyTorch itself thinks large-scale training should be written. Open torchtitan/models/llama3/model.py for a clean Llama implementation, then torchtitan/distributed/parallel_dims.py to see FSDP2, tensor parallel, pipeline parallel, and context parallel composed as an explicit device mesh.
6. NVIDIA/Megatron-LM is where tensor and pipeline parallelism were invented and is still the most complete implementation. Open megatron/core/tensor_parallel/layers.py and read ColumnParallelLinear and RowParallelLinear. They are the column-then-row construction with one all-reduce, exactly as derived above.
7. EleutherAI/gpt-neox is a Megatron-DeepSpeed derivative built by a group that publishes its failures. Open configs/ first. The YAML files for their released models are the most honest record available of what hyperparameters people actually use at each scale.
8. allenai/OLMo is the most completely open model line, including data, code, intermediate checkpoints, and the training logs. Open olmo/train.py for the trainer, but the real value is the released checkpoint series, which lets you study how representations develop over training rather than inferring it from the endpoint.
9. mlfoundations/open_lm is a compact research-oriented trainer designed for controlled scaling experiments. Open open_lm/model.py. It is the codebase to use if you want to run the isoFLOP methodology yourself rather than read about it.
10. vllm-project/vllm is the reference implementation of paged attention and continuous batching. Open vllm/core/block_manager.py for the block table and copy-on-write prefix sharing, then vllm/core/scheduler.py for continuous batching and preemption.
11. huggingface/transformers is the interoperability layer everything else targets. Open src/transformers/models/llama/modeling_llama.py. It is the most-read transformer implementation in existence and a useful check on any detail you are unsure about, including the exact RoPE interleaving convention.

Common misconceptions

"The architecture is where the wins are." Measured here, at a matched 65.5M-token budget, the spread across LayerNorm versus RMSNorm, GELU versus SwiGLU, MHA versus GQA versus MQA, and tied versus untied embeddings is about 0.1 nats, and several of those differences are within seed noise. The spread between a quality-filtered corpus and a raw one, with everything else identical, is 0.21 nats. Data beats architecture, and it is not close.

"Bigger vocabulary is always better because it compresses more." Compression grows logarithmically in \( V \) while embedding parameters grow linearly. The optimality condition derived above balances the two, and the optimum depends on model size. 8,192 is right for a 27M-parameter model and about 108,000 for an 8B one. A 128K vocabulary on a small model spends most of its parameters on embedding rows that almost never receive gradient.

"\( 6ND \) is exact." It ignores attention, which costs \( 12LdT \) per token and becomes the dominant term when \( T > 6d \). For the model here at \( T = 512 \) it is 16% of the total, and for an 8B model at 128K context it is five times the parameter term. Using \( 6ND \) to plan a long-context run underestimates the compute by a large factor.

"An 8B model fits on an 80GB GPU." For inference, yes, at 16GB in bf16. For training with Adam it needs 16 to 18 bytes per parameter, or 119 to 134 GiB of static state before any activations, which is why ZeRO-3 or FSDP is mandatory. The confusion between inference and training memory is the single most common sizing error.

"Decode is slow because the model is big, so a faster GPU will fix it." Decode at small batch is bandwidth-bound, not compute-bound. The arithmetic intensity is roughly the batch size in FLOPs per byte, against this H100's ridge point of 249. A GPU with twice the FLOPs and the same bandwidth would not decode any faster at batch 1. The fixes are batching, quantization, and speculation, in that order of impact.

"Mixture-of-experts models are cheap because they have fewer active parameters." They are cheap in FLOPs and expensive in everything else. All parameters must be resident in memory, the all-to-all routing adds communication proportional to the batch, load imbalance wastes capacity, and MFU is typically 10 to 20 points lower than a dense model. DeepSeek-V3's 671B total parameters occupy 671B parameters' worth of HBM regardless of the 37B that are active per token.

"A new optimizer gives 2× speedups." Almost every such claim is measured against an Adam baseline that was not tuned. The measurements on this page, with each optimizer's learning rate swept, put the spread at a few percent of validation loss. Muon is the one recent method that has survived a frontier-scale deployment, and even there the honest claim is tens of percent in tokens-to-loss, not a factor of two.

"Perplexity is comparable across models." Only if they share a tokenizer. Perplexity is per token, and a model with better compression has fewer, harder tokens. Bits per byte, \( L/(\ln 2 \cdot \beta) \), is the tokenizer-independent quantity, and any table comparing perplexities across differently tokenized models is measuring the tokenizers as much as the models.

"Emergent abilities appear discontinuously at scale." The underlying loss falls smoothly and predictably. What jumps is a thresholded metric such as exact-match accuracy on a multi-step task. Measuring the same tasks with a continuous metric, as Schaeffer and colleagues did, largely removes the discontinuity. This matters practically because it means capability at a future scale is more predictable than the emergence framing suggests.

Self-check

References

1. Goodfellow, I., Bengio, Y., Courville, A. (2016). Deep Learning. MIT Press. deeplearningbook.org. Background for the optimization and representation-learning material.
2. Jurafsky, D., Martin, J. H. (2025). Speech and Language Processing, 3rd edition draft. Chapters on subword tokenization and large language models. slp3.
3. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., Polosukhin, I. (2017). Attention Is All You Need. NeurIPS. arXiv:1706.03762.
4. Sennrich, R., Haddow, B., Birch, A. (2016). Neural Machine Translation of Rare Words with Subword Units. ACL. Edinburgh. arXiv:1508.07909.
5. Kudo, T. (2018). Subword Regularization: Improving NMT Models with Multiple Subword Candidates. ACL. Google. arXiv:1804.10959.
6. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I. (2019). Language Models are Unsupervised Multitask Learners. OpenAI. technical report.
7. Brown, T. B., et al. (2020). Language Models are Few-Shot Learners. NeurIPS. OpenAI. arXiv:2005.14165.
8. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., Amodei, D. (2020). Scaling Laws for Neural Language Models. OpenAI and Johns Hopkins. arXiv:2001.08361.
9. Hoffmann, J., Borgeaud, S., Mensch, A., et al. (2022). Training Compute-Optimal Large Language Models. DeepMind. arXiv:2203.15556.
10. Touvron, H., Lavril, T., Izacard, G., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. Meta AI. arXiv:2302.13971.
11. Grattafiori, A., et al. (2024). The Llama 3 Herd of Models. Meta AI. arXiv:2407.21783.
12. Su, J., Lu, Y., Pan, S., Wen, B., Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. Zhuiyi Technology. arXiv:2104.09864.
13. Press, O., Smith, N. A., Lewis, M. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. Washington, Facebook AI, Allen AI. arXiv:2108.12409.
14. Peng, B., Quesnelle, J., Fan, H., Shippole, E. (2023). YaRN: Efficient Context Window Extension of Large Language Models. Nous Research and EleutherAI. arXiv:2309.00071.
15. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. Google. arXiv:1911.02150.
16. Shazeer, N. (2020). GLU Variants Improve Transformer. Google. arXiv:2002.05202.
17. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebron, F., Sanghai, S. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Google Research. arXiv:2305.13245.
18. Zhang, B., Sennrich, R. (2019). Root Mean Square Layer Normalization. Edinburgh and Zurich. NeurIPS. arXiv:1910.07467.
19. Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., Liu, T.-Y. (2020). On Layer Normalization in the Transformer Architecture. Microsoft Research Asia and Peking. ICML. arXiv:2002.04745.
20. Xiao, G., Tian, Y., Chen, B., Han, S., Lewis, M. (2023). Efficient Streaming Language Models with Attention Sinks. MIT, Meta AI, CMU. arXiv:2309.17453.
21. Rajbhandari, S., Rasley, J., Ruwase, O., He, Y. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. Microsoft. SC20. arXiv:1910.02054.
22. Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., Catanzaro, B. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. NVIDIA. arXiv:1909.08053.
23. Narayanan, D., Shoeybi, M., Casper, J., et al. (2021). Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. NVIDIA, Stanford, Microsoft. SC21. arXiv:2104.04473.
24. Zhao, Y., Gu, A., Varma, R., et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. Meta. VLDB. arXiv:2304.11277.
25. Huang, Y., Cheng, Y., Bapna, A., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. Google. NeurIPS. arXiv:1811.06965.
26. Liu, H., Zaharia, M., Abbeel, P. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. Berkeley. arXiv:2310.01889.
27. Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., Chen, Z. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. Google. arXiv:2006.16668.
28. Fedus, W., Zoph, B., Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Google. JMLR. arXiv:2101.03961.
29. DeepSeek-AI (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.
30. Jiang, A. Q., Sablayrolles, A., Roux, A., et al. (2024). Mixtral of Experts. Mistral AI. arXiv:2401.04088.
31. Yang, A., et al. (2025). Qwen3 Technical Report. Alibaba. arXiv:2505.09388.
32. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Stanford. NeurIPS. arXiv:2205.14135.
33. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., Stoica, I. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. Berkeley. SOSP. arXiv:2309.06180.
34. Leviathan, Y., Kalman, M., Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. Google Research. ICML. arXiv:2211.17192.
35. Frantar, E., Ashkboos, S., Hoefler, T., Alistarh, D. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. ETH Zurich and IST Austria. arXiv:2210.17323.
36. Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., Han, S. (2023). AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration. MIT. arXiv:2306.00978.
37. Micikevicius, P., Narang, S., Alben, J., et al. (2018). Mixed Precision Training. NVIDIA and Baidu. ICLR. arXiv:1710.03740.
38. Loshchilov, I., Hutter, F. (2019). Decoupled Weight Decay Regularization. Freiburg. ICLR. arXiv:1711.05101.
39. McCandlish, S., Kaplan, J., Amodei, D., et al. (2018). An Empirical Model of Large-Batch Training. OpenAI. arXiv:1812.06162.
40. Penedo, G., Malartic, Q., Hesslow, D., et al. (2023). The RefinedWeb Dataset for Falcon LLM. Technology Innovation Institute. arXiv:2306.01116.
41. Penedo, G., Kydlicek, H., allal, L. B., Lozhkov, A., Mitchell, M., Raffel, C., Von Werra, L., Wolf, T. (2024). The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale. Hugging Face. arXiv:2406.17557.
42. Soldaini, L., Kinney, R., Bhagia, A., et al. (2024). Dolma: an Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. Allen AI. arXiv:2402.00159.
43. Lee, K., Ippolito, D., Nystrom, A., Zhang, C., Eck, D., Callison-Burch, C., Carlini, N. (2022). Deduplicating Training Data Makes Language Models Better. Google and Pennsylvania. ACL. arXiv:2107.06499.
44. Muennighoff, N., Rush, A. M., Barak, B., Le Scao, T., Piktus, A., Tazi, N., Pyysalo, S., Wolf, T., Raffel, C. (2023). Scaling Data-Constrained Language Models. Hugging Face, Harvard, Turku. NeurIPS. arXiv:2305.16264.
45. Chen, X., Liang, C., Huang, D., et al. (2023). Symbolic Discovery of Optimization Algorithms (Lion). Google and UCLA. arXiv:2302.06675.
46. Liu, H., Li, Z., Hall, D., Liang, P., Ma, T. (2023). Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. Stanford. arXiv:2305.14342.
47. Jordan, K., Jin, Y., Boza, V., You, J., Cesista, F., Newhouse, L., Bernstein, J. (2024). Muon: An optimizer for hidden layers in neural networks. project page.
48. Schaeffer, R., Miranda, B., Koyejo, S. (2023). Are Emergent Abilities of Large Language Models a Mirage? Stanford. NeurIPS. arXiv:2304.15004.
49. Gu, A., Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. CMU and Princeton. arXiv:2312.00752.
50. Groeneveld, D., Beltagy, I., Walsh, P., et al. (2024). OLMo: Accelerating the Science of Language Models. Allen AI and Washington. ACL. arXiv:2402.00838.