Case Study: OLMo 3

Every section so far built one piece. This one takes a real, recent model and shows where each piece went. OLMo 3 (Allen Institute for AI (artificial intelligence), released in late 2025) is a family of 7B and 32B models (7 and 32 billion parameters) whose authors published everything needed to rebuild them: the data, the training code (OLMo-core), the exact configuration scripts, the intermediate checkpoints and a 118-page report. Every number on this page comes from the report (Olmo Team, 2025), the official training scripts in OLMo-core, or the released model configuration on Hugging Face. Where a number is the result of a calculation of ours, the page says so.

This page covers the base model (pre-training, mid-training, long-context extension); the post-training of OLMo 3 into its Instruct and Think models belongs to section 12. Each part has a short example, and the exercises rebuild the numbers: the parameter count from the configuration, the attention pattern, one OLMo 3 block at toy size, the z-loss and its gradient, and the arithmetic of the training run.

The recipe in three stages

A base model is not trained in one go on one kind of data. It goes through three stages, each on different data, each with its own goal.

(Tokens per stage: T is trillions, B is billions. The context is the number of tokens the model reads at once. The data mixes are called Dolma 3 Mix, Dolma 3 Dolmino Mix and Dolma 3 Longmino Mix.)

• Pre-training is one pass (one epoch) over the mix: the official 32B script sets max_duration=Duration.epochs(1). The 32B run follows the same cosine schedule as a full epoch but stops early, at 5.5T tokens.
• Mid-training continues from that checkpoint on 100B tokens selected to boost math, code and question answering, and includes instruction data and reasoning traces to prepare for post-training. For the 32B it was run twice with different data orders, and the two resulting models were averaged weight by weight: the new model's weight is the mean of the two models' weights, matrix entry by entry.
• Long context teaches the model to use 65K-token contexts.

Example. Averaging two models weight by weight: if one model has the weight $0.30$ at some position of some matrix and the other has $0.34$, the averaged model has $0.32$ there, and the same is done for every weight. The result has the size of one model, not two.

Adding up the three stages of the 7B gives $5.93\text{T} + 0.1\text{T} + 0.05\text{T} \approx 6.08\text{T}$ tokens (our sum).

The data

Dolma 3 Mix is built from a 9.31-trillion-token pool of cleaned, deduplicated text (a document that appears many times is kept once). The mix is the part of the pool actually used:

"One epoch" over the mix does not mean every document is seen once. The mix itself upsamples some sources: code is 137B tokens in the pool and 409B in the mix, so it is repeated about three times, while most of the web pool is left out.

Example. The web share of the mix is $76.1\% \cdot 5.93\text{T} \approx 4.51\text{T}$ tokens (our product). Code: $409 / 137 \approx 2.99$, so each token of the code pool appears about 3 times in the mix.

Tokenizer

The tokenizer is the one of OLMo 2, derived from OpenAI's cl100k byte-level BPE (byte-pair encoding, section 6): 100,278 tokens in the released model, padded to a multiple of 128 during training so that the embedding matrix has dimensions that suit a GPU (graphics processing unit).

Example. Rounding 100,278 up to the next multiple of 128 gives $784 \cdot 128 = 100{,}352$: 74 extra rows in the embedding matrix that no token uses (our arithmetic).

The architecture

OLMo 3 is a decoder-only transformer (a stack of blocks like that of section 7, with causal attention), with the same block for every layer:

The parameter counts are those of the released checkpoints; one exercise derives them from the table. With 32 key/value heads for 32 query heads, the 7B has multi-head attention (MHA); with 8 for 40, the 32B has grouped-query attention (GQA), both from section 7. The feed-forward part of each block is an MLP (multilayer perceptron). Each head has $D / H = 128$ numbers, for both sizes. Inside each block, in the order the data flows, there are seven ingredients. We go through them one at a time.

1. Projections without biases

Queries, keys and values come from linear layers with no bias terms: only the weight matrix. The 32B uses grouped-query attention with 8 key/value heads for 40 query heads, so each key/value head serves $40/8 = 5$ query heads (section 7).

2. QK-norm

QK-norm is an RMSNorm (root mean square layer normalization, section 4) on the projected queries and on the projected keys, before attention. It keeps the attention logits from becoming too large, which can make the training loss diverge (Dehghani et al., 2023).

Example. Suppose a query vector grew to $q = (30, 40)$ and the key to the same. Their dot product is $30^2 + 40^2 = 2500$, and after the softmax all the attention goes to one token. After RMSNorm, $q \approx (0.85, 1.13)$, and so the key, and their dot product is $2.0$. Whatever the size of the projection, the score stays bounded.

3. Rotary embeddings with $\theta = 500{,}000$

RoPE (rotary position embedding, section 7) with $\theta = 500{,}000$ instead of the original 10,000: OLMo 2 raised it to increase the resolution of the positional encoding. A larger base makes the slower pairs of numbers turn even more slowly.

Example. Head size $d = 128$ means 64 pairs. The slowest pair turns by $\theta^{-126/128}$ radians per position: with $\theta = 10{,}000$ that is $1.15 \times 10^{-4}$, a full turn every $54{,}000$ positions; with $\theta = 500{,}000$ it is $2.5 \times 10^{-6}$, a full turn every $2.6$ million positions (our arithmetic). Positions much farther apart than 54,000 tokens would otherwise look alike to the slowest pair.

The released code pairs dimension $i$ with $i + d/2$ (the "rotate half" layout), which is equivalent to the adjacent-pair layout of section 7 up to a permutation of the weights.

4. Sliding-window attention

In three layers out of four, tokens attend only to the previous 4,096 tokens (Beltagy et al., 2020): this is sliding-window attention. Every fourth layer, and always the last, uses full attention. This keeps the cost and the KV (key-value) cache of long contexts manageable (the cache is the memory of the keys and values of all previous tokens, section 11).

Example. In a model of 8 layers the types are: sliding, sliding, sliding, full, sliding, sliding, sliding, full. With a window of 3, the token at position 4 sees positions 2, 3 and 4 in a sliding layer, and positions 0 to 4 in a full layer.

For a sequence of 65,536 tokens, a sliding layer keeps 4,096 keys instead of 65,536: $1/16$. In the 32-layer model, 24 layers are sliding and 8 are full, so the cache is $(24 \cdot 4096 + 8 \cdot 65536) / (32 \cdot 65536) \approx 0.30$ of what it would be with full attention everywhere. The 64-layer model has the same ratio (our arithmetic).

5. A SwiGLU MLP

The MLP of each block is the SwiGLU (Swish-gated linear unit) of section 7: three matrices, the second branch gating the first.

6. Reordered norm

RMSNorm is applied to the output of attention and of the MLP, before the residual addition, instead of to their input as in the LLaMA (Large Language Model Meta AI) block of section 7 (introduced for OLMo 2, following Liu et al., 2022):

where $x$ is the stream entering the block, $h$ the stream after attention and $y$ the output of the block.

Example. Suppose attention outputs $(300, 400)$. With the norm on the output, what is added to the stream is $(0.85, 1.13)$ times the learned gain: the correction of each sub-layer has a controlled size, whatever the attention produced.

7. Embeddings, output layer and long-context scaling

The input embedding and the output layer are separate matrices (not tied), and a final RMSNorm precedes the output layer. For the long-context stage, RoPE on the full-attention layers is rescaled with YaRN (Peng et al., 2023), from the 8,192 positions seen in training to 65,536.

Counting the parameters

With these ingredients one can count every parameter from the configuration. A tiny example (toy sizes, our own): $V = 10$ tokens, $D = 8$, 2 layers, 2 query heads and 1 key/value head of size $d = 4$, MLP size $F = 16$.

Example. Per layer:

• attention: $W_q$ is $8 \times 8 = 64$, $W_k$ and $W_v$ are $8 \times 4 = 32$ each, $W_o$ is $8 \times 8 = 64$: total $192$;
• QK-norm gains: $8 + 4 = 12$;
• MLP: three matrices of $8 \times 16 = 128$: $384$;
• two RMSNorms of size 8: $16$.

That is $604$ per layer, $1208$ for two. Add the two embedding matrices $2 \cdot 10 \cdot 8 = 160$ and the final norm $8$: $1376$ parameters.

The same count for the 7B (32 layers of $202{,}391{,}552$ parameters each, two embedding matrices of $100{,}278 \cdot 4096$, and the final norm of $4096$) gives exactly $7{,}298{,}011{,}136$, the number of the released checkpoint.

The loss

The training loss is the next-token cross-entropy of section 2, plus a small z-loss:

where $\mathrm{CE}$ is the cross-entropy, $\log Z = \mathrm{LSE}(z)$ is the log-sum-exp of each position's logits, and the bar is the mean over positions. Softmax does not change if every logit is shifted by a constant, so nothing stops the logits from drifting to huge values; the z-loss pulls $\log Z$ towards 0 and keeps them bounded. PaLM (Pathways Language Model) introduced it with weight $10^{-4}$ (Chowdhery et al., 2022); OLMo 3 uses $10^{-5}$.

Example. Logits $(2, 1, 0)$, right class 0: $\log Z = 2.408$ and $\mathrm{CE} = 0.408$. The z-loss term is $10^{-5} \cdot 2.408^2 \approx 5.8 \times 10^{-5}$.

Now add 100 to every logit: the softmax and the cross-entropy are unchanged ($0.408$), but $\log Z = 102.408$, and the z-loss term becomes $10^{-5} \cdot 102.408^2 \approx 0.105$: a penalty against the drift.

Optimizer and schedule (32B pre-training)

From the official OLMo-3-1025-32B-pretrain.py script and the report's Table 35:

• AdamW (Adam with decoupled weight decay, section 3) with $\beta = (0.9, 0.95)$ and weight decay 0.1, except on the embeddings, which get no weight decay.
• Learning rate: linear warmup over 2,000 steps to a peak of $6 \times 10^{-4}$, then cosine decay towards $6 \times 10^{-5}$ (truncated at 5.5T tokens). The 7B uses a peak of $3 \times 10^{-4}$.
• Batch: 1,024 sequences of 8,192 tokens, 8,388,608 tokens per step: about 656,000 steps for 5.5T tokens. The 7B uses half that batch.
• Gradient clipping at a global norm of 1.0 (section 3).
• Skipped steps: the optimizer (SkipStepAdamW) skips an update when the loss or the gradient norm is more than 6 standard deviations above its mean over the last 128 steps, a guard against single bad batches.
• Precision: BF16 (brain floating point) weights for computation, gradients reduced in FP32 (32-bit floating point) (section 8).

Example. The batch: $1024 \cdot 8192 = 8{,}388{,}608$ tokens. Steps: $5.5 \times 10^{12} / 8{,}388{,}608 = 655{,}651.09$, rounded up to $655{,}652$ steps: the "about 656,000" above.

The skip rule: if the loss over the last 128 steps had a mean of $2.5$ and a standard deviation of $0.1$, a step whose loss is above $2.5 + 6 \cdot 0.1 = 3.1$ would be skipped.

Hardware, parallelism and cost

The 32B was trained on a cluster of NVIDIA H100 GPUs (80 GB (gigabytes) each), 8 per node, with the OLMo-core library on PyTorch, torch.compile, FlashAttention kernels (Dao, 2024) and asynchronous checkpointing.

• Parallelism: HSDP (hybrid sharded data parallelism), the variant of FSDP (Fully Sharded Data Parallel) of section 9: the model states are sharded across groups of 64 GPUs, and 16 such groups are replicas, 1,024 GPUs in total. Activation checkpointing is limited to a memory budget of 50%. The long-context stage adds context parallelism, which splits each 65K-token sequence across 8 GPUs.
• Throughput: 1,960 tokens per second per GPU for the 32B and 7,700 for the 7B, which the report puts at about 41% and 43% MFU (model FLOPs utilization, section 8).
• Time: pre-training took about 9.5 days on 512 GPUs and then 35 days on 1,024 GPUs, crashes and restarts included; with mid-training and the long-context stage, about 47 days. From the start of training to the evaluated 32B Think model, about 56 days on 1,024 GPUs, which at $2 per H100-hour the report prices at $2.75M.

Example. The numbers check each other (our arithmetic).

• At 1,960 tokens per second per GPU on 1,024 GPUs the cluster processes $2.0$ million tokens per second, so 5.5T tokens would take $31.7$ days if nothing ever went wrong. The real run took $9.5 + 35 = 44.5$ days, with half the GPUs at the start, plus crashes and restarts.
• The pre-training used $9.5 \cdot 512 + 35 \cdot 1024 = 40{,}704$ GPU days.
• The cost: 56 days times 1,024 GPUs times 24 hours times $2 is $2,752,512, about $2.75M: the figure of the report.

References

• Olmo Team, Olmo 3, arXiv, 2025. Sections 2.4 (costs) and 3.2 (architecture and training), Tables 4 (data), 33 (architecture), 34 (parallelism and throughput) and 35 (hyperparameters).
• Olmo Team, 2 OLMo 2 Furious, arXiv, 2024 (COLM 2025). Section 2.1: RMSNorm, reordered norm, QK-norm, z-loss and RoPE $\theta$, inherited by OLMo 3.
• Allen Institute for AI, OLMo-core: the training library and the official scripts in src/scripts/official/OLMo3/.
• Released configurations: Olmo-3-1025-7B and Olmo-3-1125-32B; the reference implementation in Hugging Face Transformers, modeling_olmo3.py.
• M. Dehghani et al., Scaling Vision Transformers to 22 Billion Parameters, ICML, 2023. QK-norm.
• Z. Liu et al., Swin Transformer V2: Scaling Up Capacity and Resolution, CVPR, 2022. Normalizing the output of each sub-layer.
• I. Beltagy, M. E. Peters, A. Cohan, Longformer: The Long-Document Transformer, arXiv, 2020. Sliding-window attention.
• B. Peng, J. Quesnelle, H. Fan, E. Shippole, YaRN: Efficient Context Window Extension of Large Language Models, arXiv, 2023.
• A. Chowdhery et al., PaLM: Scaling Language Modeling with Pathways, arXiv, 2022. Section 5: the z-loss.
• T. Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, ICLR, 2024.