Lesson 2.4: Context windows, batches and the data loader
Turn a stream of token ids into inputs and targets, cut it into context windows, draw random batches, and see how batch size changes noise, coverage and memory.
By the end of this lesson you will be able to build the training pairs for a language model from a long list of token ids, explain the context window and why the target is the input shifted by one, write a batch loader in a dozen lines, check it for off-by-one errors, and choose a batch size by looking at the noise it produces. You will show that a window of 8 tokens gives 8 predictions at once, that a loss computed through non-overlapping blocks equals the full-batch loss to six digits (3.54219 for both), that the loss of a random batch gets steadier as 1 over the square root of its size (a standard deviation of 0.0313 for 8 tokens and 0.0025 for 1,024), that random batches cover 66 per cent of the training text after one epoch of draws and 99.5 per cent after nearly five, and that training with batches of 256 tokens reaches 1.167 against the full-batch 1.162. You should have read Lessons 1.4 to 2.3. The code is NumPy. This course is education about how language models work, not a guide to any commercial product.
Module 2 began with the question of how to turn text into numbers. We have answered it twice: with tokens (Lessons 2.1 and 2.2), and with vectors for tokens (Lesson 2.3). One step is left before a model can learn: arranging the long stream of ids into training examples, a few hundred at a time, in the shapes a model expects. It is the least glamorous part of the course, and a surprising share of bugs in real training code are found here.
From a stream to examples
After tokenisation, the whole training text is one long array of integers: trainids, 26,547 of them in our Module 1 text. A language model is trained to predict each token from the ones before it. So for every position t, the input is the tokens up to t and the target is the token at t + 1.
In Lesson 1.4 we built the pairs for a bigram model by shifting the array by one:
python
x = trainids[:-1]
y = trainids[1:]
Position t of x is the token at t, and position t of y is the token at t + 1: the next token. That is the rule for every language model in this course, for every context length. For the toy array [0, 1, 2, 3, 4, ...], a window of four gives:
python
toy = np.arange(10)
block = 4
x = toy[0:block] # [0 1 2 3]
y = toy[1:block + 1] # [1 2 3 4]
Read it as four separate predictions, made from one slice of the data:
Position · Context the model may use · Target0 · 0 · 1
1 · 0 1 · 2
2 · 0 1 2 · 3
3 · 0 1 2 3 · 4
This is the key efficiency of language-model training. A window of block tokens provides block training examples at once, with contexts of every length from 1 up to block. A bigram model uses only the last token of each context. From Module 3 on, a model will use the whole of it. The window length is called the block size or context length, and it sets the longest context the model ever sees during training. A model trained on windows of 8 cannot be expected to use 100 tokens of history.
A batch of windows
Training does not use a single window per step. We stack several windows into one array, called a batch, so that the arithmetic runs on a big block of numbers at a time (Lesson 1.2) and so that the gradient averages over many examples. A batch has two shapes to remember: x and y both have shape (batch, block).
To pick the windows, we choose random start positions and slice. This is the whole data loader:
python
def getbatch(data, batch, block, rng):
starts = rng.integers(0, len(data) - block, size=batch)
x = np.stack([data[s : s + block] for s in starts])
y = np.stack([data[s + 1 : s + block + 1] for s in starts])
return x, y
The upper limit len(data) - block is deliberate. The last allowed start, len(data) - block - 1, gives a target slice ending at index len(data) - 1, the final element. A limit one larger would slice past the end and silently return a short row, which np.stack would then refuse; one smaller would never use the last token. These off-by-one errors are the common ones, so we test the loader the moment we write it:
python
rng = np.random.defaultrng(0)
x, y = getbatch(trainids, 4, 8, rng)
print(x.shape, y.shape) # (4, 8) (4, 8)
assert (x[:, 1:] == y[:, :-1]).all() # y is x shifted left by one, in every row
print([vocab[i] for i in x[0]])
print([vocab[i] for i in y[0]])
# ['small', 'boat', 'slowly', '.', 'the', 'cat', 'watched', 'a']
# ['boat', 'slowly', '.', 'the', 'cat', 'watched', 'a', 'small']
The assert states the defining property: every row of y is the same row of x moved one place to the left, with the true next token at the end. Reading one window as words is a second check. Here the window begins in the middle of a sentence, "small boat slowly ." and the next sentence follows. The model's target after "a" is "small", the correct next word in the text.
One more test is worth doing. Our loader draws random windows, but the loss over a full pass through the data can be computed with non-overlapping blocks, and it must agree with the loss from the one-long-array method of Lesson 1.4. For the untrained model of that lesson:
python
blocks = (len(xall)) // 8
xb = xall[: blocks 8].reshape(-1, 8)
yb = yall[: blocks 8].reshape(-1, 8)
print(lossfull, lossblocks) # 3.5421911539 3.5421896213
The two numbers agree to six digits, 3.54219 against 3.54219 (the tiny difference is the few tokens left over at the end that do not fill a block). If they disagreed, the loader would be producing the wrong pairs. A consistency check of this kind is cheap and catches the whole class of shift bugs.
Random windows and what they cover
Drawing start positions at random, with replacement, is simple and good enough, but it does not visit the data evenly. Some stretches are drawn twice before others are drawn once. With our batches of 32 windows of 8 tokens, 256 tokens per step, an epoch (one pass worth of tokens, 26,547) takes 103.7 steps. How much of the text have we actually seen after that many steps?
Steps · Tokens drawn · Epochs of tokens · Share of text seen at least once10 · 2,560 · 0.10 · 10.3%
50 · 12,800 · 0.48 · 41.8%
100 · 25,600 · 0.96 · 66.3%
200 · 51,200 · 1.93 · 88.2%
500 · 128,000 · 4.82 · 99.5%
After "one epoch" of random draws, a third of the text has not been seen and some of the rest has been seen twice. This is normal and harmless for a small model that trains for thousands of steps. For large runs, where the data is much bigger than the training budget, the choice matters in the opposite direction: the model may see each document only once or twice, and care goes into shuffling a whole pass so that nothing repeats before everything has appeared. Both are variations on the same loader.
Why not use all the data every step?
In Lesson 1.4 we used full-batch gradient descent: every step used all 26,546 pairs. That cannot scale. A real training set has billions or trillions of tokens, and one step on all of them is out of the question. The alternative is mini-batch training: compute the gradient on a small random sample, take a step, and draw a new sample. The gradient from a sample is a noisy estimate of the full gradient, but it costs a tiny fraction as much, and the steps add up.
How noisy? We can measure it. Take the untrained bigram model, draw 300 batches of each size, and compute each batch's loss. The full-data loss is 3.5422. The spread of the batch losses is:
Batch (windows × block) · Tokens per batch · Standard deviation of the batch loss1 × 8 · 8 · 0.0313
8 × 8 · 64 · 0.0102
32 × 8 · 256 · 0.0049
128 × 8 · 1,024 · 0.0025
32 × 64 · 2,048 · 0.0019
Two patterns appear. The average batch loss is right: the mean of 300 batches of 256 tokens was 3.5424, against 3.5422 for all the data, so a random batch is an unbiased estimate. And the spread falls as the batch grows, by roughly 1 over the square root of the number of tokens: multiplying the batch by 4 (from 64 to 256 tokens) halves the standard deviation, from 0.0102 to 0.0049. This is the same square-root law as the standard error of an average, applied to the gradient.
Both facts matter for choosing a batch size. A bigger batch gives a more reliable gradient, but a step costs proportionally more, and returns are slow: to halve the noise we must quadruple the work. A small batch gives noisy steps, but many more of them in the same time.
What it does to training
Train the bigram model for 5,000 steps at a learning rate of 10, with three batch sizes, and score each on the validation set:
Step · 4 × 8 (32 tokens) · 32 × 8 (256 tokens) · 128 × 8 (1,024 tokens)100 · 1.362 · 1.329 · 1.324
500 · 1.230 · 1.193 · 1.185
1,000 · 1.206 · 1.174 · 1.171
2,000 · 1.222 · 1.171 · 1.166
5,000 · 1.201 · 1.167 · 1.164
For reference, full-batch gradient descent reached 1.162 in Lesson 1.4.
Look at the first column. With only 32 tokens in a batch, the loss at 2,000 steps (1.222) is worse than at 1,000 (1.206). The model is jittering around the answer instead of settling. The learning rate that suited a smooth gradient is too large for a noisy one. With 256 tokens the curve settles at 1.167, within 0.005 of full-batch, and with 1,024 it is 1.164. A standard remedy for a noisy batch is a smaller learning rate, which Lesson 5.3 develops as a schedule.
The trade is not only about quality. Mini-batches let a model train on data that does not fit in memory, and let the computer spend its time on the most useful thing: more steps on fresh data.
The memory cost of a batch
A batch also decides how much memory a step needs. The output of the model for a batch is a score for every vocabulary entry, at every position of every window: a tensor of shape (batch, block, V). For the bigram model this is small: 32 × 8 × 34 floats of 4 bytes, about 35 kilobytes. For a real model it is not.
Setting · Logits shape · Memory in float32Our bigram: 32 windows of 8, V = 34 · 32 × 8 × 34 · 35 KB
8 windows of 1,024, V = 50,257 · 8 × 1,024 × 50,257 · 1.65 GB
The second row is for the output scores alone, before the gradient and the layers behind them. This is one reason real training uses small batches of long sequences or accumulates the gradient over several small ones. Use the calculator to see how batch settings translate into steps, coverage and memory.
Edges and leaks
Four details separate a loader that works from one that quietly misleads you.
1. Split before windowing. We cut the token array into a training part and a validation part first, and draw windows only inside each part. If we drew windows from the whole array and split afterwards, a window could straddle the boundary and put validation tokens into training inputs. Even a few tokens of leak can make a score look better than it is.
2. Documents, not just a stream. Real data is many documents. Joining them into one long array and sliding windows over it means that some windows cross a document boundary and ask the model to predict the start of an unrelated text. The usual fix is to put an end-of-text token between documents (Lesson 2.1), so the model learns the boundary is a boundary. Our toy text has one kind of sentence, so the full stop does that job.
3. The last block. Never let a window run past the end of the data. The limit in getbatch handles it, and the shape check would catch a mistake.
4. Context at the start of a window. The first token in a window has no history inside it. The model learns to predict from a short context, because that is what the first positions of every window train it for. Evaluation should use the same windows, or the scores are not comparable.
What the loader gives the model
The pipeline for a text is now complete, and short enough to hold in the head:
1. Tokenise the text into ids (Lessons 2.1 and 2.2).
2. Split the ids into a training part and a validation part, in order.
3. Draw random windows of block tokens from the training part and stack them into x; make y the same windows shifted by one.
4. Look up a vector for each id in x (Lesson 2.3), add the position vectors, and feed the result to the model.
5. Compute the loss between the model's scores and y, take the gradient, update the parameters, and draw a new batch.
Module 3 begins at step 4, with the part of the model that decides how each token should use the ones before it.
Key takeaways
The target is the input shifted one place to the left. A window of block tokens gives block predictions, with contexts of every length from 1 to block.
A batch is an array of shape (batch, block) for the inputs and the same for the targets. getbatch picks random starts below len(data) - block and slices.
Test the loader: x[:, 1:] == y[:, :-1] in every row, and a loss through non-overlapping blocks equal to the loss on the whole array (3.54219 for both).
A random batch gives an unbiased estimate of the full loss. Its spread shrinks as 1 over the square root of its size: 0.0313 with 8 tokens, 0.0049 with 256, 0.0025 with 1,024.
Random windows cover 66% of the text after one epoch's worth of draws and 99.5% after nearly five.
Bigger batches reach lower loss in the same number of steps with diminishing returns: 1.206, 1.174 and 1.171 after 1,000 steps for 32, 256 and 1,024 tokens. A noisy batch with a large learning rate jitters (1.222 at step 2,000 against 1.206 at step 1,000).
A batch of 8 windows of 1,024 tokens with a 50,257 vocabulary needs 1.65 GB just for the output scores.
Split before windowing, separate documents with an end-of-text token, and never let a window run off the end.
Common mistakes
Misaligning inputs and targets. If y is the same as x, the model learns to copy. If the shift is two, it learns to skip. Always assert the shift.
Windowing then splitting. The training inputs then contain validation tokens. Split first.
Comparing losses from different batch sizes. A single batch's loss is noisy. Report the loss over a large fixed validation set.
Using a huge batch because it is smoother. The noise falls only as the square root, and each step gets more expensive. The extra work often buys less than more steps.
Forgetting the document boundary. Windows that run across unrelated documents teach the model false continuations.
Counting examples wrongly. A text of N tokens does not contain N windows of every length. Count windows as N minus the block size, and predictions as windows times block.
Check yourself
1. A text has 100,000 tokens and the block size is 50. How many distinct windows can the loader draw, and how many predictions does a batch of 16 windows provide?
Starts run from 0 to 99,949, so 99,950 windows (at most). A batch of 16 gives 16 × 50 = 800 predictions.
2. The batch loss has a standard deviation of 0.02 for a batch of 64 tokens. About what will it be for 1,024 tokens?
1,024 is 16 times larger, so the spread falls by a factor of 4, to about 0.005.
3. For a window [10, 11, 12, 13, 14] of length 4 starting at the first element, what are x and y?
x = [10, 11, 12, 13] and y = [11, 12, 13, 14].
4. Why is it a mistake to draw training windows from the whole token array and split into train and validation afterwards?
A window near the boundary would contain tokens from both sides. Validation tokens would appear as part of the training input, and the validation loss would be too good.
5. A model with a vocabulary of 32,000 processes 16 windows of 2,048 tokens. How much memory do its output scores need in 32-bit numbers?
16 × 2,048 × 32,000 × 4 bytes = 4.19 GB.
Try this
Set aside ninety minutes. You need Python, NumPy and the code from Lessons 1.4 and 2.3.
1. Type in getbatch. Print one batch as words and verify the shift by eye and with the assert.
2. Reproduce the consistency check: compute the loss of the untrained model with the long arrays and with non-overlapping blocks. Make them differ on purpose by shifting y by one extra place, and see how large the gap is.
3. Measure the standard deviation of the batch loss for batch sizes 1, 4, 16, 64 and 256 windows of 8. Plot it against 1 over the square root of the tokens. Is it a straight line?
4. Train the bigram model with batches of 32 × 8 for 5,000 steps at learning rates 1, 10 and 30. Which settles best? Which jitters?
5. Replace the random starts by a sequential pass: shuffle the non-overlapping blocks once, and use each exactly once per epoch. Compare the validation loss after one epoch (about 104 steps) to random sampling.
6. Add an end-of-text token to the word list and put it between every ten sentences. Measure how much the loss on the end-of-text token improves as the model learns where it occurs.
A good result is a loader that passes its tests, a measured relationship between batch size and noise, and a feel for the memory a batch needs. The next module takes the vectors from this module and builds the central mechanism of a transformer: attention.
Read more
Andrej Karpathy's lecture "Let's build GPT: from scratch, in code, spelled out" contains a getbatch function of the same form and discusses the block size and batch dimensions. Léon Bottou's article "Stochastic Gradient Descent Tricks" explains the trade-offs of mini-batches. The paper "An Empirical Model of Large-Batch Training" by McCandlish and colleagues measures how the useful batch size depends on the noise of the gradient. Goodfellow, Bengio and Courville's Deep Learning, section 8.1.3, covers mini-batch algorithms.
Tags: Build an LLM Data loader Context window Mini-batch Training data