Total LayerNorm parameters in GPT-2

Background

Each nn.LayerNorm(n_embd) learns two length-C vectors — γ (scale) and β (shift) — so 2C parameters per LN. A GPT-2 stack has two LNs per block (before attention and before the MLP) plus one final LN before the unembedding:

$$\text{LN count} = 2 \cdot n\_layer + 1, \qquad \text{params} = (2\,n\_layer + 1)\cdot 2C$$

Problem statement

Implement layernorm_params(n_embd, n_layer) returning the total LayerNorm parameter count.

Input

• n_embd — feature dim C.
• n_layer — number of transformer blocks.

Output

Returns an int: (2*n_layer + 1) * 2 * n_embd.

Examples

Example 1 — GPT-2 small

Input:  n_embd = 768, n_layer = 12
Output: 38400

Explanation: 25 LNs × 1536 each.

Constraints

• 2*n_layer + 1 LayerNorms, each 2*n_embd params.
• Return the product.

Notes

• Tiny (<0.1% of 124M) but architecturally non-negotiable — without it, deep stacks diverge.