Count OLMo 3's Parameters

The problem

The released config.json file of a model lists its sizes: the number of layers, the model dimension, the number of heads, and so on. That is enough to count its parameters exactly, and a good test that you understand the architecture: if your count matches the published checkpoint to the last digit, you have not forgotten any matrix.

Implement count_params(cfg), where cfg is a dict with the keys of the Hugging Face configuration of OLMo 3:

• vocab_size ($V$), the number of tokens;
• hidden_size ($D$), the model dimension;
• num_hidden_layers ($L$), the number of layers;
• num_attention_heads ($H$) and num_key_value_heads ($G$);
• intermediate_size ($F$), the hidden size of the feed-forward part;
• tie_word_embeddings: whether the input and output layers share one matrix.

The head size is $d = D / H$. OLMo 3 has no bias anywhere.

What to count

Once for the whole model:

• the input embedding, $V \times D$, and the output layer, another $V \times D$ unless tie_word_embeddings is true (then they share one matrix);
• a final RMSNorm (root mean square layer normalization) gain of size $D$.

In each of the $L$ layers:

• attention projections: $W_q$ is $D \times Hd$, $W_k$ and $W_v$ are $D \times Gd$, $W_o$ is $Hd \times D$;
• QK-norm gains: QK-norm normalizes the queries and the keys before attention, with a learned gain per feature: one of size $Hd$ for the queries and one of size $Gd$ for the keys;
• the SwiGLU (Swish-gated linear unit) MLP (multilayer perceptron): gate and up projections of $D \times F$ each, and a down projection of $F \times D$;
• two RMSNorm gains of size $D$ (one after attention, one after the MLP).

Example. A toy configuration: $V = 10$, $D = 8$, $L = 2$, $H = 2$, $G = 1$, $F = 16$, not tied. Then $d = 4$.

part	count
attention: $8 \cdot 8 + 2 \cdot (8 \cdot 4) + 8 \cdot 8$	$192$
QK-norm: $2 \cdot 4 + 1 \cdot 4$	$12$
MLP: $3 \cdot 8 \cdot 16$	$384$
two norms: $2 \cdot 8$	$16$
per layer	$604$

Total: $2 \cdot 10 \cdot 8$ (two embedding matrices) $+ \; 2 \cdot 604$ (layers) $+ \; 8$ (final norm) $= 160 + 1208 + 8 = 1376$.

What to implement

count_params(cfg): the total number of parameters.

Requirements

• Return a Python int.
• The tests compare against the parameter counts of the published checkpoints: 7,298,011,136 for the 7B and 32,233,522,176 for the 32B.

Hints

• With grouped-query attention, $W_k$, $W_v$ and the key norm depend on $G$, not on $H$. If your count does not change when you change num_key_value_heads, one of them uses the wrong size.
• With tie_word_embeddings false, the model has two $V \times D$ matrices: forgetting the second is the most common mistake.