11 Train a miniature language model from random weights

What this category does

A generative language model predicts a continuation of a token sequence. Useful application patterns include drafting a constrained message, completing a structured record, generating a response conditioned on evidence, and proposing a tool call that another program validates. A language model is useful when the output is a sequence with flexible wording or structure. If the output is one of a few fixed labels, a classifier may be simpler.

This project teaches the mechanics by training a very small byte-level transformer on original synthetic records. It will not become a general assistant. The expected learning target is a narrow text format and recurring relationships within that format. Its limited scope makes failures easier to inspect.

You have already trained a three-parameter classifier. The genuinely new elements are ordered byte sequences, shifted next-byte targets, and causal attention. Automatic differentiation, mini-batches, and tensor checkpoints build on the neural image projects you have already used. The input-to-target idea remains the same.

First generate the raw material

From the companion root, run:

python examples/tiny-transformer/make_data.py --output data/tiny-text

The generator creates records about colored household objects in rooms. A record is followed by questions and answers about that record. The generated text is original and deterministic. The training split contains 800 entity IDs, validation contains 100 different IDs, and test contains another 100. No entity is assigned to two splits.

Open train.txt before proceeding. A typical block looks like this, with values determined by the generator:

Record item0042: the blue lamp is in the office. Its switch is on.
Question: Where is item0042? Answer: the office.
Question: What color is item0042? Answer: blue.

The exact record above illustrates the format rather than claiming a particular generated item has those values. Inspect your actual file. The manifest records split sizes, seeds, and SHA-256 hashes. In the prepared CPU generator check, the training file contained 132,852 bytes. This is intentionally tiny. It is enough to exercise the pipeline, not enough to pretrain a broadly capable model.

Why this architecture needs this kind of data

The model accepts integer byte IDs and predicts the next byte. Every adjacent pair in the text provides a supervised target. It does not directly receive a label saying 'understand room location'. Instead, the loss rewards continuations that match the records and answers. If you want reliable question answering, you must separately test whether it uses the relevant earlier record rather than repeating frequent answers.

The input context is 128 bytes by default. A target depending on information farther away cannot be solved through that context window. If a record and question become separated by a long paragraph, the model may lack the necessary evidence. This connects architecture to data design: examples must expose the dependencies you want the model to learn within the representation it can use.

The training program samples fixed-length windows from each split's continuous byte stream. A window can cross a record boundary within that split. That is acceptable for this small continuation exercise because the stream format includes separators and no split boundary is crossed. For tasks where documents must be independent, use explicit boundary tokens and appropriate masking rather than blindly joining unrelated records.

Install only the small project environment

The reference API pin is torch 2.8.0 with Python 3.11. This older fixed release is intentional for the educational implementation; it is not a claim about the newest PyTorch version. Use the official PyTorch installation selector or previous-version instructions to select the correct CPU or CUDA wheel for your platform. The release and versioned attention API are documented by PyTorch. F10 F11

Create a separate virtual environment, install the matching official wheel, and verify the version and device. The project requires no Transformers package, no tokenizer download, and no remote model code. Do not copy a CUDA-specific wheel command to a machine with incompatible hardware or driver.

For a Linux or Windows Python 3.11 environment, the official 2.8.0 CPU wheel can be installed with the first command below. The second is the CUDA 12.6 build for a compatible NVIDIA driver and GPU; choose one, not both. A newer GPU may require a different supported CUDA build. The official previous-version page lists the alternatives. F15

# CPU environment
python -m pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu

# Separate compatible NVIDIA CUDA environment
python -m pip install torch==2.8.0 --index-url https://download.pytorch.org/whl/cu126

Run the local model checks first:

python examples/tiny-transformer/check_model.py

That script tests output shapes, future-token isolation, finite gradients, a real parameter update, and save/load equivalence. It is provided for execution in an environment with PyTorch. During preparation of this edition, the source was syntax-checked but PyTorch execution was unavailable; the validation appendix states that boundary explicitly.

Run a smoke test on the CPU

python examples/tiny-transformer/train.py \
  --data data/tiny-text --output runs/byte-smoke \
  --device cpu --steps 20 --warmup 2 \
  --width 32 --layers 2 --heads 4 \
  --context 32 --batch-size 2 --eval-every 10

This reduces the model and sequence sizes so a basic CPU test is practical. The first run is a correctness check. Look for a finite loss, a printed parameter count, validation output, and checkpoint.pt. Do not judge the model's usefulness after twenty updates.

The script never silently substitutes CPU for a requested unavailable GPU. A device error should be fixed directly. Likewise, BF16 requires an explicitly compatible CUDA device. These checks prevent a long run from proceeding under an unintended environment.

Move the same pipeline to the GPU

After the CPU checks pass, a small GPU experiment can use the default architecture:

python examples/tiny-transformer/train.py \
  --data data/tiny-text --output runs/byte-gpu \
  --device cuda --steps 1000 --warmup 50 \
  --context 128 --width 128 --layers 4 --heads 4 \
  --batch-size 8 --accumulation 2 --eval-every 100

Add --bf16 only after checking device support. The script keeps model parameters and AdamW states in FP32 and uses BF16 autocast for selected operations when requested. This is a clear numeric policy, not a claim that every tensor occupies two bytes. The default architecture has about 0.84 million distinct parameters by direct shape accounting, well below one billion.

The batch is eight windows of 128 bytes. Accumulating two microbatches presents 2,048 target bytes per optimizer update. One thousand updates therefore present about 2.048 million training bytes, sampled with replacement from the small training stream. That is repeated exposure, not two million unique bytes.

The code logs throughput including evaluation time, so it is an end-to-end run measure rather than a pure kernel benchmark. It also logs peak allocated and reserved GPU memory. Measure your own results; no particular RTX throughput or wall-clock duration is asserted here.

Understand the model one piece at a time

The vocabulary contains 256 byte values. A learned embedding maps each ID to a vector of width 128. A learned positional table adds information about the position inside the current window. Four transformer blocks repeatedly combine earlier information through causal attention and transform it through a feed-forward network. A final normalization and vocabulary projection produce 256 logits for each position.

The input and output vocabulary weights are tied, meaning the same parameter matrix is reused. The model uses learned positions and ordinary multi-head attention for clarity. It is not an implementation of Qwen, Gemma, Granite, or a modern production architecture. Those families differ in tokenization, normalization, position handling, feed-forward choices, attention layout, and other details.

PyTorch's scaled-dot-product attention receives is_causal=True. During evaluation, dropout_p is explicitly zero. This matters because the functional attention operation applies the dropout probability it is passed rather than automatically reading the surrounding module's evaluation state. F11

The next chapter explains the same components more deeply. First connect each component to the tensors printed by your own small run. Architecture becomes easier to understand when it names something you have already used.

Every training control in this project

--context sets the number of input bytes per window and the size of the learned position table. --width sets representation size. --layers sets the number of transformer blocks. --heads divides the hidden width into attention heads; width must be divisible by heads. --dropout randomly suppresses selected intermediate contributions during training and is disabled during evaluation.

--batch-size is windows per microbatch. --accumulation is microbatches per update. --steps is optimizer updates, not microbatches. --learning-rate is the peak update rate, default 0.0003. --weight-decay is AdamW's decoupled shrinkage strength, default 0.1. The educational implementation applies it uniformly; production recipes often use parameter groups that exempt biases and normalization parameters.

--warmup ramps the learning rate over early updates. Afterward a cosine schedule decreases it to one tenth of the peak. --eval-every controls evaluation and checkpoint intervals. --seed controls initialization and sampling seeds. --device selects CPU or CUDA. --bf16 enables compatible CUDA autocast. --data and --output select input and artifact directories. --resume loads an existing trusted checkpoint and requires the checked training configuration and dataset hashes to match.

The optimizer also sets betas to 0.9 and 0.95, epsilon to 1e-8, and clips the global gradient norm at 1.0. These are explicit starting choices for this experiment. The training-control chapter explains their meanings. They are not a recommendation to use the same values for every fine-tuning job.

Evaluate the behavior you actually want

The logged validation objective is natural-log loss per byte, accompanied by bits per byte. Lower is better for the byte-prediction objective. Do not compare byte-model perplexity directly with the perplexity of a model using a different subword tokenizer. The prediction units differ.

Evaluate a final selected checkpoint on the test file only after choosing the procedure:

python examples/tiny-transformer/use_model.py \
  runs/byte-gpu/checkpoint.pt \
  --device cpu --evaluate-file data/tiny-text/test.txt

Then create fresh records with object IDs and values not present in training. Put the necessary fact and question within the context window. Score whether the answer matches the provided fact. Compare against always choosing the most common room and against a deterministic parser for this artificial format. A transformer is educational here; a parser would be the practical solution if this exact rigid format were the entire real task.

Inference on the CPU and GPU

The same trusted checkpoint can be loaded on either device:

python examples/tiny-transformer/use_model.py \
  runs/byte-gpu/checkpoint.pt --device cpu \
  --prompt "Record item9999: the red lamp is in the hall."

python examples/tiny-transformer/use_model.py \
  runs/byte-gpu/checkpoint.pt --device cuda \
  --prompt "Record item9999: the red lamp is in the hall."

The inference program reconstructs the exact saved architecture, loads the state dictionary, calls eval, converts the UTF-8 prompt into byte IDs, moves tensors to the selected device, samples one next byte at a time, and decodes bytes back to text. Invalid UTF-8 sequences are replaced on decoding, which is a visible limitation of an immature byte model rather than a hidden text-cleaning success.

CPU and GPU use the same learned artifact but can differ slightly in floating-point results. Sampling can amplify small probability differences, so identical seeds do not guarantee identical text across devices. The tiny architecture may not benefit much from a GPU because launch overhead is significant. Larger models and batches generally change that balance. Measure latency for your intended request size.

This educational generator does not implement a key-value cache, so it recomputes the retained context at each generated byte. It also crops to the learned context length. Those are deliberate simplicity choices and become performance limits when scaling. Do not benchmark it as if it were an optimized production serving engine.

Checkpoint recovery and limits

checkpoint.pt contains weights, optimizer state, configuration, update number, and random-number-generator state. Saving uses a temporary file followed by a rename to reduce the chance that an interrupted write replaces the last complete checkpoint. Resume only trusted local checkpoints; serialized model files from unknown sources deserve security scrutiny even when safer loading modes are available.

The script checks key configuration fields and data hashes before resuming. It does not claim exact reproducibility across different PyTorch releases, operating systems, GPU architectures, or kernels. Its purpose is to make the important state explicit and reject obvious accidental changes.

Exercises

Use check_model.py to see how changing future input bytes should leave earlier logits unchanged. Explain why this is stronger evidence about causality than a low loss.

Train two small configurations with the same token-presentation budget. Double width in one and double updates in the other, keeping other choices controlled. Compare validation loss, fresh-record accuracy, elapsed time, and memory. Report the tradeoff rather than choosing a winner by parameter count.

Replace the synthetic corpus with a small collection of text you own or are licensed to use. Preserve document-level splits. State what behavior the corpus can plausibly teach and what it cannot. A collection of maintenance manuals may teach terminology; it does not automatically teach a conversational format or reliable tool use.