Parameter breakdown of a GPT-2 block

Background

A full no-bias GPT-2 block has four components. Per embedding dim C:

Component	Params
ln1	2C
attn	4C²
ln2	2C
mlp	8C²

So the block totals 12C² + 4C. For GPT-2 small (C=768) that's 7,080,960, and the MLP alone is twice the attention.

Problem statement

Implement block_param_breakdown(n_embd) returning each component's parameter count and the total.

Input

• n_embd — embedding dim C.

Output

Returns a dict with keys "ln1", "attn", "ln2", "mlp", "total":

• ln1 = ln2 = 2*C, attn = 4*C**2, mlp = 8*C**2, total = ln1+attn+ln2+mlp.

Examples

Example 1 — GPT-2 small

Input:  n_embd = 768
Output: {"ln1": 1536, "attn": 2359296, "ln2": 1536, "mlp": 4718592, "total": 7080960}

Constraints

• ln1 = ln2 = 2*n_embd, attn = 4*n_embd**2, mlp = 8*n_embd**2.
• total is the sum of all four.
• Return a dict with those five keys.

Notes

• The two C² terms (attention 4C², MLP 8C²) dominate; the LayerNorms are a rounding error.