LM2-T6 — Batching, Data Loaders and Data Splits

Learning objectives

After this lesson, you should be able to:

• Describe the shape of an input and a target batch and compute tokens per step.
• Compute the number of steps per epoch for a given dataset size and batch size, with and without drop_last.
• Explain why training data is shuffled and why a validation set must be split off before windowing.
• Use the PyTorch DataLoader to batch a window dataset and check the result with assertions.

The idea: many windows at once, in an order that does not teach the wrong thing

A batch is a stack of B windows. The input tensor has shape (B, L) with integer entries (token IDs), and the target tensor has the same shape. The model produces logits of shape (B, L, V) from them (Module 4). Batching matters for two reasons: hardware does much more work per second when many sequences are processed together, and the average gradient over a batch is a less noisy estimate than the gradient from a single window.

Shuffling

If windows are fed in the order they occur, consecutive batches contain neighbouring text: a whole stretch about one topic, then a whole stretch about another. The gradient then drifts with the topic, and training oscillates. Shuffling the windows each epoch makes every batch a random mixture, closer to the assumption that examples are independent draws from the data. Shuffle only the training data. Validation data is read in a fixed order, because you want the same number every time you measure.

The last, smaller batch

If the number of windows is not a multiple of B, the final batch is smaller. Its loss is averaged over fewer examples and is noisier. The drop_last option discards it. For large training runs this is the usual choice; for small datasets and for evaluation you may keep every window, remembering that the final batch carries less weight if you average batch losses naively.

The validation split

Training loss measures how well the model fits text it has seen. To know whether it will do well on new text, hold some text back. A validation set is data the model never trains on, used to monitor progress and pick settings. A test set is touched rarely, to report the final result honestly. For pretraining on one text, the book splits the token stream at a fixed fraction, such as the first 90% for training and the rest for validation.

Contiguous splitting also has a weakness: if the end of the text differs from the beginning (the last chapter of a book, a different year), the validation text is not typical. For corpora of many documents, assign whole documents to a set at random instead.

The mathematics: steps, tokens per step and the cost of an epoch

Let the dataset have n windows (Lesson LM2-T5), batch size B and context length L. The number of optimizer steps in one epoch is

n

number of input/target windows in the training set.

B

batch size: windows per optimizer step.

L

context length: tokens per window.

floor, ceil

round down or up. The difference is one smaller final batch.

Each step processes

and an epoch processes about n · L input tokens. With non-overlapping windows that is close to the stream length T. For instance n = 20, B = 8: with drop_last there are floor(20/8) = 2 steps per epoch and 4 windows are skipped each time (a different 4 after shuffling); without it, ceil(20/8) = 3 steps, the last with 4 windows.

Batch size and gradient noise. The gradient is an average of B window gradients; its noise (standard deviation) shrinks roughly as 1/√B. A batch of 64 has about one third of the noise of a batch of 8 (8/64 under the square root is 1/8, so √(1/8) ≈ 0.35). The rate of learning per step does not rise accordingly, and memory grows linearly with B · L. Practical batch sizes are set by memory: if a batch does not fit, use gradient accumulation (sum gradients over several small batches before one update).

The split. With a stream of T tokens and training fraction r (such as 0.9), the training part has ⌊rT⌋ tokens, validation has T − ⌊rT⌋. Apply the window formula to each part separately: a part must have more than L tokens to give even one window. For T = 5,145 and r = 0.9, training has 4,630 tokens and validation 515; with L = 256 and s = 256 these give floor((4630 − 257)/256) + 1 = 18 and floor((515 − 257)/256) + 1 = 2 windows. Two validation windows is a noisy estimate; a small corpus often needs a smaller L or a larger validation fraction.

A worked example: from windows to batches

A training stream of 100 token IDs, L = 4, s = 4, B = 8. Windows: floor((100 − 4 − 1) / 4) + 1 = 24. Batches with drop_last=True: 24 / 8 = 3 exactly. Each input batch has shape (8, 4) and each target batch has shape (8, 4); tokens per step = 32. With B = 10, we get 24 / 10 = 2.4: two full batches and one of 4 windows if drop_last=False, or two batches if True.

import torch
from torch.utils.data import Dataset, DataLoader

class Windows(Dataset):
    def __init__(self, ids, L, stride):
        self.starts = list(range(0, len(ids) - L, stride))
        self.ids, self.L = torch.tensor(ids), L
    def __len__(self):
        return len(self.starts)
    def __getitem__(self, k):
        i = self.starts[k]
        return self.ids[i : i + self.L], self.ids[i + 1 : i + self.L + 1]

stream = list(range(100))
cut = int(0.9 * len(stream))                  # split the stream first ...
train, val = stream[:cut], stream[cut:]       # ... then window each part
train_ds, val_ds = Windows(train, 4, 4), Windows(val, 4, 4)
print(len(train_ds), len(val_ds))             # 22 2

loader = DataLoader(train_ds, batch_size=8, shuffle=True, drop_last=True,
                    generator=torch.Generator().manual_seed(0))
x, y = next(iter(loader))
print(tuple(x.shape), tuple(y.shape))         # (8, 4) (8, 4)
print(bool((x[:, 1:] == y[:, :-1]).all()))    # True: targets are inputs shifted by one
print(len(loader))                            # 2  (22 windows // 8, last 6 dropped)

Three checks in the code are worth copying into your own work: the shape of a batch, the shift identity (x[:, 1:] == y[:, :-1]), and the count of batches per epoch. Fixing the generator seed makes the shuffle reproducible, which matters when you compare two runs. The validation part here (10 tokens) gives only two windows, a reminder of how small a naive 90/10 split can be.

Engineering implications and limitations

Split first, then window. Create the training and validation token streams, then build separate datasets. Never cut windows and then divide them randomly: neighbouring windows share tokens whenever the stride is smaller than the context length.

Shuffle training, not validation. Evaluate in a fixed order with shuffle=False so the validation loss is comparable between runs.

Seed everything that matters. The shuffle, model initialization and dropout all use random numbers. Set seeds and pass a torch.Generator to the DataLoader so that two runs differ only in what you changed.

Mind the tail and the weighting. If you average per-batch losses and the last batch is smaller, it counts too much. Either use drop_last or weight by batch size.

Know your throughput. The DataLoader can be a bottleneck. For in-memory tensors, num_workers=0 is fine; for large datasets read from disk, use workers and memory-mapped token files. Measure before changing.

Size the validation set sensibly. A validation set with a handful of windows gives a noisy number, so apparent improvements may be chance. Use enough tokens to make small differences in loss visible, and report the number of validation tokens.

Do not tune on the test set. If you try ten settings and keep the one with the best test score, the test score has become a validation score. Keep a final set untouched.

Memory arithmetic. Token IDs are integers: 8 bytes each in PyTorch’s default int64. A batch of 8 × 256 IDs is 16 KB, negligible. The activations inside the model are what fill GPU memory, and they scale with B · L (and with L² for attention scores).

Common misconceptions

Engineering takeaways

1. A batch is a stack of windows: inputs and targets both have shape (B, L); tokens per step are B · L.
2. Steps per epoch: floor(n/B) with drop_last, ceil(n/B) without; for n = 24, B = 8 that is 3.
3. Shuffle training windows each epoch, never validation; seed the shuffle for reproducible comparisons.
4. Split the token stream first and cut windows separately in each part so no window crosses the boundary.
5. Check batch shapes, the shift identity and the number of batches before a long run.

References and further reading

See Chapter 2 of the book for the data loader and Chapter 5 for the 90 percent training split used with a real text; the PyTorch documentation describes DataLoader.

• PyTorch documentation: the Dataset and DataLoader classes (torch.utils.data).
• Build a Large Language Model (From Scratch), Chapter 2: Working with Text Data (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

That completes Module 2. Next, Module 3 turns token IDs into vectors and builds the mathematics of attention.

Related content

Nothing below is required to finish this lesson.

• Before
  LM2-T5: Input-Target Pairs and Sliding Windows

  Build next-token training examples by shifting a token sequence, with context length and stride.

• Lab
  LM2-L3: Diagnose Sliding Windows and Batching

  Build input-target pairs and find the off-by-one, leakage and batching bugs in a data pipeline.