Data Loaders and Batching

Before a model can be trained, its data has to exist in a form the training code can use. This article picks up at that step. The goal is to prepare a training dataset and a validation dataset, and then to compute the cross entropy loss on both. Those two losses are the central quantities of model training, so the data pipeline has to be correct before any loss value means anything.

The running text is 'The Verdict', a short story by Edith Wharton. It was already used in chapter 2. Choosing it solves two problems at once. Because the story is in the public domain, there is no question about usage rights. Because it is small, the code runs on a standard laptop in a matter of minutes, even without a high-end GPU.

To see how far this example is from production work, consider the 7 billion parameter Llama 2 model, which the source describes as a relatively popular openly available LLM. Training it required 184,320 GPU hours on A100 GPUs and processed 2 trillion tokens. Dataset choice is therefore a matter of goal and constraints. A small public domain text answers the questions of licensing and speed, which matters when the aim is to learn mechanics. A larger corpus is the route toward more extensive training, at a much higher computational cost.

The first code step reads the whole story into a single string. The file is opened in read mode with UTF-8 encoding, and file.read() returns its full contents. The file name below is a placeholder used for illustration.

Once the text is in memory, the code measures it in two different units. The number of characters comes from len(text_data). The number of tokens comes from passing the same text through the tokenizer and measuring the result. Both measurements are compared directly, because they describe the same span of text in different units.

Next the text is divided into a training part and a validation part. The validation part is used for model evaluation during training, so it must be kept separate from the training part. The split is set by a train_ratio of 0.90, which means 90 percent of the characters go to training and the remaining 10 percent go to validation. The division point is called split_idx, and it is computed from train_ratio. Roughly, if N is the total number of characters, the split index sits near $\approx 0.90 \times N$, with training covering the characters before it and validation covering the characters from it onward.

A split subset is still plain text, and the model needs fixed-size tensors. The source describes the preparation as a chain of operations: the subset is tokenized, the token sequence is divided into chunks of a user-specified length, the chunks are shuffled, and the shuffled chunks are organized into batches. The batch size sets how many chunks are processed together. Follow one training subset through each stage in the diagram below.

Notice that the shape of the data changes at each stage. Text is a string, the token sequence is a flat list of token IDs, a chunk is a fixed-length window of that list, and a batch stacks several chunks together. The final batch shape is (batch_size, context_length), so with a batch size of 2 and 256-token chunks, each batch holds two samples of 256 tokens.

Each loader yields a pair: an input x and a target y. Each chunk therefore serves two roles, and the two windows are always taken from the same token sequence. The explanation below depends on that pairing.

The one-position shift is the usual convention for next-token prediction: the target is the same stream of tokens as the input, moved forward by one position, so the model is asked to predict each next token. The source's loop confirms that each batch is an input and target pair, so the shift is the part to check when you inspect the windows.

The loaders are built with the create_dataloader_v1 function from chapter 2. Both calls use the same batch size, max_length, and stride. The settings that differ are drop_last and shuffle, and those differences encode the training-versus-validation distinction. The code also fixes a random seed so the shuffling is repeatable.

Setting drop_last to True on the training loader discards an incomplete final batch, so every training batch has the full size. The validation loader sets it to False, so no validation chunk is silently thrown away. Shuffling is switched on for training, where the order of examples should vary, and off for validation, where a fixed order keeps evaluation consistent.

The source suggests an optional check. Iterate through each loader and print the shape of every input and target tensor. Each loader yields a pair, so the loop unpacks two values per batch. Try predicting the result before you read the answer.

Batch size is also a resource choice. The source uses a batch size of 2 because the dataset is small. It notes that training LLMs with batch sizes of 1,024 or larger is not uncommon in practice. The batch size used in a small exercise should therefore be read as a classroom setting that keeps runs quick, not as a recommendation for real training.