<?xml version="1.0" encoding="UTF-8"?>
<rss version="2.0" xmlns:atom="http://www.w3.org/2005/Atom" xmlns:dc="http://purl.org/dc/elements/1.1/">
  <channel>
    <title>DEV Community: Carlos Chinchilla Corbacho</title>
    <description>The latest articles on DEV Community by Carlos Chinchilla Corbacho (@cchinchilladev).</description>
    <link>https://dev.to/cchinchilladev</link>
    <image>
      <url>https://media2.dev.to/dynamic/image/width=90,height=90,fit=cover,gravity=auto,format=auto/https:%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Fuser%2Fprofile_image%2F3893016%2Fc50a9879-1fa7-4bca-8835-c95162affe58.jpg</url>
      <title>DEV Community: Carlos Chinchilla Corbacho</title>
      <link>https://dev.to/cchinchilladev</link>
    </image>
    <atom:link rel="self" type="application/rss+xml" href="https://dev.to/feed/cchinchilladev"/>
    <language>en</language>
    <item>
      <title>From GPT-2 to Llama: four changes and the reason for each</title>
      <dc:creator>Carlos Chinchilla Corbacho</dc:creator>
      <pubDate>Thu, 15 Oct 2026 13:00:00 +0000</pubDate>
      <link>https://dev.to/cchinchilladev/from-gpt-2-to-llama-four-changes-and-the-reason-for-each-4326</link>
      <guid>https://dev.to/cchinchilladev/from-gpt-2-to-llama-four-changes-and-the-reason-for-each-4326</guid>
      <description>&lt;p&gt;&lt;em&gt;Originally published at &lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-to-llama-what-changed/" rel="noopener noreferrer"&gt;cchinchilla.dev&lt;/a&gt;. Part 7 of From code to weights, a 12-part series on ML fundamentals for engineers.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-transformer-plus-five-decisions/" rel="noopener noreferrer"&gt;Part 6&lt;/a&gt; ended on GPT-2's five decisions and the loop that serves it. Open a model from this decade and the block is still Part 5's: two stages, each reading a normalised copy of the bus and adding its result back. Four things inside it changed.&lt;/p&gt;

&lt;p&gt;I read &lt;code&gt;modeling_llama.py&lt;/code&gt; from transformers 5.17 next to nanoGPT's &lt;code&gt;model.py&lt;/code&gt;, and loaded SmolLM2-135M, a model with Llama's architecture at GPT-2 small's size. Then I undid each change on its trained weights and measured what broke.&lt;/p&gt;

&lt;h2&gt;
  
  
  Same block, four edits
&lt;/h2&gt;

&lt;p&gt;&lt;code&gt;LlamaDecoderLayer.forward&lt;/code&gt;, &lt;code&gt;modeling_llama.py:305-323&lt;/code&gt;, condensed:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;residual&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;hidden_states&lt;/span&gt;
&lt;span class="n"&gt;hidden_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;input_layernorm&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;            &lt;span class="c1"&gt;# RMSNorm, not LayerNorm
&lt;/span&gt;&lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;_&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;self_attn&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;position_embeddings&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="p"&gt;...)&lt;/span&gt;
&lt;span class="n"&gt;hidden_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;residual&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;hidden_states&lt;/span&gt;                       &lt;span class="c1"&gt;# the bus, as in GPT-2
&lt;/span&gt;&lt;span class="n"&gt;residual&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;hidden_states&lt;/span&gt;
&lt;span class="n"&gt;hidden_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;mlp&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;post_attention_layernorm&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;
&lt;span class="n"&gt;hidden_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;residual&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;hidden_states&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;This is Part 5's four lines with more names; the plus signs are where they were, and so is pre-LN, the norm applied before each stage rather than after. What changed sits inside the calls: how positions enter attention, what the norm computes, the shape of the feed-forward, and how many keys and values attention keeps. (Llama also drops the bias vectors from its linear projections; nothing below depends on them.)&lt;/p&gt;

&lt;p&gt;The two models make a fair pair. GPT-2 small is 124 million parameters in 12 blocks of width 768. SmolLM2-135M is 135 million in 30 blocks of width 576. None of the four changes is about size; you could make either model with or without them.&lt;/p&gt;

&lt;h2&gt;
  
  
  Positions become a rotation
&lt;/h2&gt;

&lt;p&gt;GPT-2 adds a learned row to each token, one row per position, and the table ends at 1,024. Llama has no position table. It rotates the query and key instead, at &lt;code&gt;modeling_llama.py:158-159&lt;/code&gt;:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;q_embed&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;q&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;cos&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nf"&gt;rotate_half&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;q&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;sin&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;k_embed&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;k&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;cos&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nf"&gt;rotate_half&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;k&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;sin&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;This is rotary position embedding, RoPE. In Part 4, attention scores a pair of tokens by comparing a query vector with a key vector. Each query and key is split into heads of 64 numbers. Each head's vector is read as 32 pairs, each pair a point on a circle, and each point is turned like a clock hand by an angle proportional to the position. &lt;code&gt;rotate_half&lt;/code&gt; is how the code writes that turn without a matrix. Every pair has its own speed. One config number, the rotation base, sets all of them; SmolLM2's is 100,000. In SmolLM2 the fastest pair goes round once every 6.3 positions and the slowest once every 438,000.&lt;/p&gt;

&lt;p&gt;The rotation is not decoration. Rotate a query by position &lt;code&gt;m&lt;/code&gt; and a key by position &lt;code&gt;n&lt;/code&gt;, and their dot product, the score from Part 4, depends only on &lt;code&gt;m − n&lt;/code&gt;. I took a random query and key through SmolLM2's rotation. At positions 10 and 3 the score was −13.4724; at 1,010 and 1,003 it was −13.4723. One step closer, at 10 and 4, it was −10.786.&lt;/p&gt;

&lt;p&gt;The whole model inherits that. I shifted every position id in a 512-token text by 500, then by 5,000. SmolLM2's loss on it, its average surprise at the next token, was 3.381 unshifted and 3.381 after each shift. GPT-2 given the same shift of 500 went from 3.85 to 8.97, because every token now reads a different row of its table. Removing the rotation altogether, by setting every position to 0, took SmolLM2 from 3.12 to 6.28 on the full 1,024-token text. Using 10,000 as the rotation base, the value in the original RoPE paper and in Llama 2, instead of SmolLM2's 100,000 cost 3%.&lt;/p&gt;

&lt;p&gt;The original post has a widget for this that doesn't survive being pasted into a feed: &lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-to-llama-what-changed/" rel="noopener noreferrer"&gt;drag the key's position&lt;/a&gt; and watch the score move, then shift both and watch it stay. Pick pair 17, the slowest of the three, and nearby positions barely differ. The score here is one pair's share, between −1 and 1.&lt;/p&gt;

&lt;p&gt;No table also means no last row. Llama 3.1 8B's config records an original length of 8,192 positions and a maximum of 131,072. It gets there by stretching the rotation. Its heads have 128 numbers, so 64 pairs: the 29 slowest are slowed by a factor of 8, the 29 fastest left alone, and the six between blended. Meta also &lt;a href="https://arxiv.org/abs/2407.21783" rel="noopener noreferrer"&gt;pre-trained the family&lt;/a&gt; at growing lengths up to 128K, in six stages for the 405B. There is no assert to hit: past the length a model was trained at, quality is a measurement, not a guarantee.&lt;/p&gt;

&lt;h2&gt;
  
  
  The norm stops centring
&lt;/h2&gt;

&lt;p&gt;&lt;code&gt;LayerNorm&lt;/code&gt; subtracts each row's mean, divides by its standard deviation, then scales and shifts. &lt;code&gt;LlamaRMSNorm&lt;/code&gt;, &lt;code&gt;modeling_llama.py:62-67&lt;/code&gt;, skips the mean and the shift, and divides by the root mean square instead: square each number, average, take the root. If you've standardised a column in a dataframe, you've subtracted and divided; RMSNorm keeps only the division:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;variance&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;pow&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;mean&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;keepdim&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="bp"&gt;True&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;hidden_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;hidden_states&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;rsqrt&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;variance&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;variance_epsilon&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;weight&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;to&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;input_dtype&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;The bet is that the mean is small anyway, and in SmolLM2 it is. In a typical norm layer, a typical row's mean is 2% of its root mean square; in the worst layer, 4.5%. Swap full LayerNorm back in with the same learned scale and the loss moves from 3.124 to 3.175, under 2%.&lt;/p&gt;

&lt;p&gt;Be fair about the reason given for it. &lt;a href="https://arxiv.org/abs/1910.07467" rel="noopener noreferrer"&gt;The RMSNorm paper&lt;/a&gt; sells it as cheaper: one statistic instead of two, no shift. On my laptop's CPU, torch's &lt;code&gt;F.layer_norm&lt;/code&gt; took 50 to 75 microseconds on a &lt;code&gt;[1, 1024, 576]&lt;/code&gt; tensor across runs, and &lt;code&gt;F.rms_norm&lt;/code&gt; 110 to 135. The profiler says why: on CPU, &lt;code&gt;rms_norm&lt;/code&gt; still runs as a chain of small operations, square, average, add, root and two multiplies, while &lt;code&gt;layer_norm&lt;/code&gt; is a single tuned kernel, one compiled routine for the whole job.&lt;/p&gt;

&lt;p&gt;On a single thread the order flips: &lt;code&gt;layer_norm&lt;/code&gt; about 155 microseconds, &lt;code&gt;rms_norm&lt;/code&gt; about 130. Less arithmetic isn't less time until someone writes the kernel, and I didn't measure a GPU.&lt;/p&gt;

&lt;h2&gt;
  
  
  The feed-forward grows a gate
&lt;/h2&gt;

&lt;p&gt;GPT-2's feed-forward is two matrices with GELU between them. Llama's is three, &lt;code&gt;modeling_llama.py:175&lt;/code&gt;:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;down_proj&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;down_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;act_fn&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;gate_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;up_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;The input goes through two projections side by side. One passes through SiLU, a smooth activation close to GELU, and then multiplies the other element by element: a gate. This is SwiGLU, from a &lt;a href="https://arxiv.org/abs/2002.05202" rel="noopener noreferrer"&gt;2020 paper&lt;/a&gt; that tested it and reported better results without offering a theory for why.&lt;/p&gt;

&lt;p&gt;Three matrices cost more than two at the same width, so the width shrinks. With &lt;code&gt;d&lt;/code&gt; the model width, a plain feed-forward of width &lt;code&gt;4d&lt;/code&gt; holds &lt;code&gt;8d²&lt;/code&gt; numbers, biases aside. Three matrices of width &lt;code&gt;f&lt;/code&gt; hold &lt;code&gt;3·d·f&lt;/code&gt;, which matches at &lt;code&gt;f = 8d/3&lt;/code&gt;. SmolLM2 uses exactly that: 1,536 is 8/3 of 576. Llama 2 7B's 11,008 is 8/3 of 4,096 rounded up to a multiple of 256.&lt;/p&gt;

&lt;p&gt;At &lt;code&gt;8d/3&lt;/code&gt; the gate changes nothing in the budget: with full attention heads, the feed-forward still holds two thirds of a block. SmolLM2's 75% comes from sharing keys and values across query heads, covered next. Llama 3.1 8B moves the split from both sides. Its feed-forward is wider, &lt;code&gt;f = 3.5d&lt;/code&gt;, and its attention is smaller because it shares keys and values. Together they give the 81% quoted in Part 5: 176 million of each block's 218 million parameters. The wider feed-forward alone, with full attention heads, gives 72%.&lt;/p&gt;

&lt;p&gt;The gate carries weight. Replace SiLU with GELU in SmolLM2 and the loss goes from 3.12 to 3.79. Drop the gate's multiplication entirely and it goes to 17.7, well past 10.8, the loss of a model that gives each of 49,152 tokens equal odds.&lt;/p&gt;

&lt;p&gt;The feed-forward is also where mixture of experts plugs in: many feed-forwards per block and a router, a small layer that picks a few of them for each token. Llama 3 stayed dense; Llama 4 and other families didn't.&lt;/p&gt;

&lt;h2&gt;
  
  
  Keys and values get shared
&lt;/h2&gt;

&lt;p&gt;GPT-2 has as many key and value heads as query heads. SmolLM2 has 9 query heads and 3 key-value heads; each key and value head serves a group of three queries. Think of one index serving three different queries: same data, different questions. In the plain PyTorch path, &lt;code&gt;repeat_kv&lt;/code&gt;, &lt;code&gt;modeling_llama.py:179-189&lt;/code&gt;, copies the three key and value heads out to nine before the arithmetic; torch's fused attention kernel can skip the copy when there is no padding mask. This is grouped-query attention, GQA. Llama 2 used it only in its two largest sizes, 70B and an unreleased 34B. Every Llama 3 size has it.&lt;/p&gt;

&lt;p&gt;At inference the sharing is exact. I built a copy of SmolLM2 with nine full key and value heads, each a duplicate of its group's, and ran both. The logits matched to the last bit, and the copy had 147.8 million parameters instead of 134.5. What GQA gives up, it gives up in training, with fewer distinct keys and values to learn.&lt;/p&gt;

&lt;p&gt;The saving is in Part 6's cache. SmolLM2 stores 11,520 numbers per token, which I read back from transformers' own cache after 256 tokens. With full heads it would store 34,560. GPT-2 small, with 12 blocks to SmolLM2's 30, stores 18,432.&lt;/p&gt;

&lt;p&gt;At Llama 3.1 8B's size the ratio decides what fits. It has 32 query heads and 8 key-value heads, so each token costs 131 KB in bf16, the 16-bit format from Part 3. At its full 131,072 positions, one sequence's cache is 17.2 GB, next to 16.1 GB of weights. With full heads it would be 68.7 GB.&lt;/p&gt;

&lt;p&gt;The second widget on &lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-to-llama-what-changed/" rel="noopener noreferrer"&gt;the original post&lt;/a&gt; starts from GPT-2 small. Switch the three changes on one at a time; the feed-forward share barely moves. Then drag the key-value heads down to 1 and watch the cache and the share.&lt;/p&gt;

&lt;h2&gt;
  
  
  What the ML engineer watches
&lt;/h2&gt;

&lt;p&gt;Four habits, and none of them needs a GPU.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Recompute the cache for every model.&lt;/strong&gt; Parameter count does not predict it: the cache is &lt;code&gt;2 × blocks × key-value heads × head width&lt;/code&gt; numbers per token, and the key-value heads are a config field. SmolLM2 is larger than GPT-2 small and stores three eighths less per token.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Port the architecture before the weights.&lt;/strong&gt; Every change above is part of the contract the weights were trained under. The activation, the position scheme with its rotation base and scaling, and the key-value head count are config fields; the norm type comes with the architecture name. Get the head count wrong and loading fails, because the shapes don't match. Get the count right and the grouping wrong, tiling the shared heads instead of interleaving them, and nothing raises: the loss went past 11. The same silence covers the rest. A bias-free LayerNorm instead of RMSNorm cost 1.6%, GELU instead of SiLU 21%, and the wrong rotation base 3%. Rotation pairs laid out side by side instead of in halves took the loss to 7.24. If you port one by hand, compare your logits against a reference implementation before you trust it.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Test the lengths you serve.&lt;/strong&gt; Rotation removes the hard limit and keeps the question. If you run a model past the length it was trained at, you are the one who has to measure it there, on your own data.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Count active parameters separately.&lt;/strong&gt; With mixture of experts, Qwen3-30B-A3B holds 30.5 billion parameters and uses 3.35 billion per token; 95% of the weights are experts. Memory follows the first number; arithmetic per token follows the second, and you pay for both.&lt;/p&gt;

&lt;p&gt;None of this is modelling work either. A Llama block is a GPT-2 block with four parts swapped, and each swap is a line of config or code you can read, a cost you can compute, and a mismatch you can test for before a single request.&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Part 7 of twelve. Llama keeps GPT-2's block and bus and changes four parts: positions become a rotation of queries and keys, the norm drops its mean, the feed-forward gains a gate, and query heads share keys and values. Get the norm, the activation or the positions wrong and nothing raises; the loss moves instead. Rotary positions drop a table, shared keys and values shrink the cache, and the gate, at &lt;code&gt;8d/3&lt;/code&gt;, costs nothing. The next part asks where the weights come from, and lands on &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;the series index&lt;/a&gt; when it's out. If you arrived here directly, &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Part 1&lt;/a&gt; reads the tokenizer, &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;Part 2&lt;/a&gt; the table a token id indexes, &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;Part 3&lt;/a&gt; the shapes, &lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;Part 4&lt;/a&gt; attention, &lt;a href="https://cchinchilla.dev/posts/fcw-transformer-six-shapes-residual-stream/" rel="noopener noreferrer"&gt;Part 5&lt;/a&gt; the block, and &lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-transformer-plus-five-decisions/" rel="noopener noreferrer"&gt;Part 6&lt;/a&gt; GPT-2 and its two programs.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Written with AI assistance. See &lt;a href="https://cchinchilla.dev/ai-policy/" rel="noopener noreferrer"&gt;AI policy&lt;/a&gt;.&lt;/em&gt;&lt;/p&gt;

</description>
      <category>machinelearning</category>
      <category>python</category>
      <category>ai</category>
      <category>beginners</category>
    </item>
    <item>
      <title>GPT-2 is the transformer plus five decisions, and inference is a different program</title>
      <dc:creator>Carlos Chinchilla Corbacho</dc:creator>
      <pubDate>Thu, 08 Oct 2026 13:00:00 +0000</pubDate>
      <link>https://dev.to/cchinchilladev/gpt-2-is-the-transformer-plus-five-decisions-and-inference-is-a-different-program-4gib</link>
      <guid>https://dev.to/cchinchilladev/gpt-2-is-the-transformer-plus-five-decisions-and-inference-is-a-different-program-4gib</guid>
      <description>&lt;p&gt;&lt;em&gt;Originally published at &lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-transformer-plus-five-decisions/" rel="noopener noreferrer"&gt;cchinchilla.dev&lt;/a&gt;. Part 6 of From code to weights, a 12-part series on ML fundamentals for engineers.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;a href="https://cchinchilla.dev/posts/fcw-transformer-six-shapes-residual-stream/" rel="noopener noreferrer"&gt;Part 5&lt;/a&gt; ended on a block and a stack of twelve. GPT-2 is that stack plus five decisions around it. Its one checkpoint runs two different programs. Training scores a whole sequence in one call. Generation, what serving calls inference, writes one token at a time, in a loop, and most of what breaks in production breaks there.&lt;/p&gt;

&lt;p&gt;I loaded OpenAI's GPT-2 small weights into nanoGPT's &lt;code&gt;model.py&lt;/code&gt; again and ran both programs on the same weights. Then I added a cache to the loop, cast the weights to bf16, and compared what came out.&lt;/p&gt;

&lt;h2&gt;
  
  
  Five decisions around the stack
&lt;/h2&gt;

&lt;p&gt;&lt;code&gt;GPTConfig&lt;/code&gt;, &lt;code&gt;model.py:108-116&lt;/code&gt;, fits on one screen: 1,024 positions, 12 blocks, 12 heads, width 768. The vocabulary field says 50,304, which the comment on that line calls padding to a multiple of 64, for speed. OpenAI's checkpoint has 50,257 rows, and loading it sets the config to match. Around those numbers GPT-2 made five decisions, and three of them are already in this series.&lt;/p&gt;

&lt;p&gt;The tokenizer is byte-level BPE, Part 1: every string encodes, so there is no unknown token. The normalisation is pre-LN, Part 5. The token table is tied to the output layer at &lt;code&gt;model.py:138&lt;/code&gt;. One matrix serves both ends: the table Part 5 counted as 31% of the model. Untied, GPT-2 small would have 163 million parameters instead of 124 million.&lt;/p&gt;

&lt;p&gt;The fourth is positions. GPT-2 learns them: &lt;code&gt;wpe&lt;/code&gt; is a table of shape &lt;code&gt;[1024, 768]&lt;/code&gt;, one row per position, added to the token rows before the first block. A table has a last row. Ask &lt;code&gt;forward&lt;/code&gt; for 1,100 positions and the assert at &lt;code&gt;model.py:173&lt;/code&gt; stops you:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight plaintext"&gt;&lt;code&gt;AssertionError: Cannot forward sequence of length 1100, block size is only 1024
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;The fifth is the initialisation. Each block makes two writes to the bus, Part 5's name for the residual stream, so a stack of &lt;code&gt;L&lt;/code&gt; blocks adds &lt;code&gt;2L&lt;/code&gt; of them. In nanoGPT, every matrix, the token and position tables included, starts as a random draw with standard deviation 0.02 (OpenAI's own code used 0.01 for the position table); biases start at zero and LayerNorm at one. &lt;code&gt;model.py:143-145&lt;/code&gt; starts the two layers that make those writes at &lt;code&gt;0.02/√(2L)&lt;/code&gt; instead, so the sum doesn't grow with depth. I measured it on untrained models. With the scaled init, the stream's norm grew about 11 times through 12 blocks, and 11 times through 48. With the plain init it grew 55 times and 112 times. Training takes it from there.&lt;/p&gt;

&lt;h2&gt;
  
  
  Training is one call
&lt;/h2&gt;

&lt;p&gt;Hand &lt;code&gt;forward&lt;/code&gt; a sequence and its targets, the same sequence shifted by one, and it predicts every position at once. I gave it 92 positions of text. It returned logits of shape &lt;code&gt;[1, 92, 50257]&lt;/code&gt;. Those are raw scores for every token at every position, and softmax turns them into probabilities. It also returned one loss, the cross-entropy at &lt;code&gt;model.py:187&lt;/code&gt;: minus the log of the probability given to the right next token, averaged over all 92. Part 4's mask is what makes this honest. Row &lt;code&gt;t&lt;/code&gt; sees only rows &lt;code&gt;0..t&lt;/code&gt;, so no prediction can read its own answer.&lt;/p&gt;

&lt;p&gt;That's the program training runs at every step, for hundreds of thousands of steps. Every position is known in advance, so all of them go through the stack together, as a few large matrix multiplications. If you've written a batch job over a table, you know the shape of this program. On my laptop's CPU, one forward over 1,024 positions took 0.23 seconds.&lt;/p&gt;

&lt;p&gt;Without targets, &lt;code&gt;forward&lt;/code&gt; takes a shortcut at &lt;code&gt;model.py:190&lt;/code&gt;: it computes logits for the last position only. That line is the first sign of the second program.&lt;/p&gt;

&lt;h2&gt;
  
  
  Generation is a loop
&lt;/h2&gt;

&lt;p&gt;To write text, the model has to see each token before it can predict the next. There is nothing to run in parallel across positions. Here is nanoGPT's &lt;code&gt;generate&lt;/code&gt;, &lt;code&gt;model.py:312-328&lt;/code&gt;, condensed:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;_&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;max_new_tokens&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="n"&gt;idx_cond&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;[:,&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1024&lt;/span&gt;&lt;span class="p"&gt;:]&lt;/span&gt;                 &lt;span class="c1"&gt;# crop to the position table
&lt;/span&gt;    &lt;span class="n"&gt;logits&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;_&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nf"&gt;self&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;idx_cond&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;                &lt;span class="c1"&gt;# the whole sequence, again
&lt;/span&gt;    &lt;span class="n"&gt;logits&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;logits&lt;/span&gt;&lt;span class="p"&gt;[:,&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;:]&lt;/span&gt; &lt;span class="o"&gt;/&lt;/span&gt; &lt;span class="n"&gt;temperature&lt;/span&gt;   &lt;span class="c1"&gt;# [B, 50257], last position only
&lt;/span&gt;    &lt;span class="n"&gt;probs&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;F&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;softmax&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;logits&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;dim&lt;/span&gt;&lt;span class="o"&gt;=-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;         &lt;span class="c1"&gt;# top_k masking skipped here
&lt;/span&gt;    &lt;span class="n"&gt;idx_next&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;multinomial&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;probs&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;    &lt;span class="c1"&gt;# draw one token id
&lt;/span&gt;    &lt;span class="n"&gt;idx&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;cat&lt;/span&gt;&lt;span class="p"&gt;((&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;idx_next&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="n"&gt;dim&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;   &lt;span class="c1"&gt;# append it and go round
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Read the &lt;code&gt;self(idx_cond)&lt;/code&gt; line again. Every step runs the full forward over everything so far and keeps one row. If you've re-read a whole file to append one line to it, you've written this loop. With an 11-token prompt, 256 new tokens push 35,456 positions through the stack for 256 useful rows.&lt;/p&gt;

&lt;p&gt;The fix is a KV cache. Part 4 gave every position a key, which later queries match against, and a value, which they average. Under Part 4's mask, the key and value at position &lt;code&gt;t&lt;/code&gt; depend only on positions &lt;code&gt;0..t&lt;/code&gt;, so once computed they never change. Keep them, one pair of tensors per block, and each step only pushes the new token through the stack and attends over the stored rows. Same 256 tokens: 266 positions. The price is memory. For GPT-2 small in fp32, that's about 74 kB per position. nanoGPT doesn't have one; I wrote one in about thirty lines on top of its modules.&lt;/p&gt;

&lt;p&gt;The timer agrees, but not in proportion. Generating 64, 256 and 512 tokens took 2.9, 5.0 and 7.4 times longer without the cache, for 37, 133 and 261 times the positions. At batch 1, each step still reads all 124 million weights to produce one token, and the cache doesn't shrink that. Writing 1,023 tokens with the cache took 9.3 seconds, about forty times the training forward over the same length.&lt;/p&gt;

&lt;p&gt;The cache also breaks at the edge of decision four. Past 1,024 tokens, &lt;code&gt;generate&lt;/code&gt; crops the left end. From the second block on, every stored key and value was computed while the dropped token was still there, and with learned positions every token also moves to a new row of &lt;code&gt;wpe&lt;/code&gt;. The whole cache is stale. nanoGPT's loop gets away with it because it recomputes everything anyway; a cached loop has to handle it.&lt;/p&gt;

&lt;p&gt;The original post has a widget for this that doesn't survive being pasted into a feed: &lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-transformer-plus-five-decisions/" rel="noopener noreferrer"&gt;drag the new tokens&lt;/a&gt; and watch the two counts part. Then pick XL in bf16 and raise the batch until the cache bar passes the weights.&lt;/p&gt;

&lt;h2&gt;
  
  
  Same weights, different bugs
&lt;/h2&gt;

&lt;p&gt;Two programs over the same weights should agree. Mine did, almost. Over 200 greedy steps, always taking the highest-scoring token, both loops wrote the same text. Their logits differed by up to 0.011. Almost all of that was the same shift for every token, and softmax ignores it. The gap between the two best candidates moved by at most 0.00005; the narrowest gap at any step was 0.004, so the tokens held. A longer run or a bigger model can close that margin, and nothing will raise an error when it does.&lt;/p&gt;

&lt;p&gt;The difference is floating point. After the prompt, every matrix multiplication inside the cached loop's blocks sees one row instead of the whole sequence, so the library picks different kernels, the routines that do the arithmetic, and those add in a different order. Float addition isn't associative, so the sums differ in the last bits.&lt;/p&gt;

&lt;p&gt;Then change something your deployment changes routinely. I cast the same weights to bf16, Part 3's 16-bit format, and ran the same greedy loop. The fp32 model continued the prompt with " were not working."; the bf16 one with " were not being run on the same system". They split at the third token and never came back together. Running the prompt eight times in one batch kept the tokens but moved the logits by up to 0.005 against the single-prompt run.&lt;/p&gt;

&lt;p&gt;The usual fix people reach for is temperature 0. It doesn't fix this, and in nanoGPT it isn't even a value. &lt;code&gt;generate&lt;/code&gt; divides the logits by the temperature, so 0 turns all 50,257 logits into infinities, softmax turns those into NaN, the float for an undefined result, and &lt;code&gt;torch.multinomial&lt;/code&gt; raises. Most APIs that accept 0 treat it as a switch to argmax, picking the top score. That's greedy decoding again. It's only as deterministic as the weights, dtype, batch shape and kernels under it.&lt;/p&gt;

&lt;p&gt;Be fair about the scale of this. Every run here was repeatable: the same seed on the same machine sampled the same 40 tokens twice. The drift appears when something underneath changes, and in a serving stack those things change without asking you.&lt;/p&gt;

&lt;h2&gt;
  
  
  Sampling is a policy on top
&lt;/h2&gt;

&lt;p&gt;Greedy decoding is a choice, and GPT-2 small shows why people avoid it: within 200 tokens, the loop above was repeating its own sentence. Sampling draws from the distribution instead, and three knobs shape it before the draw.&lt;/p&gt;

&lt;p&gt;Temperature divides the logits before softmax; below 1 it sharpens the distribution, above 1 it flattens it. Top-k keeps the &lt;code&gt;k&lt;/code&gt; highest-scoring tokens and drops the rest. Top-p, also called nucleus sampling, keeps the smallest set whose probabilities add up to at least &lt;code&gt;p&lt;/code&gt;. nanoGPT's &lt;code&gt;sample.py&lt;/code&gt; defaults to temperature 0.8 and top-k 200; it has no top-p.&lt;/p&gt;

&lt;p&gt;What the knobs do depends on the prompt. After "The capital of France is", GPT-2 small's top candidate is " the", at 8.5%; " Paris" is fifth, at 3.2%. At temperature 1, covering 90% of the probability takes 1,504 tokens. After "import numpy as", " np" alone has 85%, and 90% takes three. Entropy, measured in bits, puts a number on the spread: 8.7 for the capital prompt at temperature 1, as uncertain as a fair pick among about 400 tokens, and 1.5 for numpy, about three.&lt;/p&gt;

&lt;p&gt;Top-k can't tell those apart, so any single &lt;code&gt;k&lt;/code&gt; you pick is wrong for one of them. At nanoGPT's defaults, those 200 tokens hold 89% of the mass after the capital prompt. After the numpy prompt they keep 199 options that share under 5%. Top-p adapts: at the same temperature and &lt;code&gt;p = 0.9&lt;/code&gt;, it keeps 237 tokens for one and a single token for the other.&lt;/p&gt;

&lt;p&gt;The second widget on &lt;a href="https://cchinchilla.dev/posts/fcw-gpt2-transformer-plus-five-decisions/" rel="noopener noreferrer"&gt;the original post&lt;/a&gt; starts at temperature 1 and &lt;code&gt;p = 0.9&lt;/code&gt;, where the capital prompt keeps 1,504 tokens and numpy three. Drag the temperature to 2 and watch numpy's count grow. The top-40 and top-200 readouts show what a fixed &lt;code&gt;k&lt;/code&gt; keeps.&lt;/p&gt;

&lt;h2&gt;
  
  
  What the ML engineer watches
&lt;/h2&gt;

&lt;p&gt;Four habits, and the first is a test.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Test the cached loop against the full forward, on logits.&lt;/strong&gt; Run both programs for a few hundred steps and compare the greedy tokens and, with a tolerance, the logits. Token equality alone can pass while the logits drift, as mine did.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Budget the cache like the weights.&lt;/strong&gt; Per sequence it holds &lt;code&gt;2·L·T·d&lt;/code&gt; numbers for &lt;code&gt;L&lt;/code&gt; blocks, &lt;code&gt;T&lt;/code&gt; positions and width &lt;code&gt;d&lt;/code&gt;: a key and a value for every block and position. GPT-2 small at 1,024 positions in fp32 is 75 MB. GPT-2 XL in bf16 is 315 MB per sequence, so about ten full-length sequences weigh as much as its 3.1 GB of weights. When you raise the batch size, you are making a memory decision, not only a throughput one.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Pin what you compare.&lt;/strong&gt; Dtype, batch shape, library version and kernel all move logits. If your regression test diffs model outputs, fix all of them, or compare with a tolerance and log what changed.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Log the sampling settings with every output.&lt;/strong&gt; Temperature, top-k, top-p and the seed are part of the program. Two runs that differ only in them behave differently, and an output logged without them can't be reproduced.&lt;/p&gt;

&lt;p&gt;None of this is modelling work either. The weights are one artifact. The code around them decides whether a position is scored or written, what survives between steps, and how a distribution becomes a token, and each of those is ordinary code you can test.&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Part 6 of twelve. GPT-2 is Part 5's stack plus five decisions: byte-level BPE, pre-LN, a tied token table, learned positions, and a scaled init for the writes to the bus. Training runs one parallel forward over every position; generation runs a loop that needs a cache, a sampling policy and its own tests. The next part reads what Llama changed, and lands on &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;the series index&lt;/a&gt; when it's out. If you arrived here directly, &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Part 1&lt;/a&gt; reads the tokenizer, &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;Part 2&lt;/a&gt; the table a token id indexes, &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;Part 3&lt;/a&gt; the shapes, &lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;Part 4&lt;/a&gt; attention, and &lt;a href="https://cchinchilla.dev/posts/fcw-transformer-six-shapes-residual-stream/" rel="noopener noreferrer"&gt;Part 5&lt;/a&gt; the block.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Written with AI assistance. See &lt;a href="https://cchinchilla.dev/ai-policy/" rel="noopener noreferrer"&gt;AI policy&lt;/a&gt;.&lt;/em&gt;&lt;/p&gt;

</description>
      <category>machinelearning</category>
      <category>python</category>
      <category>ai</category>
      <category>beginners</category>
    </item>
    <item>
      <title>A transformer block is six tensors and a bus: read it like a pipeline</title>
      <dc:creator>Carlos Chinchilla Corbacho</dc:creator>
      <pubDate>Sun, 04 Oct 2026 16:00:00 +0000</pubDate>
      <link>https://dev.to/cchinchilladev/a-transformer-block-is-six-tensors-and-a-bus-read-it-like-a-pipeline-34a8</link>
      <guid>https://dev.to/cchinchilladev/a-transformer-block-is-six-tensors-and-a-bus-read-it-like-a-pipeline-34a8</guid>
      <description>&lt;p&gt;&lt;em&gt;Originally published at &lt;a href="https://cchinchilla.dev/posts/fcw-transformer-six-shapes-residual-stream/" rel="noopener noreferrer"&gt;cchinchilla.dev&lt;/a&gt;. Part 5 of From code to weights, a 12-part series on ML fundamentals for engineers.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;Part 4&lt;/a&gt; ended on one stage. A decoder is a stack of blocks, twelve in GPT-2 small, and attention is one of two stages in each. This post reads a whole block the way you'd read a pipeline: what flows in, what each stage writes, and where the parameters and the cost end up.&lt;/p&gt;

&lt;p&gt;They don't end up in the same place.&lt;/p&gt;

&lt;p&gt;I read &lt;code&gt;Block&lt;/code&gt; in nanoGPT's &lt;code&gt;model.py&lt;/code&gt; with the shapes in the margin and loaded GPT-2's weights into it. Then I broke it on purpose: I took blocks out one at a time, then made each one overwrite its input instead of adding to it.&lt;/p&gt;

&lt;p&gt;This is part of &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;a series&lt;/a&gt; for engineers moving into ML engineering. Each part names one thing that does not behave the way software behaves. This one is where the cost lives: not where the code is.&lt;/p&gt;

&lt;h2&gt;
  
  
  The block is four lines
&lt;/h2&gt;

&lt;p&gt;nanoGPT is a minimal GPT-2 in two files. Its README now points to a successor, but &lt;code&gt;model.py&lt;/code&gt; is still the shortest honest read of a GPT-2 block. &lt;code&gt;Block.forward&lt;/code&gt; is &lt;code&gt;model.py:103-106&lt;/code&gt;:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;forward&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;                   &lt;span class="c1"&gt;# x: [B, T, 768], the stream
&lt;/span&gt;    &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;attn&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;ln_1&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;     &lt;span class="c1"&gt;# read a normalised copy, add attention's output
&lt;/span&gt;    &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;mlp&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;ln_2&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;      &lt;span class="c1"&gt;# read again, add the feed-forward's output
&lt;/span&gt;    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;                            &lt;span class="c1"&gt;# [B, T, 768]: the shape it came in with
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Inside the two calls, &lt;code&gt;model.py:52-92&lt;/code&gt;, six tensors are worth naming. The code below is condensed, without scaling, masking, or dropout. The shapes in the margin come from running a model with GPT-2 small's shape, &lt;code&gt;d = 768&lt;/code&gt; and 12 heads:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;qkv&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;attn&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;c_attn&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;ln_1&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;         &lt;span class="c1"&gt;# [B, T, 2304]    q, k, v side by side
&lt;/span&gt;&lt;span class="n"&gt;q&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;k&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;v&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;qkv&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;split&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;768&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;dim&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;              &lt;span class="c1"&gt;# three views, [B, T, 768] each
&lt;/span&gt;&lt;span class="n"&gt;q&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;q&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;view&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;B&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;T&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;12&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;64&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;transpose&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;     &lt;span class="c1"&gt;# [B, 12, T, 64]  12 heads of 64; k, v alike
&lt;/span&gt;&lt;span class="n"&gt;att&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;q&lt;/span&gt; &lt;span class="o"&gt;@&lt;/span&gt; &lt;span class="n"&gt;k&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;transpose&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)).&lt;/span&gt;&lt;span class="nf"&gt;softmax&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;  &lt;span class="c1"&gt;# [B, 12, T, T]   Part 4's square
&lt;/span&gt;&lt;span class="n"&gt;y&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;att&lt;/span&gt; &lt;span class="o"&gt;@&lt;/span&gt; &lt;span class="n"&gt;v&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;transpose&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;reshape&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;B&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;T&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;768&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;y&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;attn&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;c_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;                      &lt;span class="c1"&gt;# [B, T, 768]     attention out
&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;                                    &lt;span class="c1"&gt;# [B, T, 768]     the stream, attention added
&lt;/span&gt;&lt;span class="n"&gt;h&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;F&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;gelu&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;mlp&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;c_fc&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;ln_2&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;)))&lt;/span&gt;      &lt;span class="c1"&gt;# [B, T, 3072]    the feed-forward's hidden layer
&lt;/span&gt;&lt;span class="n"&gt;out&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;mlp&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;c_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;h&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;                     &lt;span class="c1"&gt;# [B, T, 768]     feed-forward out
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;With the stream itself, that's the block: &lt;code&gt;x&lt;/code&gt;, the projections, the scores, attention's output, the hidden layer, and the feed-forward's output. Three of them share the stream's shape. Everything else is a view, an elementwise step, or the per-head &lt;code&gt;[B, 12, T, 64]&lt;/code&gt; layout.&lt;/p&gt;

&lt;p&gt;One division of labour is worth reading twice. Attention is the only stage that mixes positions: row &lt;code&gt;t&lt;/code&gt; of its output is built from rows &lt;code&gt;0..t&lt;/code&gt; of the stream. LayerNorm and the feed-forward work on one position at a time.&lt;/p&gt;

&lt;p&gt;&lt;code&gt;LayerNorm&lt;/code&gt; subtracts each 768-number row's mean and divides by its standard deviation, then applies a learned scale and shift; &lt;code&gt;model.py:27&lt;/code&gt; is one call to &lt;code&gt;F.layer_norm&lt;/code&gt;. The feed-forward widens each row to 3,072, applies GELU (a smooth version of &lt;code&gt;max(0, x)&lt;/code&gt;), and narrows it back. If you've written a map over rows, you've written the feed-forward's control flow.&lt;/p&gt;

&lt;p&gt;The original post has a widget for this that doesn't survive being pasted into a feed: &lt;a href="https://cchinchilla.dev/posts/fcw-transformer-six-shapes-residual-stream/" rel="noopener noreferrer"&gt;pick the feed-forward hidden layer&lt;/a&gt; and read its parameters, then pick the scores, which have none. Drag &lt;code&gt;T&lt;/code&gt; and the parameters don't move while the arithmetic does. That split is the whole point of the last two sections.&lt;/p&gt;

&lt;h2&gt;
  
  
  The plus sign is the architecture
&lt;/h2&gt;

&lt;p&gt;Read the two lines again. Neither stage replaces &lt;code&gt;x&lt;/code&gt;. Each reads a normalised copy and adds what it computed to the original. The stream that enters a block leaves it with two edits on top, and the next block reads the result.&lt;/p&gt;

&lt;p&gt;That stream, &lt;code&gt;[B, T, 768]&lt;/code&gt; from the embedding to the last block, is what the literature calls the residual stream; this post calls it the bus. Every stage reads it and writes to it; no stage owns it. If you've written &lt;code&gt;total += delta&lt;/code&gt; in a loop, you know the contract: every stage adds to the running total, and none can overwrite it; undoing an earlier write takes a later stage adding its negative.&lt;/p&gt;

&lt;p&gt;The numbers show the contract holding. I ran 1,024 tokens of &lt;a href="https://cchinchilla.dev/posts/llm-as-judge-variance-problem/" rel="noopener noreferrer"&gt;an earlier post&lt;/a&gt; through GPT-2 small and measured lengths, a vector's norm, as medians over positions. In blocks 1 to 10, each stage adds a vector 16% to 40% as long as the stream it reads. The stream that comes out points almost where it went in: the cosine similarity, where 1 means the same direction, is 0.93 to 0.97. The norm also grows as it goes, from 49 after block 0 to 480 at the end. Nothing on the bus is ever normalised in place. Each stage normalises its own copy, and &lt;code&gt;ln_f&lt;/code&gt; normalises once more before the output layer turns the last row into logits, the next-token scores.&lt;/p&gt;

&lt;p&gt;That was a decision. Call GPT-2's order pre-LN: normalise inside the branch, &lt;code&gt;x + attn(ln(x))&lt;/code&gt;, and leave the bus alone. The original transformer, and GPT-1, were post-LN: normalise after the add, &lt;code&gt;x = ln(x + attn(x))&lt;/code&gt;, which rescales the bus itself at every stage. GPT-2's paper records the move in one sentence, the same one that adds &lt;code&gt;ln_f&lt;/code&gt;, and cites a precedent but gives no reason of its own.&lt;/p&gt;

&lt;p&gt;A small test shows one. I trained two six-block models on this blog's text, one character at a time, with the same settings and no warmup (warmup: starting with a small step size and ramping it up). Loss is the model's average surprise at each next token; lower is better, and 0 would mean certain and right.&lt;/p&gt;

&lt;p&gt;After 300 steps the pre-LN model reached about 2.4, better than the 2.56 a model gets from knowing only the previous letter. The post-LN one stalled at 3.2, roughly what letter frequencies alone give, and one seed run to 600 steps was still there. Given a 100-step warmup, it broke loose between steps 200 and 250. It's a toy, with two seeds and losses read off training batches, but the direction matches &lt;a href="https://arxiv.org/abs/2002.04745" rel="noopener noreferrer"&gt;the published analysis&lt;/a&gt; of why post-LN needs a warmup.&lt;/p&gt;

&lt;p&gt;Back to GPT-2 small: I deleted blocks. The intact model scores a loss of 3.59 on those 1,024 tokens. Remove any one of blocks 1 to 10, so the stream passes straight through, and the loss comes out between 3.61 and 3.89. Remove block 0 and it jumps to 9.15.&lt;/p&gt;

&lt;p&gt;Overwriting is worse than removing. Keep a block, but let its output replace the stream instead of adding to it. For blocks 1 to 8 the loss is above 9.6. For blocks 2 to 7 it is above 10.8, the score of a uniform guess over the vocabulary. Worse than guessing.&lt;/p&gt;

&lt;p&gt;The second widget on &lt;a href="https://cchinchilla.dev/posts/fcw-transformer-six-shapes-residual-stream/" rel="noopener noreferrer"&gt;the original post&lt;/a&gt; plots that: switch it to overwriting and blocks 2 to 7 cross the dotted line that marks a uniform guess. Then pick block 0, which barely moves.&lt;/p&gt;

&lt;p&gt;Block 0 is the exception, and the norms say why. The embedding enters with a norm of 4.6, and block 0's attention writes a vector seven times that length. After block 0 the stream is mostly block 0's output, so it behaves less like an edit than like a second embedding. Overwriting there costs nothing, 3.59, because adding already amounted to replacing.&lt;/p&gt;

&lt;h2&gt;
  
  
  Two thirds of the weights
&lt;/h2&gt;

&lt;p&gt;You can count one block's parameters by hand, with &lt;code&gt;d = 768&lt;/code&gt; and the feed-forward width &lt;code&gt;f = 4d&lt;/code&gt;. Attention's two projections are &lt;code&gt;4d² + 4d&lt;/code&gt;: 2,362,368 numbers. The feed-forward's two are &lt;code&gt;2·d·f + f + d&lt;/code&gt;, which is &lt;code&gt;8d² + 5d&lt;/code&gt;: 4,722,432. The two LayerNorms are 3,072. The feed-forward holds two thirds of every block in every GPT-2 size, because the ratio comes from the &lt;code&gt;4 * config.n_embd&lt;/code&gt; at &lt;code&gt;model.py:82&lt;/code&gt; and not from the width. Llama 3.1 8B changes both stages, and its feed-forward holds 81% of a block; Part 7 reads why.&lt;/p&gt;

&lt;p&gt;Twelve blocks are 85 million parameters, and GPT-2 small has 124,439,808. The rest is almost all the token table: 38.6 million numbers, 31% of the model. It's used twice, as Part 2's lookup at the bottom and as the projection to logits at the top; &lt;code&gt;model.py:138&lt;/code&gt; ties them, the same array and not a copy. In GPT-2 XL the same table is 5.2% of the model. Small models spend nearly a third of their weights on the vocabulary; large ones barely notice it.&lt;/p&gt;

&lt;h2&gt;
  
  
  The cost is somewhere else
&lt;/h2&gt;

&lt;p&gt;In software, a function's size and its running time come from the same code. Here they come apart. The parameters sit mostly in the feed-forward. The arithmetic follows &lt;code&gt;T&lt;/code&gt;.&lt;/p&gt;

&lt;p&gt;A forward pass costs about two operations per token for each number in a weight matrix, a multiply and an add; Part 4 counted them as FLOPs. Call the count of those numbers &lt;code&gt;N&lt;/code&gt;, so the parameters cost &lt;code&gt;2N&lt;/code&gt; per token. Add the attention term, &lt;code&gt;4·L·T·d&lt;/code&gt; with &lt;code&gt;L&lt;/code&gt; the number of blocks, for the scores and the weighted sum over the full square. For GPT-2 small at 1,024 tokens that's 247 million operations per token from the parameters and 38 million from attention. torch's &lt;code&gt;FlopCounterMode&lt;/code&gt; agrees when attention is written out, Part 4's eager path: 284.8 million per token, exactly the formula's figure (285.1 if &lt;code&gt;N&lt;/code&gt; is nanoGPT's &lt;code&gt;get_num_params()&lt;/code&gt;, which also counts the biases and the LayerNorms).&lt;/p&gt;

&lt;p&gt;On sdpa, the default, the same counter reported 247 million at both 256 and 1,024 tokens. In torch 2.14, the version I ran, it has no entry for the fused CPU attention kernel, the single routine sdpa calls instead of the four lines. The attention term never reaches its total, and nothing tells you so. PyTorch's main branch now registers that kernel (&lt;a href="https://github.com/pytorch/pytorch/pull/195801" rel="noopener noreferrer"&gt;pull request #195801&lt;/a&gt;), so the next feature release should count the attention term on sdpa too.&lt;/p&gt;

&lt;p&gt;Inside one block, the attention term is 18% of the arithmetic at 1,024 tokens, and it passes the linear layers at 4,608 tokens, six times &lt;code&gt;d&lt;/code&gt;. The clock agrees. On my laptop's CPU, one block's attention stage, projections included, took 38% of the block's time at 256 tokens, 47% at 1,024, and 57% at 4,096, while holding a third of its parameters.&lt;/p&gt;

&lt;p&gt;Be fair about the split. At the lengths GPT-2 was built for, the linear layers do most of the arithmetic, and the feed-forward is the larger half of that. The attention term only takes over past a length GPT-2 never accepts. Still, when you pick a model you pick two budgets. The parameters set the memory and are fixed at load; the context sets the cost and a growing share of the memory, and it arrives with each request.&lt;/p&gt;

&lt;h2&gt;
  
  
  What the ML engineer watches
&lt;/h2&gt;

&lt;p&gt;Four habits, and the first two are arithmetic.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Count parameters by component, not in total.&lt;/strong&gt; &lt;code&gt;4d²&lt;/code&gt; for attention, &lt;code&gt;2·d·f&lt;/code&gt; for the feed-forward, per block. Add the vocabulary times &lt;code&gt;d&lt;/code&gt;, once, for the table. If you compare a 124M model with a 1.5B one, ask how much of each is vocabulary.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Estimate cost as 2N per token, and 6N to train.&lt;/strong&gt; Training is a forward and a backward pass, and the backward, which computes the gradients, costs about two forwards. That's 6N per token and &lt;code&gt;6·N·D&lt;/code&gt; for a run of &lt;code&gt;D&lt;/code&gt; tokens. nanoGPT's config for GPT-2, &lt;code&gt;config/train_gpt2.py&lt;/code&gt;, runs 491,520 tokens per iteration for 600,000 iterations: 295 billion tokens. &lt;code&gt;6·N·D&lt;/code&gt; is 2.2 × 10²⁰ operations, 195 hours of one A100 at its bf16 peak, a day of an eight-GPU node. The file says to expect about five days and the README four, which puts the run at a fifth to a quarter of peak. nanoGPT's own estimator adds the attention term, &lt;code&gt;model.py:296&lt;/code&gt;, and at 1,024 tokens that term is another 15%.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Depth is latency.&lt;/strong&gt; Spend the same 85 million block parameters four ways: 6 blocks of width 1,088, 12 of 768, 24 of 544, and 48 of 384. At batch 1, one token through the stack took 3.2, 4.0, 4.7, and 6.7 ms on my CPU. Each block adds a fixed overhead, and blocks run one after another. If you serve one request at a time, depth is what you wait for, even when the arithmetic is equal.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Measure a block before you drop it.&lt;/strong&gt; Layer pruning (shipping with some blocks removed) and early exit (stopping at a middle block when the answer looks settled) lean on the bus, and the bus is uneven. Removing a middle block of GPT-2 small raises its loss by 0.5% to 8%; removing block 0 raises it by 155%, and the last block by 41%. The numbers above took a minute on a laptop, on text from this blog. Produce them on your own traffic.&lt;/p&gt;

&lt;p&gt;None of this is modelling work either. A block is four lines around a sum, and the sum is what lets a stack of twelve behave like one pipeline. Where the weights sit and where the time goes are two different answers, and both are arithmetic you can do before loading anything.&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Part 5 of twelve. A block reads the stream through two stages, attention and a feed-forward, and adds both results back: the plus sign is why removing a middle block costs a few percent of the loss rather than the model. The feed-forward holds two thirds of each block's weights; the attention term's share of the arithmetic grows with the context. The next part reads the five decisions GPT-2 made around the stack, and lands on &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;the series index&lt;/a&gt; when it's out. If you arrived here directly, &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Part 1&lt;/a&gt; reads the tokenizer, &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;Part 2&lt;/a&gt; the table a token id indexes, &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;Part 3&lt;/a&gt; the shapes, and &lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;Part 4&lt;/a&gt; the stage this block is built around.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Written with AI assistance. See &lt;a href="https://cchinchilla.dev/ai-policy/" rel="noopener noreferrer"&gt;AI policy&lt;/a&gt;.&lt;/em&gt;&lt;/p&gt;

</description>
      <category>machinelearning</category>
      <category>python</category>
      <category>ai</category>
      <category>beginners</category>
    </item>
    <item>
      <title>Attention is a learned weighted average, and the cost is in the square</title>
      <dc:creator>Carlos Chinchilla Corbacho</dc:creator>
      <pubDate>Fri, 25 Sep 2026 13:00:00 +0000</pubDate>
      <link>https://dev.to/cchinchilladev/attention-is-a-learned-weighted-average-and-the-cost-is-in-the-square-2e0o</link>
      <guid>https://dev.to/cchinchilladev/attention-is-a-learned-weighted-average-and-the-cost-is-in-the-square-2e0o</guid>
      <description>&lt;p&gt;&lt;em&gt;Originally published at &lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;cchinchilla.dev&lt;/a&gt;. Part 4 of From code to weights, a 12-part series on ML fundamentals for engineers.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;Part 3&lt;/a&gt; ended on a square. Two tensors in one decoder block carry &lt;code&gt;T&lt;/code&gt; twice, &lt;code&gt;[B, H, T, T]&lt;/code&gt;: a &lt;code&gt;T × T&lt;/code&gt; square per head, &lt;code&gt;H&lt;/code&gt; heads, &lt;code&gt;B&lt;/code&gt; sequences. At the widget's defaults they were 89% of the block's memory. This post reads what fills them. It's a weighted average. The average is ordinary; the weights are computed from the input, and that is the new part.&lt;/p&gt;

&lt;p&gt;I wrote attention in four lines and checked it against &lt;code&gt;F.scaled_dot_product_attention&lt;/code&gt;, torch's own implementation. Then I read &lt;code&gt;LlamaAttention.forward&lt;/code&gt; with the shapes in the margin, and priced it for Llama 3.1 8B at the context lengths people promise.&lt;/p&gt;

&lt;p&gt;This is part of &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;a series&lt;/a&gt; for engineers moving into ML engineering. Each part names one thing that does not behave the way software behaves. This one is cost: you cannot read it from the code.&lt;/p&gt;

&lt;h2&gt;
  
  
  A dictionary that matches a little
&lt;/h2&gt;

&lt;p&gt;A dictionary lookup is exact. The key matches or it doesn't, and you get one value back. Attention is the same operation with the match made soft. Every key matches a little, and you get back the average of all the values, weighted by how well each key matched.&lt;/p&gt;

&lt;p&gt;Those are attention weights, not parameters. They aren't stored anywhere; the model computes them from the input on every call. Until the cost section, "weights" means these. Three learned matrices turn each input row into a query, a key, and a value. The query is compared against every key, and the scores become weights:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;q&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;k&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;v&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;@&lt;/span&gt; &lt;span class="n"&gt;Wq&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;@&lt;/span&gt; &lt;span class="n"&gt;Wk&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;@&lt;/span&gt; &lt;span class="n"&gt;Wv&lt;/span&gt;             &lt;span class="c1"&gt;# [B, T, d] each
&lt;/span&gt;&lt;span class="n"&gt;scores&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;q&lt;/span&gt; &lt;span class="o"&gt;@&lt;/span&gt; &lt;span class="n"&gt;k&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;transpose&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;/&lt;/span&gt; &lt;span class="nf"&gt;sqrt&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;d&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;  &lt;span class="c1"&gt;# [B, T, T]
&lt;/span&gt;&lt;span class="n"&gt;weights&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;scores&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;softmax&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;                &lt;span class="c1"&gt;# every row sums to 1
&lt;/span&gt;&lt;span class="n"&gt;out&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;weights&lt;/span&gt; &lt;span class="o"&gt;@&lt;/span&gt; &lt;span class="n"&gt;v&lt;/span&gt;                           &lt;span class="c1"&gt;# [B, T, d]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Read the shapes. &lt;code&gt;x&lt;/code&gt; is Part 3's &lt;code&gt;[B, T, d]&lt;/code&gt;. The three projections keep it. &lt;code&gt;scores&lt;/code&gt; is &lt;code&gt;[B, T, T]&lt;/code&gt;: one number for every pair of positions, which is the square. The softmax turns each row into weights that sum to one, and the last line is the average. The four lines agree with &lt;code&gt;F.scaled_dot_product_attention&lt;/code&gt; to &lt;code&gt;1e-5&lt;/code&gt;.&lt;/p&gt;

&lt;p&gt;Two properties follow from "average". Each output coordinate sits between the smallest and largest value at that coordinate; attention never leaves that range. And if you replace a softmax row with a one-hot row, a single 1 and zeros elsewhere, the output is exactly the value at that position. The soft dictionary contains the hard one as a limit.&lt;/p&gt;

&lt;h2&gt;
  
  
  Why the square root
&lt;/h2&gt;

&lt;p&gt;The division by &lt;code&gt;sqrt(d)&lt;/code&gt; looks like a detail. It's the difference between a model that trains from the first step and one that starts saturated. The reason is arithmetic you can check in a shell.&lt;/p&gt;

&lt;p&gt;A dot product of two random vectors of width &lt;code&gt;d&lt;/code&gt;, with unit-variance entries, has variance &lt;code&gt;d&lt;/code&gt;. I measured it: 4,104 at &lt;code&gt;d = 4096&lt;/code&gt;. Feed scores that wide into a softmax and it saturates: one weight goes to nearly 1 and the rest to nearly 0. With 1,024 keys, the largest weight in a row averaged 0.976. The row had 0.09 bits of entropy out of a possible 10. That's a hard lookup wearing a soft one's clothes, at initialisation; trained models sharpen rows on purpose, from a start they can learn from. Divide by &lt;code&gt;sqrt(d)&lt;/code&gt; and the variance is 1.0, the largest weight 0.017, the entropy 9.27 bits. Averaged over the scores, the gradient that reaches them is about an order of magnitude larger: 12 to 18 times across seeds, with the loss read through the average. Through a single weight the ratio swings from a few times to many orders of magnitude, depending on which rows that key happens to win. A saturated softmax has almost nothing to learn from: its gradient is under &lt;code&gt;1e-6&lt;/code&gt; for 99.8% of the scores, and what's left piles onto the few near-ties.&lt;/p&gt;

&lt;p&gt;Llama writes it once: &lt;code&gt;self.scaling = self.head_dim**-0.5&lt;/code&gt;, at &lt;code&gt;modeling_llama.py:226&lt;/code&gt;. The width that matters is the head's, not the model's.&lt;/p&gt;

&lt;h2&gt;
  
  
  Heads are a view, the mask is a triangle
&lt;/h2&gt;

&lt;p&gt;Multi-head attention is the four lines above run once per head, &lt;code&gt;H&lt;/code&gt; times, on slices of &lt;code&gt;d&lt;/code&gt;. The slicing is free. &lt;code&gt;q.view(B, T, H, d // H).transpose(1, 2)&lt;/code&gt; gives &lt;code&gt;[B, H, T, d/H]&lt;/code&gt; and shares memory with &lt;code&gt;q&lt;/code&gt;. The head axis is a stride. Scores are then &lt;code&gt;[B, H, T, T]&lt;/code&gt;, one square per head. Head 0 of the batched call equals attention run on head 0 alone. What the heads buy is &lt;code&gt;H&lt;/code&gt; different sets of weights over the same positions. Each lives in a slice of width &lt;code&gt;d/H&lt;/code&gt;: 128 in Llama 3.1 8B, with 32 heads.&lt;/p&gt;

&lt;p&gt;The mask is where a decoder stops being an average over everything. To predict the token after position &lt;code&gt;t&lt;/code&gt;, the model may read positions &lt;code&gt;0..t&lt;/code&gt; and nothing after, so the scores above the diagonal are set to &lt;code&gt;-inf&lt;/code&gt; before the softmax and their weights come out as zero. I tested it the direct way: attention over eight positions, then rewrite the keys and values at positions 5 to 7. Positions 0 to 4 returned the same output to &lt;code&gt;1e-6&lt;/code&gt;. Without the mask, all eight changed.&lt;/p&gt;

&lt;p&gt;The original post has a widget for this that doesn't survive being pasted into a feed: &lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;switch between the three masks&lt;/a&gt; and read the live-cell count change. The third chip, the window, is the last lever in the cost section.&lt;/p&gt;

&lt;p&gt;The triangle composes with Part 3's padding mask by addition. Left-pad a sequence by two and its first two query rows are all &lt;code&gt;-inf&lt;/code&gt;; a row of all &lt;code&gt;-inf&lt;/code&gt; is &lt;code&gt;0/0&lt;/code&gt; in the softmax, NaN. Same trap, one part later, in the four-line form. transformers fills its eager mask with the dtype's minimum instead (&lt;code&gt;masking_utils.py:602-604&lt;/code&gt;); those rows come out uniform. Quiet, and still meaningless.&lt;/p&gt;

&lt;p&gt;The eager path computes the exponentials and their sum in fp32 whatever the model's dtype, then casts the row back: &lt;code&gt;dtype=torch.float32&lt;/code&gt; at &lt;code&gt;modeling_llama.py:208&lt;/code&gt;. A normaliser summed in bf16's eight significant bits is the wrong place to save memory.&lt;/p&gt;

&lt;h2&gt;
  
  
  Reading &lt;code&gt;LlamaAttention.forward&lt;/code&gt;
&lt;/h2&gt;

&lt;p&gt;&lt;code&gt;modeling_llama.py:243-281&lt;/code&gt; in transformers 5.17.0 is the four lines with the heads, the mask, and a cache written out. Condensed, with shapes for the 8B model:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;query_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;q_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;view&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_shape&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;transpose&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;  &lt;span class="c1"&gt;# [B, 32, T, 128]
&lt;/span&gt;&lt;span class="n"&gt;key_states&lt;/span&gt;   &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;k_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;view&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_shape&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;transpose&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;  &lt;span class="c1"&gt;# [B, 8, T, 128]
&lt;/span&gt;&lt;span class="n"&gt;value_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;v_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_states&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;view&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;hidden_shape&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;transpose&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;  &lt;span class="c1"&gt;# [B, 8, T, 128]
&lt;/span&gt;&lt;span class="n"&gt;query_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;key_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nf"&gt;apply_rotary_pos_emb&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;query_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;key_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;cos&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;sin&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="k"&gt;if&lt;/span&gt; &lt;span class="n"&gt;past_key_values&lt;/span&gt; &lt;span class="ow"&gt;is&lt;/span&gt; &lt;span class="ow"&gt;not&lt;/span&gt; &lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
    &lt;span class="n"&gt;key_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;value_states&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;past_key_values&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;update&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;key_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;value_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;layer_idx&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;attn_output&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;attn_weights&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nf"&gt;attention_interface&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;query_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;key_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;value_states&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
                                                &lt;span class="n"&gt;attention_mask&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;scaling&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;scaling&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;attn_output&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;o_proj&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;attn_output&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;reshape&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;*&lt;/span&gt;&lt;span class="n"&gt;input_shape&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Three things in it aren't in the four lines. &lt;code&gt;past_key_values.update&lt;/code&gt; is the cache. At generation time the keys and values of every earlier token are kept, so each new token computes one query row and reads &lt;code&gt;T&lt;/code&gt; stored keys.&lt;/p&gt;

&lt;p&gt;The key and value have 8 heads where the query has 32: &lt;code&gt;k_proj&lt;/code&gt; maps 4,096 to 1,024. That's grouped-query attention (GQA), against one key-value head per query head (MHA). It exists for the cache. And &lt;code&gt;attention_interface&lt;/code&gt; is a dispatch whose fallback is &lt;code&gt;eager_attention_forward&lt;/code&gt;, &lt;code&gt;modeling_llama.py:191-213&lt;/code&gt;: the name means the four lines run as written, plus &lt;code&gt;repeat_kv&lt;/code&gt;, which copies the 8 key heads to 32 at runtime so the matmul lines up. Set &lt;code&gt;_attn_implementation&lt;/code&gt; to &lt;code&gt;sdpa&lt;/code&gt;, the &lt;code&gt;scaled_dot_product_attention&lt;/code&gt; from the first section, or to &lt;code&gt;flash_attention_2&lt;/code&gt;, and a different kernel (the routine that runs the arithmetic) computes the same maths. The rotary line is Part 7's.&lt;/p&gt;

&lt;h2&gt;
  
  
  The cost is in the data, not the code
&lt;/h2&gt;

&lt;p&gt;In software you read a function's cost from its code. A loop over the input is linear; two nested loops are quadratic; you can tell before it runs. Attention is one function; its code is the four lines. There is no loop in them. The square is inside a single &lt;code&gt;@&lt;/code&gt;, and nothing in the code's shape says quadratic. Its cost is set by &lt;code&gt;T&lt;/code&gt; in the request: positions squared, per head, per layer.&lt;/p&gt;

&lt;p&gt;Take Llama 3.1 8B's shape. The scores for one sequence in one layer, in bf16, are 268 MB at 2,048 tokens, 4.3 GB at 8,192, and 1.1 TB at the 131,072 the config declares; the eager path holds three such squares at its peak, two of them in fp32: about five times those bytes. The weights of the whole model are 16 GB. Nobody materialises that last number. The kernel is why.&lt;/p&gt;

&lt;p&gt;The second widget on &lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;the original post&lt;/a&gt; is that multiplication, live: drag &lt;code&gt;T&lt;/code&gt; to 16k with &lt;code&gt;eager&lt;/code&gt; selected and one layer's scores outgrow the whole model's weights; at 131k they are about seventy times larger. Switch to &lt;code&gt;sdpa&lt;/code&gt; and the number is struck out — never allocated.&lt;/p&gt;

&lt;p&gt;I timed the last three lines on my laptop's CPU with 8 heads of width 64. From 512 tokens to 8,192, sixteen times the input, the eager form went from 1.9 ms to 284 ms: 150 times. A linear layer over the same input went from 0.23 ms to 2.5 ms, eleven times. Doubling &lt;code&gt;T&lt;/code&gt; from 4,096 to 8,192 multiplied the linear layer by 1.95 and attention by 4.3: the square, within timing noise.&lt;/p&gt;

&lt;p&gt;Three levers move the number. The first is the kernel. &lt;code&gt;F.scaled_dot_product_attention&lt;/code&gt; with its flash or memory-efficient backend computes the same numbers in blocks and never holds the &lt;code&gt;T × T&lt;/code&gt; matrix; its math fallback still does. At 4,096 tokens, in fp32, the profiler charged 545 MB to the eager form's matmuls, the 537 MB square plus its output, and 8.5 MB to sdpa, and it ran 2.4 times faster.&lt;/p&gt;

&lt;p&gt;Flash attention is that idea on a GPU. The arithmetic is still quadratic. What changes is what must exist in memory at once. The line that picks it is one config field, &lt;code&gt;_attn_implementation&lt;/code&gt;. In 5.17.0 the default at load is &lt;code&gt;sdpa&lt;/code&gt;; &lt;code&gt;eager&lt;/code&gt; is what you get by asking for it. With padding in the batch, sdpa still holds the mask itself, &lt;code&gt;[B, 1, T, T]&lt;/code&gt; booleans, 17 MB at 4,096; without padding, transformers passes no mask and lets the kernel apply the triangle (&lt;code&gt;is_causal=True&lt;/code&gt;).&lt;/p&gt;

&lt;p&gt;The other two change what's computed. Grouped-query attention shrinks the cache. With 8 key-value heads the 8B model stores 131 kB per token across its 32 layers, 17 GB for a 128k context, already a little more than the weights, and with 32 heads it would be 524 kB per token and 69 GB, four times the weights. A sliding window caps how far back a query reads. Llama 3.1 has none; Mistral 7B v0.1 had one, and Gemma 2 and 3 interleave windowed layers with global ones. transformers writes it as &lt;code&gt;kv_idx &amp;gt; q_idx - sliding_window&lt;/code&gt;, at &lt;code&gt;masking_utils.py:100&lt;/code&gt;. At 32,768 tokens a 4,096 window leaves 126 million live cells where the causal triangle has 537 million. Below the window the cost is a triangle. Above it, a band of width &lt;code&gt;W&lt;/code&gt;. A band grows with &lt;code&gt;T&lt;/code&gt;, not its square.&lt;/p&gt;

&lt;p&gt;Be fair about the square. Count operations rather than bytes, over the prompt pass, and the linear layers of the 8B model do more of them than attention until roughly 30,000 tokens, about twice that if you count only the triangle a causal kernel computes. At 8,192, attention is about a fifth of the block's floating-point operations, its FLOPs, counting the full square, and about an eighth counting the triangle.&lt;/p&gt;

&lt;p&gt;The square is a memory problem long before it's a compute problem. The number to compute before promising a context length is the memory one.&lt;/p&gt;

&lt;h2&gt;
  
  
  What the ML engineer watches
&lt;/h2&gt;

&lt;p&gt;Four things, and the first is a multiplication.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Price &lt;code&gt;[B, H, T, T]&lt;/code&gt; and the cache before promising a context.&lt;/strong&gt; Heads times &lt;code&gt;T&lt;/code&gt; squared times two bytes, per layer. Add the cache per token times the longest sequence times the batch. Do it at the batch you'll serve. The config's &lt;code&gt;max_position_embeddings&lt;/code&gt; is what the model accepts after its rotary scaling (131,072 here: an 8,192 pre-training window, scaled and then trained out to 128k), not what your card can hold.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Know which kernel runs.&lt;/strong&gt; &lt;code&gt;eager&lt;/code&gt; returns the weights and materialises the square; &lt;code&gt;sdpa&lt;/code&gt; and flash don't. The choice is &lt;code&gt;_attn_implementation&lt;/code&gt;, and it decides whether the 545 MB above exists. A debugging tool that wants the weights can't get them from &lt;code&gt;sdpa&lt;/code&gt;: &lt;code&gt;output_attentions=True&lt;/code&gt; returns an empty tuple there, and the fix, loading with &lt;code&gt;attn_implementation="eager"&lt;/code&gt;, materialises the square on every call.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Compose the masks on purpose.&lt;/strong&gt; Causal, padding, and window are ANDed; added, in the &lt;code&gt;-inf&lt;/code&gt; form. A query row with nothing left to read is NaN in that form and uniform in the library's, and neither is a number you want. Part 3's habit applies twice now: the last-token index comes from the mask, and the mask travels with the batch.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Measure latency at the longest prompt you'll accept.&lt;/strong&gt; The linear layers scale with &lt;code&gt;T&lt;/code&gt;. Attention scales with its square. The cache read at each generated token scales with everything already in it. A p99 measured at 500 tokens says nothing about 8,000. The 1.9 ms to 284 ms above is the shape of the surprise.&lt;/p&gt;

&lt;p&gt;None of this is modelling work either. A weighted average is four lines, and the weights come from three learned matrices. The square they live in is the first cost in this series that the input sets and the code doesn't.&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Part 4 of twelve. Attention is a weighted average whose weights the model computes from its own input: the average is four lines, the &lt;code&gt;T × T&lt;/code&gt; matrix of weights is the cost, and the kernel decides whether that square is ever held in memory. Attention is one stage. The next part reads the whole block around it — six tensors, and a residual stream that every stage adds to — and lands on &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;the series index&lt;/a&gt; when it's out. If you arrived here directly, &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Part 1&lt;/a&gt; reads the tokenizer, &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;Part 2&lt;/a&gt; the table a token id indexes, and &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;Part 3&lt;/a&gt; the shape both of them hand to this one.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Written with AI assistance. See &lt;a href="https://cchinchilla.dev/ai-policy/" rel="noopener noreferrer"&gt;AI policy&lt;/a&gt;.&lt;/em&gt;&lt;/p&gt;

</description>
      <category>machinelearning</category>
      <category>python</category>
      <category>ai</category>
      <category>beginners</category>
    </item>
    <item>
      <title>Shapes are the type system of ML, and nobody checks them for you</title>
      <dc:creator>Carlos Chinchilla Corbacho</dc:creator>
      <pubDate>Thu, 24 Sep 2026 13:00:00 +0000</pubDate>
      <link>https://dev.to/cchinchilladev/shapes-are-the-type-system-of-ml-and-nobody-checks-them-for-you-2f8k</link>
      <guid>https://dev.to/cchinchilladev/shapes-are-the-type-system-of-ml-and-nobody-checks-them-for-you-2f8k</guid>
      <description>&lt;p&gt;&lt;em&gt;Originally published at &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;cchinchilla.dev&lt;/a&gt;. Part 3 of From code to weights, a 12-part series on ML fundamentals for engineers.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;Part 2&lt;/a&gt; ended on a signature. A batch of sequences of embedding rows is a tensor of shape &lt;code&gt;[B, T, d]&lt;/code&gt;. That line is the closest thing this stack has to a type. This post reads it the way you'd read a function's signature, then shows what the runtime does when the signature is wrong: usually nothing.&lt;/p&gt;

&lt;p&gt;I read the broadcasting rule as torch writes it. Then I wrote four shape bugs on purpose and watched which ones the runtime caught. It caught one.&lt;/p&gt;

&lt;p&gt;This is part of &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;a series&lt;/a&gt; for engineers moving into ML engineering. Each part names one thing that does not behave the way software behaves. This one is the type system: there isn't one.&lt;/p&gt;

&lt;h2&gt;
  
  
  A shape is a signature
&lt;/h2&gt;

&lt;p&gt;A tensor is an N-dimensional array plus a description of how it sits in memory. &lt;code&gt;[B, T, d]&lt;/code&gt; says: &lt;code&gt;B&lt;/code&gt; sequences in the batch, &lt;code&gt;T&lt;/code&gt; positions in each, &lt;code&gt;d&lt;/code&gt; numbers at each position. Read it the way you'd read &lt;code&gt;f(batch: List[List[Vec]])&lt;/code&gt;, with one difference that matters later: the runtime never looks at the names, only at the sizes.&lt;/p&gt;

&lt;p&gt;The memory description is the part software engineers already know. A tensor of shape &lt;code&gt;[4, 8, 16]&lt;/code&gt; has strides &lt;code&gt;(128, 16, 1)&lt;/code&gt;: move one step along the first axis and you skip 128 numbers. Transpose the first two axes and the shape becomes &lt;code&gt;[8, 4, 16]&lt;/code&gt; with strides &lt;code&gt;(16, 128, 1)&lt;/code&gt;, the same memory read in a different order; &lt;code&gt;data_ptr&lt;/code&gt; doesn't change. That's why &lt;code&gt;transpose&lt;/code&gt; is free and why a &lt;code&gt;view&lt;/code&gt; that flattens the result refuses, with a message that says so, while &lt;code&gt;reshape&lt;/code&gt; copies when it has to. A tensor is a pointer, a shape, and a stride tuple; the view operations you'll call are arithmetic on the last two.&lt;/p&gt;

&lt;p&gt;Every operation has a shape signature of its own: what goes in, what comes out, which axis disappears. &lt;code&gt;mean(dim=1)&lt;/code&gt; on &lt;code&gt;[B, T, d]&lt;/code&gt; returns &lt;code&gt;[B, d]&lt;/code&gt;; the &lt;code&gt;T&lt;/code&gt; is gone. A linear layer maps &lt;code&gt;[B, T, d]&lt;/code&gt; to &lt;code&gt;[B, T, d_out]&lt;/code&gt;; only the last axis changes. Writing the signature next to the call is the discipline this post is about; nothing else will write it.&lt;/p&gt;

&lt;h2&gt;
  
  
  The rule is five lines
&lt;/h2&gt;

&lt;p&gt;Two tensors of different shapes can still be added, and the rule that decides how is broadcasting. torch keeps a readable copy of it in Python, under &lt;code&gt;torch.broadcast_shapes&lt;/code&gt;, which calls &lt;code&gt;_broadcast_shapes&lt;/code&gt; in &lt;code&gt;torch/_refs/__init__.py&lt;/code&gt; (lines 388-470 at v2.13.0). Stripped of the symbolic-shape machinery, the rule is this:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;common_shape&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="nf"&gt;reduce&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;max&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nf"&gt;len&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;shape&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;shape&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="n"&gt;shapes&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;
&lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;arg_idx&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;shape&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="nf"&gt;enumerate&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;shapes&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;idx&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt; &lt;span class="nf"&gt;len&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;shape&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;          &lt;span class="c1"&gt;# align from the right
&lt;/span&gt;        &lt;span class="k"&gt;if&lt;/span&gt; &lt;span class="n"&gt;common_shape&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;==&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
            &lt;span class="n"&gt;common_shape&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;shape&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;              &lt;span class="c1"&gt;# 1 stretches to anything
&lt;/span&gt;        &lt;span class="k"&gt;if&lt;/span&gt; &lt;span class="n"&gt;shape&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;==&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
            &lt;span class="k"&gt;continue&lt;/span&gt;                                    &lt;span class="c1"&gt;# ... in either direction
&lt;/span&gt;        &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;_check&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;common_shape&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;==&lt;/span&gt; &lt;span class="n"&gt;shape&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt;   &lt;span class="c1"&gt;# otherwise sizes must match
&lt;/span&gt;            &lt;span class="k"&gt;lambda&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="sa"&gt;f&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;Attempting to broadcast a dimension of length &lt;/span&gt;&lt;span class="si"&gt;{&lt;/span&gt;&lt;span class="n"&gt;shape&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;&lt;span class="si"&gt;}&lt;/span&gt;&lt;span class="s"&gt; at &lt;/span&gt;&lt;span class="si"&gt;{&lt;/span&gt;&lt;span class="n"&gt;idx&lt;/span&gt;&lt;span class="si"&gt;}&lt;/span&gt;&lt;span class="s"&gt;!&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;common_shape&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Align the shapes from the right. A size of 1 stretches to match the other. Equal sizes pass. Anything else raises. That's the whole rule. The eager path runs the same rule in C++ and raises with a shorter message. It exists so that you can add a bias of shape &lt;code&gt;[d]&lt;/code&gt; to &lt;code&gt;[B, T, d]&lt;/code&gt; without a loop or a copy.&lt;/p&gt;

&lt;p&gt;The rule has one property that the rest of this post is about: it decides by size, not by meaning. &lt;code&gt;[8, 16] + [16]&lt;/code&gt; works because 16 equals 16. &lt;code&gt;[8, 16] + [8]&lt;/code&gt; raises, because 8 is not 16 and 8 is not 1. Whether the 16 on the right is a sequence length or a hidden size or a batch, the rule has no idea, and neither will you when it's wrong. When two of your axes happen to be the same size, it can't tell them apart either. The data scientist has leaned on this rule in numpy for years; the software engineer is meeting a coercion rule with no declared types. Both should read it as the second thing, not the first.&lt;/p&gt;

&lt;h2&gt;
  
  
  The bug that only happens at batch size sixteen
&lt;/h2&gt;

&lt;p&gt;The first bug. Attention scores have shape &lt;code&gt;[B, T, T]&lt;/code&gt;, one &lt;code&gt;T × T&lt;/code&gt; matrix per sequence. A padding mask has shape &lt;code&gt;[B, T]&lt;/code&gt;: which positions in each sequence are real. In its additive form (0 for real, &lt;code&gt;-inf&lt;/code&gt; for padding), adding it to the scores switches padding off. The correct form is &lt;code&gt;scores + mask[:, None, :]&lt;/code&gt;: the &lt;code&gt;None&lt;/code&gt; makes the mask &lt;code&gt;[B, 1, T]&lt;/code&gt;, so it stretches across the query axis and masks the keys.&lt;/p&gt;

&lt;p&gt;Drop the &lt;code&gt;None&lt;/code&gt;. With &lt;code&gt;B = 8&lt;/code&gt; and &lt;code&gt;T = 16&lt;/code&gt;, &lt;code&gt;scores + mask&lt;/code&gt; raises: 16 is not 8. With &lt;code&gt;B = 16&lt;/code&gt; and &lt;code&gt;T = 16&lt;/code&gt;, it runs. At &lt;code&gt;B = 1&lt;/code&gt; it runs and is correct by accident: 1 is what the &lt;code&gt;None&lt;/code&gt; would have inserted.&lt;/p&gt;

&lt;p&gt;The rule aligns &lt;code&gt;[16, 16]&lt;/code&gt; against &lt;code&gt;[16, 16, 16]&lt;/code&gt; from the right and reads the mask as &lt;code&gt;[1, B, T]&lt;/code&gt;. The mask meant for sequence &lt;code&gt;b&lt;/code&gt; lands on query row &lt;code&gt;b&lt;/code&gt; of every sequence. On random lengths I measured a third of the positions masked wrong and a fifth of the attention weight on padding. The softmax after it is finite, the loss is finite, and the run continues. No warning, because there is nothing to warn about: the rule was followed.&lt;/p&gt;

&lt;p&gt;The original post has a widget for this that doesn't survive being pasted into a feed: &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;move the two sliders&lt;/a&gt; until &lt;code&gt;B&lt;/code&gt; and &lt;code&gt;T&lt;/code&gt; meet and the verdict flips from an error to a silent result. Then pick &lt;code&gt;[B, T, 1]&lt;/code&gt;: that mask runs at every size and is wrong at every size, because it switches off queries instead of keys. It's the loud version — a padded query's row is all &lt;code&gt;-inf&lt;/code&gt; and its softmax is NaN. The rule is the same in all three cases.&lt;/p&gt;

&lt;p&gt;The second bug is the classic one. Predictions of shape &lt;code&gt;[256]&lt;/code&gt;, targets of shape &lt;code&gt;[256, 1]&lt;/code&gt;, subtracted for a squared error. The rule stretches both to &lt;code&gt;[256, 256]&lt;/code&gt;, every prediction against every target. The mean of that is a number. It trains.&lt;/p&gt;

&lt;p&gt;I trained a one-feature linear regression through it for two thousand steps. The loss fell and settled at 9.7. The fitted slope was 0.000, the intercept was the mean of the targets, and the R² was 0.0. The model learned to ignore its input. That's the best it can do when the loss compares it with everyone else's target. The library loss, &lt;code&gt;F.mse_loss&lt;/code&gt;, warns about this mismatch. A hand-written &lt;code&gt;((pred - target) ** 2).mean()&lt;/code&gt; doesn't, and the hand-written form is the one that ends up in notebooks.&lt;/p&gt;

&lt;p&gt;The third bug never raises where it happens. On a tensor of shape &lt;code&gt;[16, 16, 64]&lt;/code&gt;, &lt;code&gt;mean(dim=0)&lt;/code&gt; and &lt;code&gt;mean(dim=1)&lt;/code&gt; both return &lt;code&gt;[16, 64]&lt;/code&gt;. Different numbers. Same shape. Nothing downstream can tell. On &lt;code&gt;[3, 5, 7]&lt;/code&gt; the same two calls return &lt;code&gt;[5, 7]&lt;/code&gt; and &lt;code&gt;[3, 7]&lt;/code&gt;, and the first thing that tries to combine them raises. The bug was there in both cases; only the sizes decided whether it was visible.&lt;/p&gt;

&lt;p&gt;The fourth is the control. &lt;code&gt;[3, 5, 7]&lt;/code&gt; plus &lt;code&gt;[3, 5]&lt;/code&gt; fails because 7 is not 5. It was caught for the one reason the others weren't: the sizes differ.&lt;/p&gt;

&lt;p&gt;Be fair about the rule. It's doing what it was asked to. The alternative exists: a runtime that carries axis names and refuses to add &lt;code&gt;T&lt;/code&gt; to &lt;code&gt;B&lt;/code&gt;. torch shipped one as a prototype, named tensors. It never left prototype, and 2.13 removed it to cut overhead and code bloat. Part of why it never caught on: every library in the stack would have had to agree on the names, and sizes are the only thing all of them share. The rule is right. It is not a type check, and nothing in the stack is.&lt;/p&gt;

&lt;h2&gt;
  
  
  Padding is a shape decision
&lt;/h2&gt;

&lt;p&gt;Sequences don't come in rectangles. Three sequences of lengths 5, 3, and 1 become a &lt;code&gt;[3, 5]&lt;/code&gt; tensor only by padding, and for those lengths 40% of the rectangle is padding. Every operation over &lt;code&gt;T&lt;/code&gt; then needs to know which cells are real. That knowledge is a second tensor, the mask, and it has to travel with the first.&lt;/p&gt;

&lt;p&gt;Which side you pad on is part of the signature. Pad on the right and the last real token sits at index &lt;code&gt;length - 1&lt;/code&gt;: 4, 2, 0 for the three above. Pad on the left and it sits at 4 for all of them. Decoders that generate one token at a time prefer that; the position to read is always the last. Either choice works. Mixing them is the bug. Code written for right padding and fed a left-padded batch reads padding, or the wrong token, where it expects the last one.&lt;/p&gt;

&lt;p&gt;The mask is not optional even when the model doesn't attend. Mean-pool that &lt;code&gt;[3, 5]&lt;/code&gt; batch without it and the length-1 sequence's vector averages one real row with four of padding. On &lt;code&gt;arange(1, 31).reshape(3, 5, 2)&lt;/code&gt; the unmasked mean of that sequence is &lt;code&gt;[25, 26]&lt;/code&gt;; the masked mean is &lt;code&gt;[21, 22]&lt;/code&gt;, its only real row. Part 2's pooling code divides by the mask sum for this reason. The division is the signature, written out.&lt;/p&gt;

&lt;h2&gt;
  
  
  Memory is shape times bytes
&lt;/h2&gt;

&lt;p&gt;A tensor's cost is arithmetic. Activations of shape &lt;code&gt;[32, 2048, 4096]&lt;/code&gt; are 268 million numbers: 1.07 GB in fp32, 537 MB in bf16, 268 MB in int8. That's the bill, and you need nothing about the model to compute it: the shape and the width of one number. It decides whether a batch fits on a card.&lt;/p&gt;

&lt;p&gt;The second widget on &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;the original post&lt;/a&gt; is that multiplication, live: drag the three sizes and switch the format. At the defaults one block costs 77 GB in fp32 and the two square rows are 89% of it; pull &lt;code&gt;T&lt;/code&gt; down to 128 and they fall to a third. Those two rows are the next part's subject.&lt;/p&gt;

&lt;p&gt;What each format gives up is also arithmetic. fp32 keeps 23 bits of mantissa. bf16 keeps 7 and the whole fp32 exponent: it writes π as 3.140625 and rounds &lt;code&gt;1 + 0.004&lt;/code&gt; to &lt;code&gt;1.0078125&lt;/code&gt;, but its range is fp32's, to within the last rounding step. fp16 keeps 10 bits and only five of exponent: it tops out at 65,504, and small gradients underflow to zero unless the loss is scaled first. That's why training moved to bf16 and serving reaches for int8. Each is a decision about which precision a tensor doesn't need.&lt;/p&gt;

&lt;p&gt;The batch axis is the one you control at serving time, and it moves two numbers in opposite directions. I timed a &lt;code&gt;[B, 4096]&lt;/code&gt; by &lt;code&gt;[4096, 4096]&lt;/code&gt; matmul on my laptop's CPU. It costs 447 µs per row at &lt;code&gt;B = 1&lt;/code&gt; and 22 µs per row at &lt;code&gt;B = 512&lt;/code&gt;: twenty times the throughput. The call itself went from 0.45 ms to 11.4 ms; each request waited twenty-five times longer, plus whatever it spent waiting for the batch to fill. Five repetitions on one laptop: the ratio is the finding, not the digits. Dynamic batching is the knob between those two numbers, and its setting is a product decision wearing a shape.&lt;/p&gt;

&lt;h2&gt;
  
  
  What the ML engineer watches
&lt;/h2&gt;

&lt;p&gt;Four habits, and the first two are one habit.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Annotate shapes at function boundaries.&lt;/strong&gt; &lt;code&gt;jaxtyping&lt;/code&gt; with &lt;code&gt;beartype&lt;/code&gt; turns the signature into a check that runs on every call:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="nd"&gt;@jaxtyped&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;typechecker&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="n"&gt;beartype&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;pool&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;Float&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;Tensor&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;B T d&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="n"&gt;mask&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;Bool&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;Tensor&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;B T&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;])&lt;/span&gt; &lt;span class="o"&gt;-&amp;gt;&lt;/span&gt; &lt;span class="n"&gt;Float&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;Tensor&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;B d&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]:&lt;/span&gt;
    &lt;span class="n"&gt;m&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;mask&lt;/span&gt;&lt;span class="p"&gt;[:,&lt;/span&gt; &lt;span class="p"&gt;:,&lt;/span&gt; &lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;].&lt;/span&gt;&lt;span class="nf"&gt;float&lt;/span&gt;&lt;span class="p"&gt;()&lt;/span&gt;
    &lt;span class="nf"&gt;return &lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;m&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;sum&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;/&lt;/span&gt; &lt;span class="n"&gt;m&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;sum&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Reduce the wrong axis inside and it raises on the return, because &lt;code&gt;[T, d]&lt;/code&gt; is not &lt;code&gt;[B, d]&lt;/code&gt;. Pass the column-shaped target into the squared error and it raises before the rule can stretch anything. Two of the bugs above, caught at the boundary (the reduce only when &lt;code&gt;B&lt;/code&gt; ≠ &lt;code&gt;T&lt;/code&gt;), with the shape written where the next reader looks. The batch-sixteen bug, and the wrong-axis reduce when &lt;code&gt;B&lt;/code&gt; equals &lt;code&gt;T&lt;/code&gt;, need the next habit.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Test with sizes that can't collide.&lt;/strong&gt; The bug at batch size sixteen is invisible to the annotation. When &lt;code&gt;B&lt;/code&gt; and &lt;code&gt;T&lt;/code&gt; are both 16, the names bind to the same number and the check passes; so does the wrong-axis reduce, because &lt;code&gt;[T, d]&lt;/code&gt; and &lt;code&gt;[B, d]&lt;/code&gt; are the same &lt;code&gt;[16, 64]&lt;/code&gt;. The fix is in the tests, not the code: use &lt;code&gt;B = 3&lt;/code&gt;, &lt;code&gt;T = 5&lt;/code&gt;, &lt;code&gt;d = 7&lt;/code&gt;. Every dropped &lt;code&gt;None&lt;/code&gt; and every mix-up between &lt;code&gt;B&lt;/code&gt;, &lt;code&gt;T&lt;/code&gt; and &lt;code&gt;d&lt;/code&gt; then produces a size that matches nothing, and the runtime raises. Size-1 axes still stretch at any size, and the two &lt;code&gt;T&lt;/code&gt;s of the scores are always equal; those need the annotation or a test that checks values. Round numbers in tests are how the bug at batch size sixteen survives to production. That's why the first two are one habit: the annotation is only as strict as the sizes it sees.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Mask and padding side are one decision.&lt;/strong&gt; Record the side next to the tokenizer config. Compute the last-token index from the mask, never from &lt;code&gt;T&lt;/code&gt;. Reject a batch whose mask disagrees with its padding. The pooling code in Part 2 already divides by the mask sum; make sure yours does.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Price a tensor before allocating it.&lt;/strong&gt; Shape times bytes, at the batch you'll serve, in the format you'll serve it in. That multiplication has an answer before a single weight is loaded.&lt;/p&gt;

&lt;p&gt;None of this is modelling work either. A shape is the one type this stack never checks: a device mismatch raises, a wrong axis broadcasts. The check is yours to write.&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Part 3 of twelve. A shape is a signature that nobody verifies: broadcasting decides by size, so a wrong axis raises only when the sizes happen to differ, and a bug can live at one batch size alone. The attention scores in the first widget and the two amber rows in the second carry &lt;code&gt;T&lt;/code&gt; twice, one &lt;code&gt;T × T&lt;/code&gt; matrix per sequence (per head, in the second). &lt;a href="https://cchinchilla.dev/posts/fcw-attention-learned-weighted-average/" rel="noopener noreferrer"&gt;The next part&lt;/a&gt; reads where that square comes from and what it costs. If you arrived here directly, &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Part 1&lt;/a&gt; reads the tokenizer and &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;Part 2&lt;/a&gt; the table a token id indexes.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Written with AI assistance. See &lt;a href="https://cchinchilla.dev/ai-policy/" rel="noopener noreferrer"&gt;AI policy&lt;/a&gt;.&lt;/em&gt;&lt;/p&gt;

</description>
      <category>machinelearning</category>
      <category>python</category>
      <category>ai</category>
      <category>beginners</category>
    </item>
    <item>
      <title>An embedding is a lookup table, and everything else is how you fill it</title>
      <dc:creator>Carlos Chinchilla Corbacho</dc:creator>
      <pubDate>Tue, 22 Sep 2026 13:00:00 +0000</pubDate>
      <link>https://dev.to/cchinchilladev/an-embedding-is-a-lookup-table-and-everything-else-is-how-you-fill-it-50p8</link>
      <guid>https://dev.to/cchinchilladev/an-embedding-is-a-lookup-table-and-everything-else-is-how-you-fill-it-50p8</guid>
      <description>&lt;p&gt;&lt;em&gt;Originally published at &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;cchinchilla.dev&lt;/a&gt;. Part 2 of From code to weights, a 12-part series on ML fundamentals for engineers.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Part 1&lt;/a&gt; ended on an index. The tokenizer turns text into ids, an id is a row number, and the row lives in a matrix the model was born with. This post reads that matrix. It's smaller than it sounds: a lookup table, indexed by integer, with learned contents. The science is in how the contents get there. The engineering is in what you can't do with the table once they're in.&lt;/p&gt;

&lt;p&gt;I read two files for this post: the &lt;code&gt;forward&lt;/code&gt; of &lt;code&gt;torch.nn.Embedding&lt;/code&gt;, which is one call, and the pooling module of sentence-transformers, which is where the other kind of embedding, the one you put in an index, gets made. Then I built a 48-line index and broke it three ways on purpose, to measure what breaks.&lt;/p&gt;

&lt;p&gt;This is part of &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;a series&lt;/a&gt; for engineers moving into ML engineering. Each part names one thing that does not behave the way software behaves. This one is the index: it has no migration.&lt;/p&gt;

&lt;h2&gt;
  
  
  A matrix and an index
&lt;/h2&gt;

&lt;p&gt;If you've written a hash map, you've written most of an embedding layer. &lt;code&gt;nn.Embedding(V, d)&lt;/code&gt; allocates a matrix of &lt;code&gt;V&lt;/code&gt; rows and &lt;code&gt;d&lt;/code&gt; columns. Looking up id &lt;code&gt;i&lt;/code&gt; returns row &lt;code&gt;i&lt;/code&gt;. That's the whole forward pass, and torch's own source says so in &lt;code&gt;sparse.py:188-197&lt;/code&gt;:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;forward&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="nb"&gt;input&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;Tensor&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;-&amp;gt;&lt;/span&gt; &lt;span class="n"&gt;Tensor&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;F&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;embedding&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;
        &lt;span class="nb"&gt;input&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
        &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;weight&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
        &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;padding_idx&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
        &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;max_norm&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
        &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;norm_type&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
        &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;scale_grad_by_freq&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
        &lt;span class="n"&gt;self&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;sparse&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;
    &lt;span class="p"&gt;)&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;&lt;code&gt;F.embedding&lt;/code&gt; is a gather. That's checkable rather than a matter of trust: for a fresh &lt;code&gt;Embedding(50257, 768)&lt;/code&gt; and three ids, &lt;code&gt;torch.equal(emb(ids), emb.weight[ids])&lt;/code&gt; is &lt;code&gt;True&lt;/code&gt;. No transformation, no hashing, no similarity. A row per id.&lt;/p&gt;

&lt;p&gt;The sizes are where it stops looking like a hash map. GPT-2's input table, &lt;code&gt;wte.weight&lt;/code&gt;, is 50,257 rows by 768 columns: 38.6 million numbers, 154 MB in fp32, every one of them a trained parameter. &lt;code&gt;all-MiniLM-L6-v2&lt;/code&gt;, the sentence model the rest of this post runs on, has a 30,522 by 384 table for the same job. Row 11,241 of GPT-2's table is the vector for &lt;code&gt;" token"&lt;/code&gt; for the reason Part 1 gave: that's where the tokenizer sends it. The contents of the row were learned. The address was fixed before training started.&lt;/p&gt;

&lt;p&gt;For the engineer coming from software, the surprise is that the keys mean nothing. A hash map's key carries its own identity; row 11,241 means something only inside the model that trained it. For the data scientist, the surprise runs the other way. This table isn't a word2vec artifact you produce and then use. It's a parameter inside the model, updated by the same gradient step as everything above it. It never leaves.&lt;/p&gt;

&lt;h2&gt;
  
  
  Two ways to fill the table
&lt;/h2&gt;

&lt;p&gt;A row gets its numbers in one of two ways, and the difference decides what you can do with the result.&lt;/p&gt;

&lt;p&gt;The first is the one above. The table is a parameter of a language model. Rows start random and gradients flow into them through the layers above. After training, each row is whatever helped the model predict the next token. Nobody designs the space; it's a side effect. Part 8 reads the gradient that fills it.&lt;/p&gt;

&lt;p&gt;The second is the one most engineers actually deploy. You take a trained model, run a piece of text through it, and take the output as the text's vector. The row isn't stored anywhere; it's computed on demand, and it depends on the whole input, not on one id.&lt;/p&gt;

&lt;p&gt;The word &lt;code&gt;bank&lt;/code&gt; has one row in the input table of that sentence model, row 2,924. A position table gets added to it, but the word's row is the same in "she sat on the river bank" and in "she opened a bank account". At the output, the two vectors for that same token have a cosine of 0.77 to each other. Same lookup, different result, because a position and six layers sit between the two. That's what contextual means, and it's the whole difference between a static embedding and the kind you put in a search index.&lt;/p&gt;

&lt;p&gt;The output is a vector per token. A sentence needs one vector, so something has to collapse &lt;code&gt;T&lt;/code&gt; vectors, one per token, into one, and that something is pooling. sentence-transformers keeps every variant in one place, and on the default padded path the branch a mean-pooled model runs is &lt;code&gt;pooling.py:200-213&lt;/code&gt;:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;elif&lt;/span&gt; &lt;span class="n"&gt;mode&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;mean&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;mean_sqrt_len_tokens&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;if&lt;/span&gt; &lt;span class="n"&gt;mean_sum&lt;/span&gt; &lt;span class="ow"&gt;is&lt;/span&gt; &lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
        &lt;span class="n"&gt;mask&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;attention_mask&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;unsqueeze&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;expand_as&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;token_embeddings&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;to&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;token_embeddings&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;dtype&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
        &lt;span class="n"&gt;mean_sum&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;token_embeddings&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt; &lt;span class="n"&gt;mask&lt;/span&gt;&lt;span class="p"&gt;).&lt;/span&gt;&lt;span class="nf"&gt;sum&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;dim&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
        &lt;span class="c1"&gt;# ... mean_mask is the count of real tokens (or a supplied per-token weight sum), clamped away from zero
&lt;/span&gt;    &lt;span class="k"&gt;if&lt;/span&gt; &lt;span class="n"&gt;mode&lt;/span&gt; &lt;span class="o"&gt;==&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;mean&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
        &lt;span class="n"&gt;output_vectors&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;append&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;mean_sum&lt;/span&gt; &lt;span class="o"&gt;/&lt;/span&gt; &lt;span class="n"&gt;mean_mask&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Sum the tokens the attention mask says are real, divide by how many there were. Six modes live in that function, each a different way to collapse the same &lt;code&gt;T&lt;/code&gt; vectors: &lt;code&gt;cls&lt;/code&gt; takes the first, &lt;code&gt;lasttoken&lt;/code&gt; the last, &lt;code&gt;max&lt;/code&gt; the largest per dimension. A model is trained with one of them, and the choice is part of the artifact.&lt;/p&gt;

&lt;p&gt;So is the step after it. The small model's own &lt;code&gt;modules.json&lt;/code&gt; lists what runs when you call &lt;code&gt;encode&lt;/code&gt;:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight json"&gt;&lt;code&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="w"&gt;
  &lt;/span&gt;&lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="nl"&gt;"idx"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"name"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"0"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"path"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;""&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"type"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"sentence_transformers.models.Transformer"&lt;/span&gt;&lt;span class="p"&gt;},&lt;/span&gt;&lt;span class="w"&gt;
  &lt;/span&gt;&lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="nl"&gt;"idx"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"name"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"1"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"path"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"1_Pooling"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"type"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"sentence_transformers.models.Pooling"&lt;/span&gt;&lt;span class="p"&gt;},&lt;/span&gt;&lt;span class="w"&gt;
  &lt;/span&gt;&lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="nl"&gt;"idx"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"name"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"2"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"path"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"2_Normalize"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="nl"&gt;"type"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;&lt;span class="w"&gt; &lt;/span&gt;&lt;span class="s2"&gt;"sentence_transformers.models.Normalize"&lt;/span&gt;&lt;span class="p"&gt;}&lt;/span&gt;&lt;span class="w"&gt;
&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;&lt;span class="w"&gt;
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Transformer, then pooling, then normalisation to unit length. Not even the family agrees on the third: &lt;code&gt;paraphrase-MiniLM-L6-v2&lt;/code&gt; ships without it, and its raw vectors have norms around six. Three modules, and the last two are the ones a client tends to reimplement. Reimplemented differently, they change the answer. By how much is measured two sections down.&lt;/p&gt;

&lt;h2&gt;
  
  
  Close is a dot product
&lt;/h2&gt;

&lt;p&gt;Once every vector is unit length, cosine similarity is a dot product, and nearest-neighbour search is a matrix multiply followed by a sort. GPT-2 reads its own table the same way on the way out: the logits for the next token are a dot product against every row of &lt;code&gt;wte&lt;/code&gt;. That's the whole retrieval side, and it's why the index bill is arithmetic: ten million documents at 384 dimensions is 15.4 GB of fp32 before any index structure. At 768 it's 30.7 GB. At 1,536, the width of a popular hosted embedding, 61.4 GB. Dimension is a cost you pay per document for as long as the index exists. Whoever trained the model decided it, not you.&lt;/p&gt;

&lt;p&gt;The number the search returns is less useful than it looks. On my 48-line corpus, the best match for "how do I undo a bad release" has a cosine of 0.30 to the query, and it's the right line: the rollback procedure. The best match for "the migration is stuck waiting on a lock" scores 0.71. Both are correct. A cosine isn't a confidence; it's a position in a space whose scale the model chose. Compare ranks within one model, and never compare the raw number across two.&lt;/p&gt;

&lt;h2&gt;
  
  
  There is no schema migration
&lt;/h2&gt;

&lt;p&gt;Here is the part that's different from software. When you change a database schema you write a migration, run it, and the old rows come out the other side in the new shape. An embedding space has no migration. Two models produce two spaces, and no map between them ships with either, because neither was designed; each is whatever training settled on. You can fit one after the fact, and people do. It's a learned approximation with its own error, not a migration: a linear map fitted on four thousand unrelated sentences lifts the stale index below from 60% to 70%, and stops there. Re-embed with a new model and every vector you stored is in a coordinate system nothing else uses.&lt;/p&gt;

&lt;p&gt;I expected that to show up as garbage. It shows up as something worse.&lt;/p&gt;

&lt;p&gt;The test: 48 lines from a runbook, eight each on deploys, auth, billing, databases, observability, and on-call. I indexed them with the L6 model and queried through &lt;code&gt;all-MiniLM-L12-v2&lt;/code&gt;, the upgrade an engineer would reach for (same family, same 384 dimensions, so the index accepts the vectors without complaint).&lt;/p&gt;

&lt;p&gt;Against what the upgraded index would return, the stale one keeps 60% of the top five. Random would keep 10%, and guessing five lines from the right topic would keep 35%. The old index queried by the old model keeps 77%: the mismatch costs seventeen points, and the other twenty-three are the two models disagreeing about the answer. For "user keeps getting logged out", the first hit is still the session-expiry line, and the third is the database backup retention policy. Nothing in the response says which is which.&lt;/p&gt;

&lt;p&gt;Be fair about the number. Sixty percent isn't a scrambled index, and the reason is lineage, not English: both models were fine-tuned from the same parent checkpoint, the L6 one after dropping every second layer, so the two spaces started in the same coordinates. The same line under the two models has, on average, a cosine of 0.54 to itself, far from orthogonal. Two models that only shared training data wouldn't share coordinates, and a stale index between them should land near random. Six tight topics of eight lines each is also the friendliest corpus a stale index could hope for. I haven't measured a production index, and I'd expect it to keep less. What the number does support is the shape: partial, plausible, silent.&lt;/p&gt;

&lt;p&gt;The pooling module does the same thing without a model change. Take the same weights and pool with &lt;code&gt;cls&lt;/code&gt; instead of &lt;code&gt;mean&lt;/code&gt;, at index time and at query time, and the top five keeps 50% of the mean-pooled answer. Mix them, mean in the index and &lt;code&gt;cls&lt;/code&gt; on the query path, and it keeps 70%: more than consistent &lt;code&gt;cls&lt;/code&gt;, because one side of the pair is still the pooling the model was trained with. A client that reimplements pooling has changed the model.&lt;/p&gt;

&lt;p&gt;There's an interactive version of the next bit that doesn't survive being pasted into a feed, so it lives on &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;the original post&lt;/a&gt;: drag the reindex forward, then flip which model encodes the queries. On the new model, recall climbs as the table migrates. On the old one it starts at 0.77, because two consistent models already disagree on a quarter of the answers, then falls through the middle, to 0.63, gets back to 0.77 only in the last few lines, and finishes below it, at 0.73: by the end, every line in the index was written by a model the query path isn't using.&lt;/p&gt;

&lt;p&gt;That second curve is the one production runs. The index is a batch job; the query encoder is a deployed service; the two ship on different days.&lt;/p&gt;

&lt;h2&gt;
  
  
  What the ML engineer watches
&lt;/h2&gt;

&lt;p&gt;Four habits, all cheap next to a reindex.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;The embedder is part of the index.&lt;/strong&gt; Store the model's hash, its pooling mode, its normalisation step, and its sequence cap in the index metadata. Refuse a query whose encoder doesn't match. The &lt;code&gt;modules.json&lt;/code&gt; above is half of that contract; the pooling mode sits next to it in &lt;code&gt;1_Pooling/config.json&lt;/code&gt;, and the hash you add yourself. The cap matters: the L6 model reads 256 tokens and the L12 upgrade reads 128, so the swap that broke the index also halved what each vector sees, silently. Ship all of it with the index, not in a client's config.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Cut over, don't roll.&lt;/strong&gt; Build the new index beside the old one, switch queries and index together behind one flag, and delete the old one afterwards. A rolling reindex under a live query path is the falling curve.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Measure neighbour overlap, not cosine.&lt;/strong&gt; Keep a fixed set of queries with their expected top-k, and check recall@k against it on every model change. A cosine threshold tuned on one model means nothing on the next; the neighbour set is the only observable that survives a swap. It measures drift, not relevance: the migrated top five for the logout query still has a cold-storage line in it.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Dimension is a bill.&lt;/strong&gt; N times d times four bytes, before the index structure. Half-precision halves it; some newer models let you keep only the first k dimensions. Both can be applied to vectors already stored, without re-embedding; truncated vectors have to be renormalised and queries cut to the same k, and both choices belong in the index metadata with the rest. Price a larger model's dimension against the corpus you'll have in two years, not the one you have now.&lt;/p&gt;

&lt;p&gt;None of this is modelling work either. An embedding index looks like a database with a float column, right up to the day someone upgrades the model that wrote the column.&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Part 2 of twelve. An embedding is a row of a matrix, filled either by the gradient or by a model's output, and the space it lives in has no migration: change the model and every stored vector goes stale, silently. A batch of sequences of those rows is a tensor with the signature &lt;code&gt;[B, T, d]&lt;/code&gt;. &lt;a href="https://cchinchilla.dev/posts/fcw-tensors-shapes-type-system/" rel="noopener noreferrer"&gt;The next part&lt;/a&gt; reads that signature the way you'd read a function's, and finds that the runtime never checks it.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Written with AI assistance. See &lt;a href="https://cchinchilla.dev/ai-policy/" rel="noopener noreferrer"&gt;AI policy&lt;/a&gt;.&lt;/em&gt;&lt;/p&gt;

</description>
      <category>machinelearning</category>
      <category>python</category>
      <category>ai</category>
      <category>beginners</category>
    </item>
    <item>
      <title>A model doesn't read text: what a tokenizer decides for you</title>
      <dc:creator>Carlos Chinchilla Corbacho</dc:creator>
      <pubDate>Sun, 20 Sep 2026 19:14:20 +0000</pubDate>
      <link>https://dev.to/cchinchilladev/a-model-doesnt-read-text-what-a-tokenizer-decides-for-you-1f11</link>
      <guid>https://dev.to/cchinchilladev/a-model-doesnt-read-text-what-a-tokenizer-decides-for-you-1f11</guid>
      <description>&lt;p&gt;&lt;em&gt;Originally published at &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;cchinchilla.dev&lt;/a&gt;. Part 1 of From code to weights, a 12-part series on ML fundamentals for engineers.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;A language model has never seen a letter.&lt;/p&gt;

&lt;p&gt;It sees integers. Somewhere between your string and the first matrix multiply there's a function that turns text into a list of ids, and that function was frozen the day the model was trained. It isn't preprocessing. It's part of the artifact. Swap it and the weights are still there, but the model you had is gone.&lt;/p&gt;

&lt;p&gt;I came to this from the compiler side, where the parser is something you own and can change on a Tuesday. Here you own it exactly once, before training, and never again.&lt;/p&gt;

&lt;p&gt;This is the first post of &lt;a href="https://cchinchilla.dev/series/" rel="noopener noreferrer"&gt;a series&lt;/a&gt; for engineers moving into ML engineering, from software or from data science. Each part names one thing that does not behave the way software behaves. This one is the parser: it ships inside the model.&lt;/p&gt;

&lt;h2&gt;
  
  
  A lexer with a different goal
&lt;/h2&gt;

&lt;p&gt;If you've written a compiler, you've written a tokenizer. The lexer takes characters and emits a stream of symbols: identifiers, keywords, literals. Its goal is to recognise structure so the parser can build a tree.&lt;/p&gt;

&lt;p&gt;A model's tokenizer has the same shape and a different goal. There's no structure to recognise. The goal is to compress: turn a string into as few ids as possible from a fixed vocabulary, so the model sees fewer positions and the vocabulary stays small enough to put a row of the embedding matrix behind every id.&lt;/p&gt;

&lt;p&gt;Both extremes fail. Characters give a tiny vocabulary and sequences that are far too long; every position costs attention, and attention costs the square of the length. Whole words give short sequences and an unbounded vocabulary, plus a hole for every word that isn't in it. The data scientist knows this hole as out-of-vocabulary; the software engineer knows it as an unhandled case.&lt;/p&gt;

&lt;p&gt;Subwords are the compromise. &lt;a href="https://arxiv.org/abs/1508.07909" rel="noopener noreferrer"&gt;Sennrich, Haddow and Birch (2016)&lt;/a&gt; took byte-pair encoding, a compression algorithm from 1994, and pointed it at this problem instead: learn the sub-word units from the corpus rather than declaring them up front.&lt;/p&gt;

&lt;h2&gt;
  
  
  BPE is one loop
&lt;/h2&gt;

&lt;p&gt;I read &lt;a href="https://github.com/openai/tiktoken" rel="noopener noreferrer"&gt;&lt;code&gt;tiktoken&lt;/code&gt;&lt;/a&gt;'s &lt;code&gt;_educational.py&lt;/code&gt;. Its encoder is a readable stand-in for the Rust core, which I checked by running both over every sample below and comparing ids. Its trainer is a stand-in for nothing: the Rust core ships no trainer, so &lt;code&gt;bpe_train&lt;/code&gt; (&lt;code&gt;_educational.py:119-185&lt;/code&gt;) is the only version of the loop in the package. It is one loop. The body of &lt;code&gt;bpe_train&lt;/code&gt;:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;while&lt;/span&gt; &lt;span class="nf"&gt;len&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;ranks&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;&amp;lt;&lt;/span&gt; &lt;span class="n"&gt;vocab_size&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
    &lt;span class="n"&gt;stats&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;collections&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nc"&gt;Counter&lt;/span&gt;&lt;span class="p"&gt;()&lt;/span&gt;
    &lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;piece&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="n"&gt;words&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
        &lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;pair&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="nf"&gt;zip&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;piece&lt;/span&gt;&lt;span class="p"&gt;[:&lt;/span&gt;&lt;span class="o"&gt;-&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="n"&gt;piece&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;:]):&lt;/span&gt;
            &lt;span class="n"&gt;stats&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;pair&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;+=&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;
    &lt;span class="n"&gt;most_common_pair&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nf"&gt;max&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;stats&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;key&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;stats&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;])&lt;/span&gt;
    &lt;span class="n"&gt;token_bytes&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;most_common_pair&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;most_common_pair&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
    &lt;span class="n"&gt;token&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nf"&gt;len&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;ranks&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
    &lt;span class="c1"&gt;# Add the new token!
&lt;/span&gt;    &lt;span class="n"&gt;ranks&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;token_bytes&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;token&lt;/span&gt;
    &lt;span class="c1"&gt;# ... then replace that pair through the corpus, and loop
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Start with 256 tokens, one per byte. Count every adjacent pair across the corpus. Take the most frequent pair, give it the next id, replace it through the corpus left to right. Repeat until the vocabulary is the size you asked for.&lt;/p&gt;

&lt;p&gt;The vocabulary is the list of merges, in order. That order is the whole model of the language. Encoding replays it: split the input into bytes, find the adjacent pair with the lowest merge rank, merge it, repeat until nothing merges. &lt;code&gt;bpe_encode&lt;/code&gt; (&lt;code&gt;_educational.py:83-116&lt;/code&gt;) is thirty-four lines. It never iterates the merge list; training takes the most frequent pair and encoding takes the lowest-ranked one, which comes to the same thing.&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;The original has an interactive widget that replays the merge loop on a toy corpus, one merge per click.&lt;/strong&gt; &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Open it on cchinchilla.dev&lt;/a&gt; to use it.&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;The widget runs that loop on a toy corpus, one merge per click. Twelve of its eighteen steps are ties, so the tie-break rule (first pair seen wins) is doing real work.&lt;/p&gt;

&lt;p&gt;Two details matter for the engineer. First, it works on bytes, not characters, so any input encodes and there's no unknown token. A character the merges never saw costs up to one token per UTF-8 byte instead of raising. Under GPT-2, &lt;code&gt;ß&lt;/code&gt; is one token because its two bytes merged during training, &lt;code&gt;漢&lt;/code&gt; is three, one per byte, and a Gothic letter is four. Newer vocabularies bought some of that back; &lt;code&gt;o200k_base&lt;/code&gt; has &lt;code&gt;漢&lt;/code&gt; whole. The Gothic letter still costs four everywhere. Second, before BPE runs there is a regex that splits the text into chunks, and merges never cross chunk boundaries. GPT-2's pattern keeps a leading space attached to the word that follows it, so &lt;code&gt;" token"&lt;/code&gt; and &lt;code&gt;"token"&lt;/code&gt; are different tokens. That's why the model treats the start of a line differently from the middle of one, and why a stray space in a prompt template changes the output.&lt;/p&gt;

&lt;h2&gt;
  
  
  The vocabulary is the training corpus
&lt;/h2&gt;

&lt;p&gt;GPT-2's vocabulary has 50,257 entries. That number isn't round because it's a sum: 256 byte tokens, 50,000 learned merges, and one special token, &lt;code&gt;&amp;lt;|endoftext|&amp;gt;&lt;/code&gt;, id 50256. &lt;a href="https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf" rel="noopener noreferrer"&gt;Radford et al. (2019)&lt;/a&gt; describe the byte-level design those numbers come out of. The two newer OpenAI encodings I checked, &lt;code&gt;cl100k_base&lt;/code&gt; and &lt;code&gt;o200k_base&lt;/code&gt;, hold 100,261 and 200,000 tokens spread over id ranges of 100,277 and 200,019 (tiktoken's &lt;code&gt;n_vocab&lt;/code&gt;); the gaps are unused ids.&lt;/p&gt;

&lt;p&gt;The merges encode the training corpus. GPT-2 was trained mostly on English web text, so English words are cheap: &lt;code&gt;" tokenizer"&lt;/code&gt; is two tokens. The same concept in Spanish, &lt;code&gt;" tokenizador"&lt;/code&gt;, is three. &lt;code&gt;" desvanecimiento"&lt;/code&gt; is six. A German word with an umlaut, &lt;code&gt;" Wörter"&lt;/code&gt;, is three; one of those three is the two bytes of &lt;code&gt;ö&lt;/code&gt; welded to the &lt;code&gt;r&lt;/code&gt; after it. The tokenizer has no notion of a character; it has bytes and frequencies.&lt;/p&gt;

&lt;p&gt;Numbers are the case that surprises engineers. GPT-2 tokenizes &lt;code&gt;1234567&lt;/code&gt; as &lt;code&gt;123&lt;/code&gt;, &lt;code&gt;45&lt;/code&gt;, &lt;code&gt;67&lt;/code&gt;, because that's what the merge ranks happen to produce. &lt;code&gt;cl100k_base&lt;/code&gt; changed the regex to cap digit runs at three, which bounds the damage without fixing it: the cap groups left to right, so &lt;code&gt;1000&lt;/code&gt; becomes &lt;code&gt;100&lt;/code&gt; and &lt;code&gt;0&lt;/code&gt;. That's a tokenizer decision with a direct effect on what the model can compute, and it was made in a regex.&lt;/p&gt;

&lt;p&gt;Be fair about why. GPT-2's merges were learned from frequency alone, and in web text &lt;code&gt;1000&lt;/code&gt; is common enough that &lt;code&gt;00&lt;/code&gt;, then &lt;code&gt;000&lt;/code&gt;, then &lt;code&gt;1000&lt;/code&gt; each won the pair count in turn; a merge table that respects frequency will swallow it whole. Nobody chose to break arithmetic. The cost only became visible once people started asking these models to add.&lt;/p&gt;

&lt;h2&gt;
  
  
  The same paragraph costs more in Spanish
&lt;/h2&gt;

&lt;p&gt;I took one paragraph of 184 bytes, wrote it in English and in Spanish, and built a JSON payload of roughly the same size, 182 bytes. Then I encoded all three with the three tokenizers.&lt;/p&gt;

&lt;p&gt;Spanish costs about 1.7× English under GPT-2 and still 1.2× under the newest encoding. JSON costs about 2.3× under all three: braces, quotes, and keys do not compress. The 26 quotes in that payload never appear as a token on their own; each one is welded to the brace, bracket, colon, comma, or space beside it, in tokens like &lt;code&gt;"&lt;/code&gt;, &lt;code&gt;":&lt;/code&gt;, &lt;code&gt;",&lt;/code&gt;, &lt;code&gt;{"&lt;/code&gt;, &lt;code&gt;"},&lt;/code&gt;. Twenty-six quotes, twenty-six tokens carrying one.&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;The original has an interactive widget where you switch the sample text and the encoding.&lt;/strong&gt; &lt;a href="https://cchinchilla.dev/posts/fcw-tokens-models-dont-read-text/" rel="noopener noreferrer"&gt;Open it on cchinchilla.dev&lt;/a&gt; to use it.&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;Switch the sample and the encoding. Newer vocabularies mostly close the Spanish gap; the JSON one does not move.&lt;/p&gt;

&lt;p&gt;Three things follow, and all three are billing or capacity problems, not modelling problems. A context window quoted in tokens holds less Spanish than English, and much less JSON. A per-token price is a per-language price. And a prompt template that wraps every request in structured metadata is paying the JSON rate on every call.&lt;/p&gt;

&lt;h2&gt;
  
  
  Nothing in the type system says so
&lt;/h2&gt;

&lt;p&gt;In software, the parser is code and the model is data. You can change the parser, rebuild, and read the same data. Here the parser is coupled to the weights. &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;The embedding matrix&lt;/a&gt; has one row per id, and the meaning of row 11,241 is &lt;code&gt;" token"&lt;/code&gt; only because that's what GPT-2's merges said when the model was trained. Under &lt;code&gt;cl100k_base&lt;/code&gt; the same row is &lt;code&gt;Seg&lt;/code&gt;. Encode the same text with a different tokenizer and you get ids that index the same matrix and mean something else. Whether that is loud or silent is luck. The same English paragraph under &lt;code&gt;o200k_base&lt;/code&gt; reaches id 99,665, which a 50,257-row matrix refuses. Go the other way, GPT-2 ids into a model with a larger vocabulary, and every id lands. Nothing throws. Even the refusal is softer than it looks: production models pad the embedding to a multiple of 64 or 128 for tensor-core alignment, so the matrix has more rows than the vocabulary has ids, and on a GPU an out-of-range index surfaces as an asynchronous assert at some unrelated line rather than where it happened.&lt;/p&gt;

&lt;p&gt;This is what makes the tokenizer part of the artifact. In GPT-2 that same matrix is also the output head, so a tokenizer swap corrupts the lookup on the way in and the sampling distribution on the way out, from one array of weights. The model file and the tokenizer file are two halves of one thing, and the failure mode when they drift is silent: a full swap whose ids stay in range produces garbage that no check flags, and a partial drift (a renumbered special token, a missing BOS) produces fluent text conditioned on the wrong input.&lt;/p&gt;

&lt;h2&gt;
  
  
  What the ML engineer watches
&lt;/h2&gt;

&lt;p&gt;Four things to check before any of this reaches production.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Tokenizer and model versioned together.&lt;/strong&gt; Store the tokenizer's hash next to the weights and check it at load time. A mismatch that doesn't fail on load is the worst kind.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Cost per language and per domain.&lt;/strong&gt; Measure it on your real traffic before you estimate a bill or promise a context length. The 8k window is 8k English tokens; for the JSON payload above it is closer to 3.5k of the same content.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Special tokens.&lt;/strong&gt; tiktoken's &lt;code&gt;Encoding.encode&lt;/code&gt; raises by default if the text contains something that looks like &lt;code&gt;&amp;lt;|endoftext|&amp;gt;&lt;/code&gt;, precisely because a user could type it. That default is a security decision; check what your serving layer does with it, and where your chat template inserts its own.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Length in tokens, not characters.&lt;/strong&gt; Every limit the model has is in tokens. Truncation, context budgeting, and rate limits all need the tokenizer in the request path, not a character count.&lt;/p&gt;

&lt;p&gt;None of this is modelling work. It's the part of the model that behaves most like ordinary software, which is exactly why it ships without review: a function that turns strings into integers looks like preprocessing right up until the moment it isn't.&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Part 1 of twelve. This part starts before the model exists: the tokenizer ships inside the artifact rather than in front of it, and nothing in the type system says so. A token id is an index; &lt;a href="https://cchinchilla.dev/posts/fcw-embeddings-lookup-table/" rel="noopener noreferrer"&gt;the next part&lt;/a&gt; reads what sits at that index, the embedding matrix, which is a lookup table and not much else.&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Written with AI assistance. See &lt;a href="https://cchinchilla.dev/ai-policy/" rel="noopener noreferrer"&gt;AI policy&lt;/a&gt;.&lt;/em&gt;&lt;/p&gt;

</description>
      <category>machinelearning</category>
      <category>python</category>
      <category>ai</category>
      <category>beginners</category>
    </item>
  </channel>
</rss>
