How to Train an LLM: Verify the Training Step Before Scaling It

The training run in this note lowers loss on its two training phrases from 2.199037 to 0.224955. Loss on the held-out phrase rises from 2.239899 to 4.832658. Both results matter. A working optimization loop is not evidence that the model learned a useful language task.

The experiment trains a tiny, randomly initialized GPT-2 architecture on CPU using Hugging Face Transformers and PyTorch. It downloads no pretrained weights. Its purpose is to verify labels, padding, parameter updates and saved-model behavior before those details become expensive to debug on a larger run.

Decide what “train” means

Training a language model from random initialization and adapting a pretrained model are different projects. This example does the former at toy size: one transformer layer, two attention heads, embedding width sixteen and a vocabulary of nine tokens. Calling it a trained general-purpose LLM would be misleading.

A pretrained fine-tuning job would load an existing checkpoint and its matching tokenizer, then optimize a selected set of parameters on new data. The model revision, tokenization, objective and evaluation data become part of that job's specification. None of those choices can be inferred from a successful import of Trainer or from a low training loss.

The Transformers causal-language-modeling guide provides the broader workflow. Here, a short explicit PyTorch loop keeps the model object, batch and optimizer visible. The tested versions are Python 3.11.5, Transformers 4.56.2, PyTorch 2.9.1 and NumPy 2.2.6.

Make a batch whose targets you can inspect

The training phrases are red means rouge and blue means bleu. The held-out phrase is green means vert. A fixed word-to-ID dictionary makes this a token-training fixture, not a learned tokenizer or a translation dataset. A distinct end token is appended to each phrase.

Padding has its own ID. The batch stores three tensors: input IDs, an attention mask and labels. Labels initially copy the inputs, then only padding positions become -100, which the loss ignores. On a short blue sequence, the labels are [3, 1, -100, -100]: the end token remains a target.

That distinction matters when a tokenizer shares an end-token ID with padding. Masking every occurrence of that ID can remove genuine end-of-sequence targets. This fixture uses separate IDs to make the intended behavior unambiguous.

GPT-2's language-model head shifts the labels internally; see the versioned model documentation. The test compares its returned loss against cross-entropy between logits[:, :-1] and labels[:, 1:]. Manually shifting labels before that call would train a different alignment. An attention mask alone also does not replace the loss mask: the test checks the label tensor itself.

Update the same model you evaluate

The training function receives a model, creates an optimizer over its parameters, clears gradients, computes loss, backpropagates and steps the optimizer. It rejects nonfinite loss and gradients. There is no second model initialized between training and evaluation.

def train(net,data,steps=60):
    if type(steps) is not int or steps < 1: raise ValueError('positive integer steps required')
    optimizer=torch.optim.AdamW(net.parameters(),lr=.02)
    net.train()
    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        result=net(**data)
        if not torch.isfinite(result.loss): raise RuntimeError('nonfinite loss')
        result.loss.backward()
        if any(p.grad is not None and not torch.isfinite(p.grad).all() for p in net.parameters()):
            raise RuntimeError('nonfinite gradient')
        optimizer.step()
    return net

The complete implementation, tests, pinned requirements and setup instructions are downloadable. The fixture uses CPU explicitly, one computation thread, a fixed seed and zero dropout. After installing the dependencies, run:

python3 -B -m unittest -v test_example.py

Five tests pass. Besides checking label alignment and padding, they verify that an actual parameter tensor changes after a step, the training loss falls, invalid batches fail, and a saved model produces matching logits after reload. These are useful regression checks because a loop can appear to run even when it is updating a model that later code never uses.

Read both losses

The recorded run uses sixty AdamW steps with learning rate 0.02. The values are losses averaged over the predicted, unmasked tokens in each batch; they are not accuracy percentages or a quality benchmark.

The model becomes better at the particular training sequences and worse on this held-out sequence. That observation supports a simple diagnosis for this toy: fitting the training examples did not generalize to the unseen color pair. It does not establish a statistically reliable evaluation metric from one held-out phrase.

The small vocabulary also explains why this is not a useful translator. The model receives no evidence connecting green and vert in training. Increasing the step count cannot create that missing supervision. A realistic dataset needs representative examples, documented provenance and a split made before tuning decisions start using the evaluation results.

Save the result, then test the saved result

After evaluation, the experiment saves the trained model with save_pretrained, reloads it from the temporary local directory and compares logits in evaluation mode. The recorded comparison passes. This checks model-weight persistence; it does not verify resumable training because optimizer and scheduler state are not part of this round trip.

A full training artifact also needs the tokenizer or vocabulary, preprocessing configuration and model revision. The tiny vocabulary lives explicitly in the downloadable source. In a real fine-tuning job, keep its matching tokenizer files beside the checkpoint and test a fixed input through the deployed loading path.

Scaling this exercise requires new evidence: a representative validation set, an appropriate pretrained model, measured memory use and a run configuration that records precision, sequence length and batch behavior. This CPU result supplies none of those hardware limits. What it supplies is a small testable contract: the intended targets reach the loss, the intended weights change, and the saved weights are the ones evaluated afterward.