Train a tiny language model from scratch with PyTorch

TL;DR: Train a 470,528-parameter byte-level language model on your CPU. Start with the included synthetic text, watch held-out loss, save the optimizer and random-number state with the weights, then resume and generate text. This small experiment teaches the training machinery; it will not produce a useful general-purpose assistant.

1. Decide what this run should prove

A language model predicts what comes next. For this project, a token is one byte: an integer from 0 to 255. Using bytes avoids downloading a tokenizer and makes the input visible. A word can occupy several bytes, and a non-English character can use several UTF-8 bytes. Token counts here are therefore not comparable with a Qwen tokenizer's counts.

You will implement a two-layer causal Transformer. “Causal” means a position may read earlier positions and itself, but cannot inspect the next byte it is being asked to predict. The learning signal comes from comparing the model's prediction with that next byte. PyTorch computes derivatives of the loss, and an optimizer uses them to change the model's weights.

This is training from random initialization. To adapt an already capable small language model (SLM) to your task, follow the Unsloth fine-tuning guide, which requires a free account. Training a useful SLM from scratch adds a much larger data collection, tokenizer, evaluation and compute project. Complete this miniature run before deciding whether that expense is justified.

The success criteria are concrete: the script finishes with finite losses, validation improves on the included fixture, a checkpoint reloads, and continuing a saved run agrees with an uninterrupted run in the same CPU environment. Text quality remains a separate question.

2. Create an isolated CPU environment

Use a Linux terminal with Python 3.12, a few GB of free RAM and space for a Python environment. Save tiny_lm.py from Project files into an empty project folder. These instructions keep the older, executed PyTorch version so the example can be reproduced; they do not claim it is the newest release.

mkdir tiny-language-model
cd tiny-language-model
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install 'torch==2.6.0' --index-url https://download.pytorch.org/whl/cpu
python -m pip install 'numpy==2.2.6'
python -m pip check
python -c 'import torch; print(torch.__version__, torch.cuda.is_available())'
python -m pip freeze > environment.txt

The environment check should show PyTorch 2.6.0 with a CPU build and False for CUDA availability. A GPU is unnecessary. On Windows, use a Linux environment such as WSL2 for these shell commands; a native PowerShell activation command is different.

Place the downloaded script in this folder before the next step. environment.txt records the installed dependency set. Keep it with your checkpoint when you compare experiments.

First result: train, then sample

From the folder containing tiny_lm.py, run:

python tiny_lm.py --steps 25 --checkpoint first.pt
python tiny_lm.py --sample --checkpoint first.pt

The first command prints the parameter count, an initial validation loss and a step-25 record, then creates first.pt. The second reloads that file and prints generated bytes as text. Check for finite losses and a sample beginning with Service oak:. Readable prose is not the success condition of this first run.

This short run exercises the full save/reload path before you inspect the training loop. Keep it separate from the longer tiny.pt run below. If either command fails, use the symptom table near the end before changing the model.

3. Follow one training example

Suppose the original bytes spell GPU. One input position sees G and predicts P; the next sees GP and predicts U. The script makes 64-byte input windows and shifts each target window by one byte. This alignment is the central check: if inputs and targets are identical, a model can learn to copy instead of predicting.

At position 0, the model must predict P from G. The P already present at input position 1 is blocked. Target shifting names the prediction; the mask prevents reading its answer from the future.

The script embeds each byte ID into 128 numbers, adds a position embedding and passes the result through two attention-and-feed-forward blocks. The upper-triangular attention mask blocks future positions. Although the PyTorch class is named TransformerEncoderLayer, applying this mask makes the stack causal for this experiment. Its documented mask and shape interface is what matters here; this is not an encoder-decoder translation model.

Cross entropy takes raw scores, so the training path does not apply softmax first. Generation applies softmax later because it needs a probability distribution from which to sample.

4. Connect the model to its update

The default corpus contains original synthetic service-status sentences. Training and validation use different service names, but deliberately share sentence templates. That makes the example small and inspectable. It does not measure generalization to independently written prose.

Two lines in batch() establish the prediction task:

x = torch.stack([data[i:i + CONTEXT] for i in starts])
y = torch.stack([data[i + 1:i + CONTEXT + 1] for i in starts])

These are excerpts from the downloaded file, not separate programs. x contains the context and y contains the next-byte answers. In forward(), a boolean upper triangle marks the input positions that attention must ignore. After both layers, head maps each 128-number hidden vector to 256 next-byte scores.

The loss line flattens the batch and sequence dimensions together, producing 512 predictions per update. The target vector is flattened in the same order. Swapping those orders would compare valid-looking scores against the wrong answers.

Read the loop in this order: choose windows, compute predictions, compute loss, clear the old gradients, differentiate, clip gradients, update parameters. model.eval() and torch.no_grad() make validation a measurement rather than another update. The evaluator uses a separate fixed generator, so evaluating more often does not change which training windows are sampled.

The model has 470,528 trainable parameters. Stored as 32-bit values, their raw payload is 470528 × 4 = 1,882,112 bytes, about 1.88 MB. Training also needs gradients, optimizer state, activations and the Python/PyTorch process. This arithmetic is a weight-size calculation, not a claim that the whole process uses 1.88 MB. Model memory footprint explains those additional allocations.

5. Train and interpret the result

python tiny_lm.py --steps 200 --checkpoint tiny.pt

Every 25 steps the script prints a JSON record. A step means one optimizer update from eight 64-byte windows. At 200 steps, it has processed 200 × 8 × 64 = 102400 target positions. Windows are sampled with replacement, so this does not mean 102,400 unique bytes of training data.

Our CPU run used Python 3.12 and PyTorch 2.6.0+cpu. Its measured validation values were:

Small floating-point differences across platforms are possible. These numbers describe the included synthetic exercise, not a target quality score for your own text. Byte perplexity is exp(validation_loss); at step 200 it is about 3.32. Comparing that number with a subword model's reported perplexity would mix different prediction units and datasets.

If training loss improves while validation loss worsens, stop extending the run automatically. Check whether the model is memorizing recurring examples, whether the two datasets differ unexpectedly, and whether you changed the learning rate. A lower training loss alone is insufficient evidence to keep a checkpoint.

6. Reload, generate and resume

python tiny_lm.py --sample --checkpoint tiny.pt
python tiny_lm.py --resume --steps 250 --checkpoint tiny.pt

The sample command starts with Service oak: and generates 160 more bytes. It may produce recognizable fragments and nonsense in the same line. Invalid UTF-8 sequences are replaced for display. This is expected from a tiny byte model and is a useful reminder that falling loss and coherent answers are different checks.

--steps 250 means 250 total updates, so this command adds 50 updates to the saved 200-step run. A fresh run refuses an existing checkpoint path; choose a new name or use --resume deliberately. The checkpoint contains model weights, AdamW state, both random-number states, the completed step and corpus hashes. Saving only weights permits inference, but does not restore the same optimizer trajectory. PyTorch's checkpoint guidance explains this distinction.

Checkpoint writes go through a temporary file in the same directory before replacement. This avoids replacing the last file with a half-written serialization during an ordinary interrupted write; it is not a claim of durable storage across hardware failure. Load checkpoints you created or trust.

7. Replace the fixture and diagnose failures

For your own experiment, create separate UTF-8 train.txt and valid.txt files containing text you have permission to use. Split by original document or source before producing windows. Duplicating a document and putting one copy in each split defeats the purpose of validation.

python tiny_lm.py --train train.txt --valid valid.txt --steps 200 --checkpoint custom.pt
python tiny_lm.py --train train.txt --valid valid.txt --sample --checkpoint custom.pt

Both files must contain more than 64 bytes. The script refuses identical file hashes and refuses a resume when the corpus hashes change. Different hashes do not prove the documents are independent; inspect your split process too.

Before moving on, explain why the attention mask and shifted targets are both needed. Then change the corpus while keeping the model fixed and predict which saved-run check will fail. The hash mismatch should stop the resume before any update happens.

Deactivate the environment when finished with deactivate. Keep tiny.pt, your text and environment.txt if you want to continue. There is no cloud resource to terminate in this CPU project. Next, build a validated fine-tuning dataset; training cost estimates become useful when you expand beyond this local exercise.

Check your understanding

Why can an unmasked model show a very low training loss and still fail at generation? During training it can read the future byte it is supposed to predict. At generation time that byte does not exist yet. The causal mask keeps the two situations aligned.

What should happen if you resume 200 saved steps with --steps 250? You should see start_step: 200 and finish at 250, adding 50 updates. Compare this with loading only the model weights: the latter has no saved optimizer trajectory or sampled-window state.

Complete project source

These are the same files offered under Project files. Run the downloaded files; the listings let you inspect each implementation.

tiny_lm.py

"""A CPU byte-level causal language model. Demo text is original synthetic data."""
import argparse
import hashlib
import json
import math
from pathlib import Path
import torch
from torch import nn
from torch.nn import functional as F

CONTEXT = 64

class TinyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.tokens = nn.Embedding(256, 128)
        self.positions = nn.Embedding(CONTEXT, 128)
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(128, 4, 512, dropout=0.0,
                                       batch_first=True, norm_first=True)
            for _ in range(2)
        ])
        self.norm = nn.LayerNorm(128)
        self.head = nn.Linear(128, 256, bias=False)

    def forward(self, ids):
        n = ids.shape[1]
        x = self.tokens(ids) + self.positions(torch.arange(n, device=ids.device))
        mask = torch.triu(torch.ones(n, n, device=ids.device, dtype=torch.bool), 1)
        for layer in self.layers:
            x = layer(x, src_mask=mask)
        return self.head(self.norm(x))


def corpus(path, held_out=False):
    if path:
        raw = Path(path).read_bytes()
    else:
        names = ['cedar', 'maple'] if held_out else ['birch', 'oak', 'pine']
        raw = ''.join(f'Service {name}: the request waits for a free worker.\n'
                      f'Service {name}: save the checkpoint before stopping.\n'
                      for name in names for _ in range(40)).encode()
    if len(raw) <= CONTEXT:
        raise ValueError('Each split needs more than 64 bytes.')
    return torch.tensor(list(raw), dtype=torch.long), hashlib.sha256(raw).hexdigest()


def batch(data, generator, size=8):
    starts = torch.randint(len(data) - CONTEXT, (size,), generator=generator)
    x = torch.stack([data[i:i + CONTEXT] for i in starts])
    y = torch.stack([data[i + 1:i + CONTEXT + 1] for i in starts])
    return x, y


@torch.no_grad()
def evaluate(model, data):
    model.eval()
    generator = torch.Generator().manual_seed(91)
    losses = []
    for _ in range(8):
        x, y = batch(data, generator)
        losses.append(F.cross_entropy(model(x).reshape(-1, 256), y.flatten()).item())
    return sum(losses) / len(losses)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--steps', type=int, default=200, help='Total optimizer steps, including resumed steps')
    parser.add_argument('--train')
    parser.add_argument('--valid')
    parser.add_argument('--checkpoint', default='tiny.pt')
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--sample', action='store_true')
    args = parser.parse_args()
    if args.steps < 1 or bool(args.train) != bool(args.valid):
        parser.error('Use positive steps and supply both --train and --valid, or neither.')
    if not (args.resume or args.sample) and Path(args.checkpoint).exists():
        raise FileExistsError('Choose a new checkpoint path or explicitly resume the existing run.')
    torch.set_num_threads(2)
    torch.manual_seed(17)
    model = TinyLM()
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
    generator = torch.Generator().manual_seed(18)
    train, train_hash = corpus(args.train)
    valid, valid_hash = corpus(args.valid, True)
    if train_hash == valid_hash:
        raise ValueError('Training and validation files must differ.')
    step = 0
    if args.resume or args.sample:
        state = torch.load(args.checkpoint, map_location='cpu', weights_only=True)
        if state['data_hashes'] != [train_hash, valid_hash]:
            raise ValueError('Corpus changed; use the original files or start a new run.')
        model.load_state_dict(state['model'])
        optimizer.load_state_dict(state['optimizer'])
        generator.set_state(state['batch_rng'])
        torch.set_rng_state(state['torch_rng'])
        step = state['step']
    if args.sample:
        model.eval()
        ids = torch.tensor([list(b'Service oak:')], dtype=torch.long)
        with torch.no_grad():
            for _ in range(160):
                logits = model(ids[:, -CONTEXT:])[:, -1] / 0.8
                ids = torch.cat([ids, torch.multinomial(logits.softmax(-1), 1)], dim=1)
        print(bytes(ids[0].tolist()).decode('utf-8', errors='replace'))
        return
    print(json.dumps({'parameters': sum(p.numel() for p in model.parameters()),
                      'start_step': step, 'valid_loss': evaluate(model, valid)}))
    for step in range(step + 1, args.steps + 1):
        model.train()
        x, y = batch(train, generator)
        loss = F.cross_entropy(model(x).reshape(-1, 256), y.flatten())
        if not torch.isfinite(loss):
            raise RuntimeError('Non-finite loss; stop and inspect the batch and learning rate.')
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        if step % 25 == 0 or step == args.steps:
            val = evaluate(model, valid)
            print(json.dumps({'step': step, 'train_loss': loss.item(),
                              'valid_loss': val, 'byte_perplexity': math.exp(val)}))
    state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(),
             'step': step, 'batch_rng': generator.get_state(),
             'torch_rng': torch.get_rng_state(), 'data_hashes': [train_hash, valid_hash]}
    target = Path(args.checkpoint)
    temporary = target.with_suffix(target.suffix + '.tmp')
    torch.save(state, temporary)
    temporary.replace(target)

if __name__ == '__main__':
    main()