Tensors, Weights, Layers, and State

Most arrays in a model are nested, and an array of arrays is something you already write. Here is a small one, built from plain Python lists:

# two prompts, three tokens each, four numbers per token
batch = [
    [[1, 0, 2, 5], [3, 1, 0, 0], [2, 2, 1, 4]],  # prompt 0
    [[1, 0, 2, 5], [3, 1, 0, 0], [7, 0, 3, 1]],  # prompt 1
]
print("prompts:          ", len(batch))
print("tokens per prompt:", len(batch[0]))
print("numbers per token:", len(batch[0][0]))
print("shape:", (len(batch), len(batch[0]), len(batch[0][0])))
print("prompt 1, token 2:", batch[1][2])

$ python3 nested.py
prompts:           2
tokens per prompt: 3
numbers per token: 4
shape: (2, 3, 4)
prompt 1, token 2: [7, 0, 3, 1]

The outer list holds two prompts. Each prompt holds three tokens, and each token is a list of four numbers. The two prompts differ only in their last token, the way "The capital of France is" and "The capital of Spain is" differ in one word. Asking len at each level reads off what that level counts: 2 prompts, 3 tokens per prompt, 4 numbers per token. batch[1][2] walks down two levels to prompt 1's token 2, which is one token's list of four numbers.

Each level of nesting has a name. An axisTerm axisOne level of a nested array. The first index you write picks along the first axis, the next index along the second. is one level of a nested array. The first index you write picks along the first axis, the next index along the second. The shapeTerm shapeThe sizes of a nested array, from the outermost list in. Each size counts one thing, such as prompts, tokens, or numbers per token. is the sizes of a nested array, from the outermost list in. Each size counts one thing, such as prompts, tokens, or numbers per token. Here the shape is (2, 3, 4). A tensorTerm tensorWhat libraries call an array of numbers, often nested: a list of lists of lists. Every number in a running model sits in one. is what libraries call an array of numbers, often nested: a list of lists of lists. Every number in a running model sits in one. The word sounds more technical than the thing it names.

The three axes of a model's arrays have names of their own:

• The batchTerm batchSeveral independent prompts stacked into one array, so one call handles them all. is several independent prompts stacked into one array, so one call handles them all. The batch axis says which prompt.
• The positionTerm positionWhere a token sits in the sequence: 0 for the first token, 1 for the next. is where a token sits in the sequence: 0 for the first token, 1 for the next. The position axis says which token.
• The widthTerm widthHow many numbers describe one token inside the model: 1,024 here. is how many numbers describe one token inside the model: 1,024 here. The width axis says which of a token's numbers.

Why nest them at all? Chapter 1 showed that the model turns each token into a list of numbers. An embeddingTerm embeddingThe list of 1,024 numbers the model looks up for a token ID: one row of a big learned table. is the list of 1,024 numbers the model looks up for a token ID: one row of a big learned table. A model handles many tokens in one call, and often several prompts. Lists of the same length stack into lists of lists, so one call can carry every token of every prompt. Stacking prompts does not mix them, as the next section shows.

Now the real model. The scripts in this chapter import the same load.py as chapter 1. It loads the pinned revision of the model with the Hugging Face Transformers library, which runs on PyTorch:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

REPO = "LiquidAI/LFM2.5-230M"
REV = "40cb2ad3b3044d5a41eee083a6103c8b523afa45"
tok = AutoTokenizer.from_pretrained(REPO, revision=REV)
model = AutoModelForCausalLM.from_pretrained(
    REPO, revision=REV, dtype=torch.bfloat16).eval()

The next script sends two prompts in one call. It prints the sizes of three arrays, then each prompt's top next token:

import torch
from load import tok, model

texts = ["The capital of France is",
         "The capital of Spain is"]
ids = tok(texts, return_tensors="pt").input_ids
with torch.no_grad():
    out = model(ids, output_hidden_states=True)
emb = out.hidden_states[0]          # the embedding output
print("token IDs ", tuple(ids.shape))
print("embeddings", tuple(emb.shape))
print("logits    ", tuple(out.logits.shape))
for b, text in enumerate(texts):
    top = int(out.logits[b, -1].argmax())  # prompt b, last position
    print(f"{text!r:28} -> {tok.decode([top])!r}")

$ python3 shapes.py
token IDs  (2, 6)
embeddings (2, 6, 1024)
logits     (2, 6, 65536)
'The capital of France is'   -> ' Paris'
'The capital of Spain is'    -> ' Madrid'

The first two lines are the same levels as nested.py. The token IDs are 2 prompts of 6 tokens each. The six tokens include the start token, a special token with ID 1 that the tokenizer puts before the text. The embeddings add a third level, the width. So that array is 2 prompts, 6 tokens each, and 1,024 numbers per token. The third line holds the scores, and the next section reads it.

A shape on its own does not say what an array is. (2, 3, 4) could count 2 images of 3 rows of 4 pixels, or 2 users with 3 orders of 4 items. The sizes mean something only once you say what each one counts. Before you index an array, say what each of its sizes counts.

The model's last array holds its scoresTerm scoresThe raw numbers the output layer gives every possible next token, also called logits.: the raw numbers the output layer gives every possible next token, also called logits. The vocabularyTerm vocabularyEvery token the model can score: 65,536 of them for this model. is every token the model can score: 65,536 of them for this model. Chapter 1 turned the last token's scores into probabilities and picked the next token from them. This section shows where that last list sits among the others. Here is a small version with a vocabulary of four words:

words = ["of", "is", "Paris", "Madrid"]
# scores[prompt][position][word]: one score per word
scores = [
    [[5, 1, 0, 0], [0, 6, 1, 0], [0, 0, 9, 2]],  # France
    [[5, 1, 0, 0], [0, 6, 1, 0], [0, 0, 2, 9]],  # Spain
]
def best(row):
    return words[row.index(max(row))]

print("best at each position, Spain:",
      [best(row) for row in scores[1]])
right = [p[-1] for p in scores]   # every prompt, last position
wrong = scores[-1]                # last prompt, every position
print("right:", len(right), "rows", [best(r) for r in right])
print("wrong:", len(wrong), "rows", [best(r) for r in wrong])

$ python3 slice.py
best at each position, Spain: ['of', 'is', 'Madrid']
right: 2 rows ['Paris', 'Madrid']
wrong: 3 rows ['of', 'is', 'Madrid']

The scores have the same three levels as nested.py, with a different inner level. Each prompt has three positions, and each position holds a list of four scores, one per word in the vocabulary. The best word at a position is the one with the highest score, which best finds as the index of the largest.

The first line of output is the best word at each position of the Spain prompt: of, is, Madrid. Read the Spain prompt's three positions as capital, Spain and is: after capital the best guess is of, after Spain it is is, and after is it is Madrid. Each position's list guesses the word that comes after that position. The next word of the reply is the guess at the last position, the word after the whole prompt.

The right line asks for that guess in every prompt at once. [p[-1] for p in scores] takes each prompt and keeps its last position. It returns 2 rows, one per prompt, and their best words are Paris and Madrid. A tensor library writes the same slice as scores[:, -1]. The colon means every entry on the first axis, and the -1 picks the last entry on the second. The wrong line is the next section's subject.

Why does the model score every position when the reply needs only the last? The output layer treats every token's list the same way, so one call gives a guess after each position. The reply loop reads only the last one.

On the real model, the third line of shapes.py's output is the scores. They are 2 prompts, 6 positions each, and at every position a score for each of the 65,536 tokens in the vocabulary. out.logits[b, -1] takes prompt b's last position and keeps all 65,536 of its scores, and argmax picks the highest. Prompt 0 gets Paris, and prompt 1 gets Madrid.

Batching did not mix the two prompts. In a separate run we turned each prompt's last scores into probabilities, once in the batch and once alone. France gives Paris at 95.22 percent both ways, and Spain gives Madrid at 93.66 percent both ways. France's last scores in the batch were identical to its scores alone.

A natural first guess is that the scores hold one answer per prompt. They hold one list per position. Six tokens in gives six lists of scores out, for each prompt.

The right slice took one index. Index the other level by mistake and the code still runs. The wrong line of slice.py takes scores[-1], which indexes the outermost level. That picks the last prompt and keeps all of its positions. Its output is three real words, of, is and Madrid. They are the guesses after each position of the Spain prompt, and none of them answers what comes next in each prompt. The only sign is the size: 3 rows where 2 belong.

The real model makes the same slip look even more believable:

import torch
from load import tok, model

texts = ["The capital of France is",
         "The capital of Spain is"]
ids = tok(texts, return_tensors="pt").input_ids
with torch.no_grad():
    logits = model(ids).logits          # (batch, position, vocabulary)
right = logits[:, -1]   # every prompt, last position
wrong = logits[-1]      # last prompt, every position
print("right", tuple(right.shape))
print([tok.decode([int(i)]) for i in right.argmax(-1)])
print("wrong", tuple(wrong.shape))
print([tok.decode([int(i)]) for i in wrong.argmax(-1)])

$ python3 wrong.py
right (2, 65536)
[' Paris', ' Madrid']
wrong (6, 65536)
['<|startoftext|>', ' ', ' of', ' France', ' is', ' Madrid']

logits[:, -1] keeps every batch entry and takes the last position. That gives one next token per prompt, Paris and Madrid. logits[-1] indexes the first axis instead. It takes the last batch entry, the Spain prompt, and keeps all six positions. That gives six predictions, one after each position of the Spain prompt.

Both lines run without an error, and both print real tokens. The wrong one even looks plausible. Its fourth entry is the model's guess after "The capital of", and the guess is France, at about 19 percent. A bug like this can slip past a quick look at the output.

The shape catches it. You wanted one row per prompt, so the first size should be 2. It is 6. When you index an array, write down the shape you expect, then compare it with the shape you got.

So far every array was made by running the model. The learned arrays are loaded from its file. Chapter 1 counted 229,693,184 parameters. ParametersTerm parametersThe model's learned numbers, also called weights. Training sets them. Running the model only reads them. are the model's learned numbers, also called weights. Training sets them. Running the model only reads them. They sit in named arrays. The next script lists them, then prints the names and sizes of everything outside the blocks, plus one conv block and one attention block:

from load import model

params = list(model.named_parameters())
print(len(params), "tensors")
for name, p in params:
    block = name.split(".")[2] if ".layers." in name else ""
    if block in ("", "0", "2"):          # one conv block, one attention block
        print(f"{name.removeprefix('model.'):38} {tuple(p.shape)}")

$ python3 names.py
132 tensors
embed_tokens.weight                    (65536, 1024)
layers.0.conv.conv.weight              (1024, 1, 3)
layers.0.conv.in_proj.weight           (3072, 1024)
layers.0.conv.out_proj.weight          (1024, 1024)
layers.0.feed_forward.w1.weight        (2560, 1024)
layers.0.feed_forward.w3.weight        (2560, 1024)
layers.0.feed_forward.w2.weight        (1024, 2560)
layers.0.operator_norm.weight          (1024,)
layers.0.ffn_norm.weight               (1024,)
layers.2.self_attn.q_proj.weight       (1024, 1024)
layers.2.self_attn.k_proj.weight       (512, 1024)
layers.2.self_attn.v_proj.weight       (512, 1024)
layers.2.self_attn.out_proj.weight     (1024, 1024)
layers.2.self_attn.q_layernorm.weight  (64,)
layers.2.self_attn.k_layernorm.weight  (64,)
layers.2.feed_forward.w1.weight        (2560, 1024)
layers.2.feed_forward.w3.weight        (2560, 1024)
layers.2.feed_forward.w2.weight        (1024, 2560)
layers.2.operator_norm.weight          (1024,)
layers.2.ffn_norm.weight               (1024,)
embedding_norm.weight                  (1024,)

The model's 229,693,184 learned numbers sit in 132 arrays. A name reads like a path into nested objects, the way a JSON path does. Take layers.2.self_attn.q_proj.weight:

• layers.2 is block 2, the third block. A blockTerm blockOne repeated stage that reworks every token's numbers. This model stacks 14 of them. is one repeated stage that reworks every token's numbers. This model stacks 14 of them. The code calls them layers.
• self_attn is that block's attention part, its mixer.
• q_proj is one of attention's grids of learned numbers. Chapter 8 opens what each of them does.
• weight is the array of learned numbers in that part.

Every block has a mixer. A mixerTerm mixerThe part of a block where tokens share information with each other. is the part of a block where tokens share information with each other. This model has two kinds. AttentionTerm attentionA mixer that lets each token pull information from every earlier token, weighted by content. is a mixer that lets each token pull information from every earlier token, weighted by content. ConvTerm convA mixer that blends each token with the two tokens just before it (short for convolution). is a mixer that blends each token with the two tokens just before it (short for convolution). Six of the 14 blocks use attention and the other eight use conv.

Not every array is a grid of rows. The embedding table is read as a lookup: each token ID selects one of its 65,536 rows. The norms are single lists of 1,024 or 64 numbers, and chapter 8 opens them. One name misleads. embedding_norm.weight is the final norm, which the model applies after the last block, not after the embeddings. A name is a strong hint, and the model's code is the check.

The word "layer" changes meaning from place to place. A layerTerm layerA word with two meanings. In this model's code and config a layer is a block. People also use the word for one grid of learned numbers inside a block. is a word with two meanings. In this model's code and config a layer is a block. People also use the word for one grid of learned numbers inside a block. The config's num_hidden_layers is 14, and it counts blocks, each holding 8 or 11 arrays. When someone says "layer", read the name to see which one they mean.

Why does every array need a name? Every block of one kind repeats the same arrays at the same sizes. Block 0 is a conv block, whose mixer blends each token with the two tokens just before it, and every conv block holds the same 8 arrays: the conv filter (a learned row of 3 numbers for each of the 1,024 numbers, one per token it blends), a grid on each side of it, the three feed-forward grids and two norms. Block 2 is an attention block, and every attention block holds the same 11. The name is the only way to tell block 2's q_proj from block 4's. With the embedding table and the final norm, that makes 1 + 8 × 8 + 6 × 11 + 1 = 132 arrays.

Most of the arrays in names.py have two sizes, which makes each one a grid. A grid's two sizes count different things. In code, a grid is a list of lists, and the script below reads the two lengths of a small one, two rows of three numbers:

w = [[2, 0, 1],                   # row 0
     [1, 3, 0]]                   # row 1
print("rows (numbers out):     ", len(w))
print("row length (numbers in):", len(w[0]))
print("takes a list of", len(w[0]), "numbers, gives", len(w))

$ python3 grid.py
rows (numbers out):      2
row length (numbers in): 3
takes a list of 3 numbers, gives 2

The outer length counts the rows, 2. The length of one row is 3. In a model, each row of such a grid makes one output number from the list it takes in, so a grid with 2 rows gives 2 numbers. Each row is as long as the list it takes in, so this grid takes 3. A learned grid like this has a name. A weight matrixTerm weight matrixA grid of numbers the model learned in training. Each row makes one output number and is as long as the list it takes in. is a grid of numbers the model learned in training. Each row makes one output number and is as long as the list it takes in. How a row turns a list into its one number is the whole of chapter 3. Here, read only the sizes: rows out, then row length in.

The real model's grids read the same way, because PyTorch stores each one with one row per output number. w1 is 2,560 rows of 1,024. It takes a token's 1,024 numbers and gives 2,560. w2 is 1,024 rows of 2,560, so it brings them back to 1,024. The two hold the same count of numbers, but the order of the sizes says which one widens a token's list and which narrows it back. That widening and narrowing is the feed-forwardTerm feed-forwardThe part of a block that reworks each token's numbers on its own, widening them and narrowing back.: the part of a block that reworks each token's numbers on its own, widening them and narrowing back.

Who decided 2,560, or 1,024? In code, the answer is a constructor that allocates arrays from a dict of settings. The script below builds a small set of arrays from four settings. Then it edits one setting and tries to load the saved arrays into the new sizes:

def grid(rows, cols):             # an empty grid of zeros
    return [[0] * cols for _ in range(rows)]

def build(config):                # every array, sized by config
    w, inner = config["width"], config["inner"]
    arrays = {"table": grid(config["vocab"], w)}
    for b in range(config["blocks"]):
        arrays[f"layers.{b}.w1"] = grid(inner, w)
        arrays[f"layers.{b}.w2"] = grid(w, inner)
    return arrays

def size(a):
    return f"{len(a)} rows of {len(a[0])}"

config = {"vocab": 5, "width": 3, "inner": 4, "blocks": 2}
saved = build(config)             # stands in for the saved file
for name, a in saved.items():
    print(f"{name:11} {size(a)}")

config["inner"] = 6               # edit one setting, then load
for name, a in build(config).items():
    if size(a) != size(saved[name]):
        print("cannot load", name)
        print("  saved:", size(saved[name]))
        print("  config wants:", size(a))
        break

$ python3 build.py
table       5 rows of 3
layers.0.w1 4 rows of 3
layers.0.w2 3 rows of 4
layers.1.w1 4 rows of 3
layers.1.w2 3 rows of 4
cannot load layers.0.w1
  saved: 4 rows of 3
  config wants: 6 rows of 3

build makes empty grids of zeros, sized only by the settings. It makes a table of 5 rows of 3, and in each of 2 blocks a w1 of 4 rows of 3 and a w2 of 3 rows of 4. The first build stands in for the saved file. Then inner changes from 4 to 6, and the code builds again. The first array whose size differs is layers.0.w1. It was saved as 4 rows of 3, and the edited settings want 6 rows of 3. The saved numbers have nowhere to go, so loading stops.

The settings hold no learned numbers. They only say how big each array is. The learned numbers fill the arrays, and they fit only the sizes they were saved at.

Real models are built this way for a reason. The same code runs this design at any size, and a short file picks the size. That is why a model ships with that file next to its weights. The configTerm configThe model's description file, config.json. Its numbers, like the width and the number of blocks, set the size of every learned array. It holds no learned numbers. is the model's description file, config.json. Its numbers, like the width and the number of blocks, set the size of every learned array. It holds no learned numbers. Here is the real model's:

from load import model

c = model.config
for k in ("vocab_size", "hidden_size", "intermediate_size",
          "num_hidden_layers", "num_attention_heads",
          "num_key_value_heads", "conv_L_cache"):
    print(f"{k:20} {getattr(c, k)}")

$ python3 config.py
vocab_size           65536
hidden_size          1024
intermediate_size    2560
num_hidden_layers    14
num_attention_heads  16
num_key_value_heads  8
conv_L_cache         3

The config is not the settings file from chapter 1. The settings fileTerm settings fileThe file shipped with a model that sets default picking rules and the stop token (generation config). is the file shipped with a model that sets default picking rules and the stop token (generation config). This model ships both. generation_config.json holds the picking rules and the stop token, and config.json holds the sizes.

Match each config number to the sizes in names.py:

The conv filter is 3 positions wide because conv blends each token with the two before it, three tokens in all. The two attention settings set the sizes of attention's own grids, such as k_proj's 512 rows, and chapter 8 opens them. Chapter 10 opens the conv block and its other grids.

The tempting mistake is to edit the config to change the model. The next script reads the real config, sets intermediate_size to 3,072, and loads the same weights under it:

import textwrap
import torch
from transformers import AutoConfig, AutoModelForCausalLM
from load import REPO, REV

cfg = AutoConfig.from_pretrained(REPO, revision=REV)
print("intermediate_size in config:", cfg.intermediate_size)
cfg.intermediate_size = 3072          # edit one setting
try:
    AutoModelForCausalLM.from_pretrained(
        REPO, revision=REV, config=cfg, dtype=torch.bfloat16)
    print("loaded")
except RuntimeError as e:
    print(textwrap.fill("RuntimeError: " + str(e), 60))

$ python3 wrongsize.py
intermediate_size in config: 2560
RuntimeError: You set `ignore_mismatched_sizes` to `False`,
thus raising an error. For details look at the above report!

Loading stopped with an error. With the library's logging turned on, its load report names the arrays that did not fit: w1, w2 and w3 in all 14 blocks. Each was saved with 2,560 where the edited config wanted 3,072. A bigger model needs new learned numbers, because the old ones only fit the sizes they were trained at.

Every learned array has a name and a size. Adding them all up should give the model's number of learned numbers. Done the obvious way, it does not. The reason is aliasing: two names for one list, which Python's is and id() can detect. Here it is in plain Python:

table = [[1, 0, 2], [0, 3, 1]]    # 6 numbers
saved = {"embed": table,          # the same list, two names
         "output": table}
by_name = sum(len(r) for t in saved.values() for r in t)
print("counted by name:", by_name)
print("same list:", saved["embed"] is saved["output"])
once = {id(t): t for t in saved.values()}   # one entry per list
print("counted once:  ", sum(len(r) for t in once.values()
                             for r in t))

$ python3 alias.py
counted by name: 12
same list: True
counted once:   6

The dict holds one list under two names, embed and output. Summing over the names counts its 6 numbers twice, which gives 12. saved["embed"] is saved["output"] is True, so both names point to the same list. Keying a dict by id() keeps one entry per distinct list, and that count is 6. Count arrays, not names.

The real model has the same trap. Code that saves or loads a model works with its state_dict. The state_dictTerm state_dictThe dict from name to array that saving code writes out. An array shared by two parts appears once under each name. is the dict from name to array that saving code writes out. An array shared by two parts appears once under each name. The script below counts the parameters once, then sums the state_dict:

from load import model

params = sum(p.numel() for p in model.parameters())
saved = sum(t.numel() for t in model.state_dict().values())
print(f"parameters  {params:>12,}")
print(f"state_dict  {saved:>12,}")
sd = model.state_dict()
print("lm_head.weight listed:", "lm_head.weight" in sd)
print("same storage:", sd["lm_head.weight"].data_ptr()
      == sd["model.embed_tokens.weight"].data_ptr())

$ python3 count.py
parameters   229,693,184
state_dict   296,802,048
lm_head.weight listed: True
same storage: True

The state_dict holds 67,108,864 more numbers, exactly one table of 65,536 rows of 1,024. This model's output layer reuses the embedding table instead of holding its own, and chapter 7 explains why. The state_dict lists that one array under two names, model.embed_tokens.weight and lm_head.weight, so summing it counts the table twice. model.parameters() gives each shared array once, which is how it reaches 229,693,184.

Equal values would not prove the two names are one array, because two separate copies could hold the same numbers. The last line checks the address in memory where each array's numbers start. Both names point to the same address, so they are the same array, as is showed for the small dict.

The learned arrays are what the model holds. Everything else is made while it runs. In code, the difference is a function's local variables against the data it only reads. The function below reads a table of 15 learned numbers. For each word it makes two lists of three: the word's row of the table, and that row added to the previous word's row. It runs on 5 words, then on 50:

table = {"the": [1, 0, 2], "capital": [0, 3, 1],
         "of": [2, 2, 0], "France": [1, 4, 1], "is": [0, 1, 1]}
learned = sum(len(row) for row in table.values())

def call(words):                  # reads the table, makes lists
    made, prev = [], [0, 0, 0]    # made: numbers made this call
    for w in words:
        row = list(table[w])                    # a copied row
        mixed = [a + b for a, b in zip(row, prev)]
        made += [row, mixed]
        prev = row
    return made

before = {w: list(r) for w, r in table.items()}
for text in ("the capital of France is",
             "the capital of France is " * 10):
    n = sum(len(x) for x in call(text.split()))
    print(f"{len(text.split()):>2} words  made {n:>3} numbers"
          f"  learned {learned}")
print("table unchanged:", table == before)

$ python3 made.py
 5 words  made  30 numbers  learned 15
50 words  made 300 numbers  learned 15
table unchanged: True

The function made 30 numbers for 5 words and 300 for 50. What it makes grows with the input. The table stayed at 15 numbers, and the last line confirms that no call changed it. The function only reads it.

The model's arrays split the same way. An activationTerm activationAn array the model computes during one call, like a block's output. Made fresh each call, then dropped. is an array the model computes during one call, like a block's output. Made fresh each call, then dropped. The embeddings, every block's output, the 2,560 numbers per token inside each feed-forward part and the scores at the end are all activations.

The next script counts them. It registers a callback on every layer that holds weights. PyTorch calls the callback each time that layer runs, and here the callback adds up how many numbers the layer hands on. PyTorch names this kind of callback a forward hook. The script runs a 6-token prompt and a 352-token prompt:

import torch
from load import tok, model

made = []                    # numbers each layer hands on
def hook(module, inputs, out):
    made.append(out.numel())
for m in model.modules():    # every layer that holds weights
    if list(m.parameters(recurse=False)):
        m.register_forward_hook(hook)

params = sum(p.numel() for p in model.parameters())
for text in ("The capital of France is",
             "The capital of France is Paris. " * 50):
    ids = tok(text, return_tensors="pt").input_ids
    made.clear()
    with torch.no_grad():
        model(ids)
    print(f"{ids.shape[1]:>4} IDs  activations {sum(made):>12,}"
          f"  parameters {params:,}")

$ python3 activations.py
   6 IDs  activations    1,456,128  parameters 229,693,184
 352 IDs  activations   85,426,176  parameters 229,693,184

The parameters did not change with the prompt. They are 229,693,184 both times. The activations grew with it, from about 1.5 million numbers to about 85 million, which is 242,688 numbers per position both times. Every activation has a position axis, so a longer prompt makes each one longer, and a bigger batch multiplies them again. The parameters have no position axis, so one set serves every prompt. The output layer alone hands on 65,536 numbers per position, more than a quarter of the total, because it scores every vocabulary entry.

These are numbers handed from one layer to the next over the call, not memory in use at one moment. Most activations are released as soon as the next layer has read them. The count also leaves out arrays made inside a single step, such as the grid each attention block builds to compare every token with every earlier one: 576 numbers at 6 tokens, almost 2 million at 352. Chapter 13 measures what a model needs in memory.

Running the model reads the parameters and never writes them. Chapter 1 showed that for one call. The next script makes a stronger test of it. It copies all 132 arrays, generates three slogans with the shipped picking rules, and compares:

import torch
from load import tok, model

before = {n: p.clone() for n, p in model.named_parameters()}
Q = "Write a one-line slogan for a coffee shop."
msgs = [{"role": "user", "content": Q}]
ids = tok.apply_chat_template(
    msgs, add_generation_prompt=True,
    return_tensors="pt", tokenize=True)["input_ids"]
for seed in range(3):
    torch.manual_seed(seed)
    model.generate(ids, max_new_tokens=40)
same = sum(torch.equal(before[n], p)
           for n, p in model.named_parameters())
print(same, "of", len(before), "tensors unchanged")

$ python3 readonly.py
132 of 132 tensors unchanged

Every number is exactly as it was. A model does not learn from the prompts you send it. Training changes parameters, and chapter 6 shows how. Running only makes activations, and, while a reply is written, a cache.

A reply takes many calls. Chapter 1's loopTerm loopCall the model, pick one token, append it to the input, and call again, until a stop token appears. was this: call the model, pick one token, append it to the input, and call again, until a stop token appears. Each call sent the whole text again, so the model redid the work for every earlier token at every step. Can it keep that work instead?

In code, keeping a step's result so a later step does not recompute it is called memoization. Here are two versions of a small loop. Each token's work is to add it to a running total, and the next token is the total's last digit. redo resends every token at every step. remember keeps each token's total and adds only the new token:

prompt = [3, 1, 4]
def step(prev, t):                # one token's work:
    return prev + t               # its running total

def redo(n):                      # resend everything
    seq, work = list(prompt), 0
    for _ in range(n):
        total = 0
        for t in seq:
            total = step(total, t); work += 1
        seq.append(total % 10)
    return seq, work

def remember(n):                  # keep each result: a cache
    seq, cache, work, new = list(prompt), [0], 0, list(prompt)
    for _ in range(n):
        for t in new:
            cache.append(step(cache[-1], t)); work += 1
        new = [cache[-1] % 10]
        seq += new
    return seq, work

for name, f in (("redo", redo), ("remember", remember)):
    seq, work = f(5)
    print(f"{name:8} tokens {seq}  work {work}")

$ python3 redo.py
redo     tokens [3, 1, 4, 8, 6, 2, 4, 8]  work 25
remember tokens [3, 1, 4, 8, 6, 2, 4, 8]  work 7

Both write the same tokens. Redoing costs 25 units of work, and remembering costs 7. Keeping the earlier work changed the cost, not the answer.

What a program keeps from one call to the next has a name. StateTerm stateData the model carries from one call to the next, like the cache. is data the model carries from one call to the next, like the cache. The cacheTerm cacheWhat the model keeps between steps of one reply, so earlier tokens are not recomputed. It belongs to one conversation. is what the model keeps between steps of one reply, so earlier tokens are not recomputed. It belongs to one conversation.

Why can a model keep its earlier work? In this model, each token's numbers depend only on that token and the tokens before it. Attention lets a token look back at earlier tokens, never ahead. The conv part reads a token and the two before it. So appending a token never changes the work already done for the earlier ones, and that work can be kept instead of redone. This is a property of this model's design. Chapter 21 meets a model that does not write left to right.

On the real model, the next script runs 64 steps of a reply from chapter 1's 21-token chat prompt, always taking the top token. It runs them both ways, counts the tokens sent through the model, and checks that the two replies match:

import torch
from load import tok, model

Q = "What is the capital of France? Answer in one sentence."
msgs = [{"role": "user", "content": Q}]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True,
                              return_tensors="pt")["input_ids"]

def run(steps, keep):
    cache, new, sent, toks = None, ids, 0, []
    for _ in range(steps):
        with torch.no_grad():
            out = model(new, past_key_values=cache,
                        use_cache=keep)
        sent += new.shape[1]
        nxt = out.logits[0, -1].argmax().view(1, 1)
        toks.append(int(nxt))
        if keep:              # remember: cache + one token
            cache, new = out.past_key_values, nxt
        else:                 # redo: resend everything
            new = torch.cat([new, nxt], 1)
    return toks, sent

redo, remember = run(64, False), run(64, True)
print("64 steps, redo:     tokens sent", redo[1])
print("64 steps, remember: tokens sent", remember[1])
print("same 64 tokens:", redo[0] == remember[0])

$ python3 saved.py
64 steps, redo:     tokens sent 3360
64 steps, remember: tokens sent 84
same 64 tokens: True

Resending everything sent 3,360 tokens through the model: 21 on the first step, 22 on the next, and so on up to 84. Keeping the cache sent 84: the 21-token prompt once, then one token per step. Both wrote the same 64 tokens. This counts tokens, not time. How much time a cache saves depends on the runtime and the hardware. Chapter 9 covers the reply loop with its cache, and chapter 13 covers the cost.

What a block has to keep depends on how far back it looks. In code, that is the difference between a list that keeps growing and a fixed-size window. The script below runs two rules one token at a time from the prompt [3, 1, 4]. Rule A reads every token so far and counts how many are larger than the newest. Rule B reads only the last three and takes the last digit of their sum, which becomes the next token:

everything = []                   # every token so far
window = []                       # only the last three
new = [3, 1, 4]                   # the prompt
for step in range(5):
    for t in new:
        everything.append(t)
        window = (window + [t])[-3:]
    # rule A reads every token: how many beat the newest?
    a = sum(1 for t in everything if t > everything[-1])
    # rule B reads the last three: their sum, last digit
    b = sum(window) % 10
    print(f"step {step} sent {len(new)} | A kept "
          f"{len(everything)} says {a} | B kept {len(window)}"
          f" says {b}")
    new = [b]

$ python3 lookback.py
step 0 sent 3 | A kept 3 says 0 | B kept 3 says 8
step 1 sent 1 | A kept 4 says 0 | B kept 3 says 3
step 2 sent 1 | A kept 5 says 2 | B kept 3 says 5
step 3 sent 1 | A kept 6 says 1 | B kept 3 says 6
step 4 sent 1 | A kept 7 says 1 | B kept 3 says 4

Rule A's kept list grows by one every step, from 3 to 7. Rule B's window stays at 3. After the first call, each call sends one token. What a rule must keep depends on how far back it looks.

The model has both kinds of rule. Attention is a mixer that lets each token pull information from every earlier token, weighted by content. So an attention block keeps numbers for every token so far. Keys and valuesTerm keys and valuesThe numbers an attention block keeps for every token so far, so each new token can look back at them. are the numbers an attention block keeps for every token so far, so each new token can look back at them. Conv is a mixer that blends each token with the two tokens just before it (short for convolution). So a conv block keeps a fixed window. Its conv stateTerm conv stateWhat a conv block keeps between calls: the numbers it blends, for the last three positions. is what a conv block keeps between calls: the numbers it blends, for the last three positions.

The next script runs the cached loop on chapter 1's question. At each step it prints the tokens it sent, the token it picked, and what one attention block and one conv block keep:

import torch
from load import tok, model

Q = "What is the capital of France? Answer in one sentence."
msgs = [{"role": "user", "content": Q}]
ids = tok.apply_chat_template(
    msgs, add_generation_prompt=True,
    return_tensors="pt", tokenize=True)["input_ids"]
cache, new = None, ids                   # first call: the prompt
for step in range(40):
    with torch.no_grad():
        out = model(new, past_key_values=cache, use_cache=True)
    cache = out.past_key_values          # kept for the next call
    nxt = int(out.logits[0, -1].argmax())
    keys = cache.layers[2].keys          # attention block 2
    conv = cache.layers[0].conv_states[0]  # conv block 0
    print(f"{step} in {new.shape[1]:>2} {tok.decode([nxt])!r:12}"
          f" keys {tuple(keys.shape)} conv {tuple(conv.shape)}")
    if nxt == 7:
        break
    new = torch.tensor([[nxt]])          # next call: one token

$ python3 cache.py
0 in 21 'The'        keys (1, 8, 21, 64) conv (1, 1024, 3)
1 in  1 ' capital'   keys (1, 8, 22, 64) conv (1, 1024, 3)
2 in  1 ' of'        keys (1, 8, 23, 64) conv (1, 1024, 3)
3 in  1 ' France'    keys (1, 8, 24, 64) conv (1, 1024, 3)
4 in  1 ' is'        keys (1, 8, 25, 64) conv (1, 1024, 3)
5 in  1 ' Paris'     keys (1, 8, 26, 64) conv (1, 1024, 3)
6 in  1 '.'          keys (1, 8, 27, 64) conv (1, 1024, 3)
7 in  1 '<|im_end|>' keys (1, 8, 28, 64) conv (1, 1024, 3)

The reply is the same eight tokens as chapter 1's loop: The capital of France is Paris. and the end-of-turn token. The difference is the in column. The first call sends the 21-token prompt, and every later call sends one token.

Block 2 is an attention block. Its keys are 1 prompt, 8 groups, 21 tokens so far, and 64 numbers per group. So for every token so far it keeps 512 numbers of keys, and as many of values. The token count grows by one each step, from 21 to 28. Block 0 is a conv block. Its conv state is 1 prompt, 1,024 numbers, and 3 positions. It never grows. Each step drops the oldest position and adds the newest.

These two arrays are a reason to say what each size counts instead of assuming where it sits. In the keys, the token level comes before the 64 numbers of each group. In the conv state, the 1,024 numbers come first and the 3 positions last.

Across the whole model, each of the six attention blocks keeps 2 × 8 × 64 = 1,024 numbers per token, 6,144 per token in all. The eight conv blocks keep 24,576 numbers in total, however long the reply runs. We calculated that at 10,000 tokens the cache would hold about 61 million numbers, about a quarter as many as the parameters.

The cache stands in for every earlier token. The model never sees those tokens again, only what the cache kept from them. So whatever is in the cache, the model takes as the conversation so far. In code, the danger is a cache stored under the wrong key. The server below keeps earlier words between requests. It keys them first by the prompt's word count, then by conversation:

capitals = {"France": "Paris", "Spain": "Madrid"}

def model(words):                 # reads every word it is given
    for w in words:
        if w in capitals:
            return capitals[w]

caches = {}                       # kept between requests
def serve(key, words):
    earlier = caches.setdefault(key, words[:-1])
    return model(earlier + words[-1:])   # cache + last word

fr = "The capital of France is".split()
es = "The capital of Spain is".split()
print("keyed by length:      ",
      serve(len(fr), fr), serve(len(es), es))
caches.clear()
print("keyed by conversation:",
      serve("chat 1", fr), serve("chat 2", es))

$ python3 session.py
keyed by length:       Paris Paris
keyed by conversation: Paris Madrid

Both prompts are 5 words long. Keyed by length, the Spain request finds France's stored words and answers Paris. Keyed by conversation, each request gets its own stored words, and the answers are Paris and Madrid. The function believes whatever words it is handed.

The real model does the same:

import torch
from load import tok, model

def ids(text):
    return tok(text, return_tensors="pt").input_ids
def top(out):
    return tok.decode([int(out.logits[0, -1].argmax())])

fr = ids("The capital of France is")
es = ids("The capital of Spain is")
with torch.no_grad():
    print("Spain, whole prompt:  ", repr(top(model(es))))
    cache = model(fr[:, :-1], use_cache=True).past_key_values
    out = model(es[:, -1:], past_key_values=cache)
    print("Spain, France's cache:", repr(top(out)))

$ python3 stale.py
Spain, whole prompt:   ' Madrid'
Spain, France's cache: ' Paris'

The two prompts are both six tokens long, counting the start token, and they differ only at position 4, France against Spain. The script builds a cache from the first five tokens of the France prompt. Then it sends the last token of the Spain prompt, is, with that cache. The model answers Paris, with exactly the probabilities it gives the France prompt, 95.22 percent for Paris. As far as the model can tell, it was asked about France.

A runtime that reuses a cache because the earlier tokens have the same count, without checking that they are the same tokens, makes this bug. So does one that hands one conversation's cache to another. The parameters are fine, and the model did exactly its job.

A cache is part of the input. Start each conversation with an empty cache, and reuse one only for the exact tokens it was built from.

Three questions place any array a running model touches:

1. What do its sizes count? Say it before you index the array.
2. Who made it? Training made the parameters. One call made the activations. The loop made the cache from activations. The publisher wrote the config.
3. How long does it live? As long as the model is loaded, for one call, or for one conversation.

Two kinds of array sit outside the table. The token IDs are the input the tokenizer hands the model. And when the model loads, it computes two small lists of 32 numbers from the config, which it uses to mark positions. They are not parameters, and training did not set them.

The questions also find this chapter's bugs. The wrong slice failed the first question, because its first size counted positions where it should have counted prompts. The stale cache failed the third, because it was used outside the conversation it lives for. The double count needs one more check for learned arrays: are two names the same array?