← Sprout Your own LLM from scratch Glossary Code RU

Chapter 9 of 14 30 min

The corpus

A model is what it has read. Here is where Sprout's 1.4 billion tokens came from, how we cleaned them, and how much text a model of this size actually needs.

In this chapter

  • browse the three sources of Sprout's corpus and read exactly what it read
  • clean and deduplicate text the way prepare.py does, rule by rule
  • estimate how many tokens a model of a given size needs, and why 330 million suits Sprout

In the last chapter we assembled the whole machine: eight transformer blocks, 17.31 million parameters, every wire in place. Switch it on and it produces nonsense. That is not a bug. The architecture is only a shape; what fills it is text. Everything Sprout will ever say, every name it knows, every joke and every grammatical habit, will come from the pile of text we are about to build.

People who train models like to say that the data matters more than the architecture. At our scale that is simply true: the same eight blocks trained on web pages would babble about cookies and privacy policies; trained on children's stories, they learn to tell children's stories. So before we train anything, let's choose what Sprout reads, look through its books, and clean them carefully.

Sprout's library

Sprout's corpus is made of three open datasets, 6 million documents in total. We chose them for one quality: they are written in simple, clean, consistent English, which a small model can actually master. Open the shelves and read a few pages. The switch above the page shows the same document as Sprout sees it: first as tokens, then as the bare numbers that go into training.

Real documents from the corpus, picked at random from each source. "From the mix" chooses a source in proportion to its share of tokens, just as training does.

TinyStories is the heart of the collection. These are short fairy tales that GPT-4 was asked to write using only words a three- or four-year-old understands. They have a recognisable rhythm: "Once upon a time… One day… The end". Their vocabulary is tiny, their plots are simple, and every story is complete. We use version 2, which contains only the stories written by GPT-4.

SimpleStories is synthetic too, but its stories were generated from a grid of parameters: topic, theme, style, narrative device, and so on. The result is still simple English, but longer and more varied: time travel, a grandmother's box of old photos, a story told from the point of view of a lighthouse. They make Sprout's world a little wider.

SODA is a different animal: everyday dialogues between two people, each preceded by a short narrative that sets the scene. We turned every dialogue into lines like Anna: I'm fine, thanks. Without them Sprout would know how stories go, but not how people talk to each other, and in chapter 12 we will teach it to hold a conversation.

Why children's stories?

It may seem strange to feed a model fairy tales when the internet is full of text. The reason is capacity. A model with 17 million parameters is too small to swallow the whole world. Trained on web pages, it spreads its few numbers thinly across programming, football, recipes and advertising, and learns each of them badly. The words it produces look English, but the sentences don't hold together.

The authors of TinyStories asked the opposite question: what if we shrink the world instead of growing the model? Take a language a small child can understand, with about the vocabulary of a four-year-old, and see how small a model can be and still write coherent text. Very small, as it turned out. Models with only a few million parameters learned to write fluent, grammatical stories in which the characters stay consistent from beginning to end.

There is a price to pay. Synthetic text carries the imprint of the model that wrote it: GPT-4's favourite turns of phrase ("filled with joy", "learned a valuable lesson") are repeated in hundreds of thousands of stories, and Sprout will pick them up as faithfully as grammar. At the end of the chapter we will find the fingerprints of these habits among the most frequent tokens of the corpus.

A model is a compressed portrait of its corpus. It cannot know what it has not read, and it picks up every habit of its text, good and bad.

The mixture

After cleaning and tokenization the corpus contains 1.43 billion tokens: 541.5 million from TinyStories, 607.7 million from SimpleStories and 285.2 million from SODA. Proportions are easier to feel than to read, so in the mosaic below each square is one million tokens, a thick novel's worth of text.

Each coloured square is about a million training tokens; the grey squares below are the validation part. "Train" runs all of Sprout's training steps with random windows of the real lengths (128, 256, then 512 tokens), and each square fills up in proportion to how much of its text was actually read.

Two things are worth noticing. First, switch to "As in train.bin". Before writing the file, prepare.py cuts every source into runs of about a million tokens (always on a document boundary) and shuffles those runs. As a result any stretch of the file is a fair mix of all three sources. Training picks windows at random anyway, but a shuffled file is safer: if you ever read it in order, or take just its beginning for a quick experiment, you still get the right mix.

Second, press "Train". Sprout's run takes 10,070 steps of 32,768 tokens each, 330 million tokens in total. That is only 23% of the file. By the end of training most squares are barely tinted: about four fifths of the corpus is never seen at all, and only about 2% of it is seen twice or more. Sprout never reads its library cover to cover, and it does not need to: in the section on the token budget we will see that 330 million tokens is exactly the right dose for its size.

This has a pleasant consequence. A model that almost never sees the same text twice cannot simply memorise its training set: the only way to lower the loss is to learn things that are true of text in general. In the next chapter this will show up in the curves: the loss on the training batches and the loss on validation text will go almost hand in hand.

Cleaning

Raw text always carries rubbish, even when it was written by a model. Typographic quotes “like these” and plain quotes "like these" mean the same thing, but for the tokenizer they are different bytes and therefore different tokens. A model would have to learn twice that a quote opens a line of dialogue. The same goes for three kinds of dash, the ellipsis character, the non-breaking space and Windows line endings. Then there are stray spaces, piles of blank lines, fragments too short to be a story, and texts in other languages.

Our cleaner is one small function. Try it on the examples, or paste your own text:

The same rules as clean() in prepare.py, run in the browser. The deduplication memory starts with the documents from the library above: the "duplicate" example is one of them, in capital letters.

Here is that function exactly as it appears in prepare.py. It is pure Python, so you can run it right here:

REPLACE = str.maketrans({'‘': "'", '’': "'", '“': '"', '”': '"', '–': '-', '—': ' - ', '…': '...', ' ': ' ', '\r': ''}) def clean(text): """Normalise quotes and dashes, drop stray spaces; None if the text should go.""" text = text.translate(REPLACE).strip() lines = [' '.join(line.split()) for line in text.split('\n')] text = '\n'.join(lines).strip() while '\n\n\n' in text: text = text.replace('\n\n\n', '\n\n') if len(text) < 80: return None if sum(ord(c) > 126 for c in text) > 0.002 * len(text): # other scripts, mojibake return None return text

The last rule deserves a closer look. It counts characters beyond the printable ASCII range and throws out the text if they make up more than 0.2% of it. For a story of 500 characters that means one foreign character is allowed, two are not. The rule is crude: it removes Russian, Chinese, emoji and "mojibake" (text that was decoded with the wrong encoding and turned into "café"), but it also throws out an innocent story that mentions a "café" once too often. For an English-only model of this size that is a fair trade: every such character costs two or three tokens and teaches Sprout nothing it could use.

On our data the cleaner had almost nothing to do, which is typical of synthetic text. Here is exactly what it did (we counted the reasons separately):

  • TinyStories: 2,717,700 → 2,717,373 documents. 232 were too short and 95 had too many characters outside ASCII: mostly the words café, piñata and Chloé, plus a few emoji, invisible zero-width spaces and one English story with a few sentences in Chinese. Quotes or spaces had to be fixed in about 338 thousand stories, one in eight.
  • SimpleStories: 2,115,696 → 2,115,670. Only 26 exact duplicates were removed, but spaces or line breaks had to be fixed in almost 1.28 million stories, 60% of them.
  • SODA: 1,186,423 → 1,184,876 (another 5,159 dialogues, whose lists of speakers and lines did not match, were skipped while reading). 1,545 dialogues went because of foreign characters, mostly conversations in which the characters switch to Spanish or French, or mention déjà vu. Two were too short.

For text collected from the web it is the other way round: most of a raw web crawl is thrown away at this stage, and cleaning is the most laborious part of the whole project.

Duplicates

The last step of cleaning is removing repeats. If the same story appears in the corpus a thousand times, the model sees it a thousand times more often than any other: it will memorise it word for word, and its idea of what an "ordinary" story looks like will be skewed towards it. Duplicates also simply waste the training budget.

Exact duplicates are cheap to find. For each document we compute a hash, a short fingerprint of 16 bytes (MD5), and keep a set of the fingerprints we have already seen. We hash the text in lower case, so that a story shouted in capitals counts as the same story. Checking whether a fingerprint is in a set takes the same time whether the set holds ten documents or ten million.

import hashlib messy = "“Can we go to the park?” asked Lily.\u00a0\u00a0Mom smiled — “Yes… but first, lunch!”\r\n\r\n\r\nThey ate quickly and ran outside." print(repr(clean(messy))) print(clean("The end.")) print(clean("Жили-были дед да баба, и была у них курочка Ряба. Снесла курочка яичко, да не простое.")) seen = set() for doc in ["Tom had a red ball. He liked to play with it in the park every day with his dog.", "TOM HAD A RED BALL. HE LIKED TO PLAY WITH IT IN THE PARK EVERY DAY WITH HIS DOG.", "Tom had a blue ball. He liked to play with it in the park every day with his dog."]: h = hashlib.md5(clean(doc).lower().encode()).digest() print('duplicate' if h in seen else 'new ', h.hex()[:12], doc[:28]) seen.add(h)

Notice the third document. It differs from the first by a single word, and the hash is completely different: exact deduplication does not catch near-duplicates. For our synthetic sources this hardly matters, but on the web near-copies are everywhere: the same article on a hundred sites with a different menu around it.

Documents and windows

How long is a typical document? Sprout reads text through a window of 512 tokens: that is its context. If documents were longer, the model would never see their beginnings and endings together; if they were much shorter, one window would contain several unrelated stories.

Distribution of document lengths in tokens. The dashed lines are the three window lengths Sprout trains with: 128, 256 and 512 tokens. The slider shows what share of documents fits whole into a window.

On the full validation set (30,088 documents) the average TinyStories story is 199 tokens long, a SimpleStories story 286 and a SODA dialogue 240. Almost all documents, 97%, fit into 512 tokens. That means a random 512-token window usually covers the end of one document, a whole second one and the beginning of a third.

How does the model know where one story ends and another begins? Every document in the file starts with a special token, <|endoftext|>. The model quickly learns what it means: whatever came before it has nothing to do with what follows, because a new text is starting. The same token lets us generate "from scratch": we give the model only <|endoftext|> and it starts a new story.

The chart also explains a trick from the next chapter. For the first 10% of training Sprout reads windows of only 128 tokens, then 256, and only then 512. Short windows are cheaper, because the cost of attention grows as the square of the window length, and at the start of training the model is learning local things anyway: words, punctuation, grammar. Look at the slider: few documents fit into 128 tokens whole, but that is plenty for a sentence.

From text to train.bin

Training must not waste time on text processing: the GPU would stand idle waiting for the tokenizer. So we tokenize the whole corpus once, in advance, and store the result as a flat array of numbers. Here is the part of prepare.py that does it:

# 3. encode in parallel; a small slice of every source goes to validation tok = Tokenizer.load(tok_path) eot = tok['<|endoftext|>'] train, val = [], [] with Pool(max(1, os.cpu_count() - 2), initializer=_init, initargs=(tok_path,)) as pool: for name, texts in docs.items(): n_val = int(len(texts) * args.val_frac) for split, part in ((val, texts[:n_val]), (train, texts[n_val:])): chunks = [part[i:i + 2000] for i in range(0, len(part), 2000)] split.append(np.concatenate(pool.map(_encode, chunks))) stats[name]['tokens'] = int(len(train[-1]) + len(val[-1])) print(f'{name}: {stats[name]["tokens"] / 1e6:.1f}M tokens', flush=True) # shuffle whole runs of documents between sources, so any stretch of the file is a fair mix blocks = [] for arr in train: starts = np.flatnonzero(arr == eot) cuts = starts[np.searchsorted(starts, np.arange(1_000_000, len(arr), 1_000_000))] blocks.extend(np.split(arr, np.unique(cuts))) rng.shuffle(blocks) np.concatenate(blocks).tofile(f'{args.out}/train.bin') np.concatenate(val).tofile(f'{args.out}/val.bin')

A few details are worth a comment:

  • In parallel. Documents are cut into chunks of 2,000 and handed to a pool of processes, roughly one per processor core. Each worker loads its own copy of the tokenizer and returns an array of numbers. On our Mac all 1.43 billion tokens were encoded in about a minute, and the whole of prepare.py, from raw files to train.bin, ran in about five.
  • Validation. The first 0.5% of documents of each source (the list was shuffled earlier) goes to val.bin: 7.17 million tokens that Sprout never trains on. That is our honest exam: the loss on this text shows how the model copes with text it has never seen.
  • uint16. Each token is stored in two bytes. Two bytes hold numbers from 0 to $2^{16} - 1 = 65{,}535$, and our largest id is 8,191, so every id fits comfortably. 1.43 billion tokens take 2.85 GB, half of what ordinary four-byte integers would need.

Let's see what that looks like at the byte level and how training reads such a file. Here is the beginning of a story as ids, and random windows cut from it the way the Batches class in train.py cuts them:

import numpy as np # "Once upon a time, there was a little girl." after <|endoftext|> (8188) ids = [8188, 679, 690, 258, 466, 44, 490, 307, 258, 531, 548, 46] arr = np.array(ids, dtype=np.uint16) print(arr.nbytes, 'bytes for', len(arr), 'tokens') print(arr.tobytes()[:8].hex(' ')) # little-endian: the low byte comes first # train.bin is just this, 1.4 billion times; np.memmap reads it without loading it into memory data = np.frombuffer(arr.tobytes(), dtype=np.uint16) rng = np.random.default_rng(0) T = 4 # Sprout's windows are 128-512 tokens for s in rng.integers(0, len(data) - T - 1, 3): window = data[s:s + T + 1] print('x =', window[:-1].tolist(), ' y =', window[1:].tolist())

The pair x, y is the whole secret of training data for a language model. y is x shifted by one token: at every position the model sees the tokens up to it and must guess the next one. One window of 512 tokens gives 512 such questions at once, and one training step (32,768 tokens) gives 32,768 of them. No labels, no annotators: the text is its own answer key.

Why do we store tokens as uint16 rather than as ordinary 32-bit integers?

Two bytes hold numbers up to 65,535. With a vocabulary of 8,192 that is plenty, so train.bin takes 2.85 GB instead of 5.7. Models with larger vocabularies (over 65,536 tokens) have to use four bytes.

How much text is enough?

We have 1.43 billion tokens, but we train on 330 million. Why not everything? And why not less? To answer that we need a measure of cost.

Now the question can be put precisely: we have a compute budget $C$. How should we split it between the size of the model $N$ and the amount of text $D$? A large model trained on little text is undertrained. A small model trained on a mountain of text hits the ceiling of its capacity, and the extra compute is wasted. In 2022 a team at DeepMind trained over four hundred models of different sizes on different amounts of data and found the optimum: parameters and tokens should grow together, in proportion. The rule of thumb that came out of it is about 20 tokens per parameter.

For Sprout: $17.31 \text{M} \times 20 \approx 346$ million tokens. We train on 330 million, about 19 tokens per parameter: by Chinchilla's standard this is a compute-optimal run. Drag the point in the chart and see what your own model would cost:

Log scale on both axes. The dotted diagonals are lines of equal compute $6ND$; the golden line is the Chinchilla optimum. "On our Mac" is the time at the speed actually measured on Sprout: 29.7 thousand tokens per second on a Mac with an M4 Pro chip, which is about $3 \cdot 10^{12}$ useful operations per second. The real run went somewhat slower than this benchmark; the time-lapse in the next chapter shows its actual clock.

Look at Llama 3 8B: 15 trillion tokens for 8 billion parameters, almost two thousand tokens per parameter, a hundred times more than Chinchilla recommends. That is not a mistake. Chinchilla optimises the cost of training, but a popular model is then run billions of times, and a small model is cheaper to run. So today it pays to train a small model far longer than the "optimum": the loss keeps going down, only more slowly. For Sprout we chose the classic optimum because a training run of a few hours on one Mac was our real constraint.

What Sprout reads most

Let's finish by looking at the corpus through the model's eyes. Here are the most frequent tokens of the validation text and how much of all the text they cover.

Counts over the 7.17 million tokens of the validation text. The top 400 tokens are listed; the other 7,792 share the remaining quarter of the text.

Ten tokens (the full stop, the comma, the line break and short words like " the", " and", " a" and " to") make up more than a quarter of all text. A hundred make up more than half. That is Zipf's law, which we met in chapter 6, and it matters for training: after a few hundred steps the model already guesses these tokens well, and all the rest of training is a slow fight for the rare, meaningful ones.

And look who lives in this world. " Lily" and " Tim" are more frequent than most verbs, " happy" is more common than " sad", " joy" and " together" are among the top two hundred. These are the fingerprints of the models that wrote the stories. Sprout will inherit them: it will be kind, a little sentimental and fond of the name Lily. Now you know where that comes from.

Sprout right now

This is a real Sprout at full size: eight blocks, six heads, 17.31 million parameters. But its weights are random numbers, and it has not read a single token of the corpus. Let it write and you get a random jumble of tokens. Ask how surprised it is by a real story, and the answer will be close to $\ln 8192 \approx 9.01$ nats per token: exactly what a model that spreads its probability evenly over all 8,192 tokens would say. In the next chapter we start reading the library, and in about 1,500 steps this number will fall below two (by the last step, to 1.55).

Chapters

  1. 0 Meet Sprout
  2. 1 Counting letters
  3. 2 Measuring surprise
  4. 3 Gradient descent
  5. 4 Backpropagation
  6. 5 Embeddings
  7. 6 Tokens
  8. 7 Attention
  9. 8 Transformer
  10. 9 Corpus
    1. Sprout's library
    2. Why children's stories?
    3. The mixture
    4. Cleaning
    5. Documents and windows
    6. From text to train.bin
    7. How much text is enough?
    8. What Sprout reads most
    9. Sprout right now
  11. 10 Training
  12. 11 Sampling
  13. 12 Chat
  14. 13 LoRA
  15. 14 What's next