← Sprout Your own LLM from scratch Glossary Code RU

Chapter 9 of 17 30 min

The corpus

Sprout needs a library. We will look at its sources, prepare the texts for training, and see why the size of a corpus is not the amount a model actually reads.

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 previous chapter we assembled a transformer with eight blocks and 17.31 million parameters. Its weights are still random, so it produces nonsense: a well-designed architecture does not know English by itself. To make those numbers capture patterns in language, we need examples. The texts we choose determine which words Sprout encounters often, which kinds of story it practises continuing, and which subjects it barely encounters at all.

Imagine giving two identical models different libraries: children's stories for one, recipes, forum threads and technical manuals for the other. They now face different learning tasks. Stories repeat words, characters and sentence patterns, giving a small model a manageable place to practise writing coherent text. Before training, we will choose Sprout's library, read a few pages and prepare the documents for it.

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 adds conversation to the library. These are synthetic everyday dialogues generated with InstructGPT from descriptions of social situations. Each has a short introduction; we turn the dialogue into lines such as Anna: I'm fine, thanks. These examples expose Sprout to speakers taking turns and responding to each other. That experience will be useful when we fine-tune it for chat in chapter 12.

Why children's stories?

Why stories, when the internet offers so much text? A model with 17 million parameters has limited capacity. A mixed web corpus asks it to handle code, sports reports, recipes and advertisements at once. Children's stories give it a more manageable task: a small vocabulary, recurring constructions and short plots whose beginnings and endings can fit in the same context.

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 corpus supplies the examples a model learns from. The model can combine familiar elements in new ways and apply patterns to unseen cases. But it also picks up repetitions, errors and imbalances in its sources, alongside the useful skills.

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 estimates coverage of each square. This is a sampling simulation, not a trace of positions read during the actual run.

Switch to "As in train.bin". Before writing the file, prepare.py cuts each source into runs of about a million tokens, preserving document boundaries, then shuffles the runs. Sources are spread throughout the file, though any short stretch need not contain the intended proportions: a single run still belongs to one source. Training samples windows at random across the whole file, so the expected contribution of each source follows its share of the tokens.

Now press "Train". Our run consists of 10,070 steps of 32,768 tokens, about 330 million tokens in total. That is 23% of the training file's length, but it does not mean 23% of its positions are read: random windows can overlap. An approximate calculation puts the unread share at 79%, with about 2% read at least twice. Later in this chapter we will see why we chose this budget for the first run.

A corpus much larger than the amount sampled during a run reduces repeated exposure to the same passages. It does not rule out memorisation: a model can learn some sequences from a single presentation, and similar phrases occur in different documents. That is why we still need a separate check. In the next chapter we will compare loss on training batches with loss on held-out text.

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 counts characters with a code point above 126 and rejects a document if their share exceeds 0.2%. In a 500-character text, one is allowed and two are not. This is a crude filter, not a language detector: it removes many Russian or Chinese passages, but can also reject ordinary English containing names such as Chloé or words such as café. For this experiment we chose an almost entirely ASCII corpus. That choice says nothing about the usefulness of other languages or characters.

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.

For each cleaned document, prepare.py converts the text to lower case and computes an MD5 hash: a 16-byte fingerprint. It stores previously seen fingerprints in a set, whose membership checks take approximately constant time on average. There is no need to compare a new document with every earlier one. This script keeps a separate set for each source, so it does not remove duplicates across sources. Hash collisions are also possible: matching short fingerprints are not a mathematical proof that two documents are identical.

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)

The third document differs from the first by just one word, but its hash is different. Exact deduplication therefore misses near-duplicates. These occur in synthetic stories as well as on the web, where one article may appear on a hundred sites with different menus around it. Finding them needs a separate method; Sprout's preparation pipeline does not include that step.

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.

Every document in the file starts with <|endoftext|>. This special token acts as a separator: training supplies many examples of a new text beginning after it. The token does not itself erase the context or stop attention from looking at the preceding document. It provides a learned boundary signal. To generate from scratch, we give the model just this token and ask it to continue.

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 excluded from language-model weight updates. The tokenizer is built before this split and may have seen some of these documents. This is therefore a held-out check of the language model, not a fully isolated evaluation of the entire processing pipeline.
  • 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. With more than 65,536 vocabulary entries, uint16 is insufficient; 32-bit integers are one alternative.

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.

How should we divide a compute budget $C$ between the number of parameters $N$ and the number of training tokens $D$? An oversized model may not get enough training. A small model trained for a very long time may improve so slowly that making it larger would have been a better use of the budget. DeepMind's 2022 experiments found that, near the best balance, parameters and tokens should grow roughly together. This gave rise to the familiar rule of thumb: about 20 tokens per parameter. It is an empirical guide under particular conditions, not a universal quota.

For Sprout, that guide gives $17.31 \text{M} \times 20 \approx 346$ million tokens. Our 330 million amounts to about 19 tokens per parameter. The numbers are close, but that does not establish an optimum: Sprout is smaller than the models in the original experiments, and its corpus is different. We used the rule as a starting estimate that fitted a budget of a few hours on one Mac. Drag the point in the chart to compare other choices:

Log scale on both axes. The dotted diagonals are lines of equal compute $6ND$; the golden line is the 20-tokens-per-parameter guide. "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.

For comparison, Llama 3 8B was trained on over 15 trillion tokens: well over a thousand tokens per parameter. Why go so far beyond twenty? The Chinchilla rule balances the cost of training. For a model that will be used many times, the cost of each answer matters too. Training a smaller model for longer can then make sense: spend more once to make subsequent use cheaper. Sprout has a different aim: to show a complete training cycle within one short experiment.

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, including the full stop, comma, line break and short words such as " the", " and", " a" and " to", account for more than a quarter of the text; the first hundred cover over half. This resembles the uneven frequency pattern we discussed alongside Zipf's law in chapter 6. Common tokens provide many training examples, but predicting them is not always easy: choosing between " a" and " the" also depends on context. Later training can improve both rare continuations and difficult uses of familiar words.

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 More data
  16. 15 A bigger model
  17. 16 Thinking and Python
  18. 17 What's next