mesoGPT

Building and training large language models from scratch.

I built mesoGPT to learn how the main parts of language-model training work together and along the way, I worked with tokenizers, data pipelines, model architectures, compute planning, cost estimations, and efficiency.

Spiral Galaxy illustration

The current model is a 97.7M-parameter decoder-only Transformer trained on a single NVIDIA L40S. It is not intended to compete with production language models, and it does not introduce a new architecture. The purpose is to build a system large enough for training and systems problems to appear, while remaining small enough to inspect, measure, and run without a distributed cluster.

What “from scratch” means here

“Training a language model from scratch” can mean several different things, so it is useful to be precise.

For this project, it means that the tokenizer was trained separately, the model began training from randomly initialized weights, and the Transformer architecture, data pipeline, training loop, evaluation, checkpointing, and experiment tracking were assembled in PyTorch.

It does not mean that every underlying operation was implemented from first principles. mesoGPT relies on PyTorch, CUDA, rustbpe, tiktoken, PyArrow, and an existing text dataset. I did not write custom CUDA kernels or implement Flash Attention itself.

That distinction matters. The goal was to understand how the pieces of language-model training fit together, not to reproduce every layer of the software stack.

The complete pipeline

The project follows a relatively conventional pretraining pipeline:

ClimbMix Parquet shards
        ↓
Byte-level BPE tokenizer
        ↓
Streaming token dataloader
        ↓
Decoder-only Transformer
        ↓
Pretraining and validation
        ↓
Checkpoint and metadata
        ↓
Autoregressive text generation

None of these stages is especially mysterious in isolation. The difficulty comes from making them work together.

The tokenizer’s vocabulary size determines the model’s embedding and output dimensions. Choosing the right vocabulary size is crucial for balancing model capacity and training efficiency. The dataloader must produce correctly shifted input and target sequences. The training loop must handle gradient accumulation, mixed precision, and evaluation without introducing instability. Checkpointing must capture the model state, optimizer state, and training metadata. Finally, the generation interface must correctly use the trained model to produce coherent text.

So we have to make sure all components work together without introducing bugs or numerical instability. That is the essence of training a language model from scratch.

Starting with the tokenizer

Before training the language model, I had to train a tokenizer.

The tokenizer is a byte-level BPE tokenizer with a fixed vocabulary of 16,384 tokens. rustbpe is used to learn the vocabulary, while tiktoken provides fast encoding and decoding. Nine special tokens are reserved for possible later use.

The experiment trained eleven tokenizers using corpus sizes ranging from 25 million to 500 million characters. Each tokenizer was evaluated on the same held-out shard using:

  • Characters per token
  • UTF-8 bytes per token
  • Tokens per word
  • Total validation tokens produced

The expectation was straightforward: more tokenizer-training text should improve held-out compression, but the improvement should eventually flatten.

That is roughly what happened.

The tokenizer trained on 450M characters produced the best measured result, reducing the held-out corpus to 57,613,409 tokens. However, most of the gain had already appeared much earlier. The 250M-character tokenizer captured approximately 97.2% of the total token reduction observed between the 25M and 450M conditions.

Training on 250M characters took 19.05 seconds. Training on 450M took 32.75 seconds, while reducing the validation output by only another 6,246 tokens.

The results were also not perfectly monotonic. The 400M and 500M conditions were slightly worse than nearby smaller conditions.

The useful result was not that 450M was a universally correct training size. It was that tokenizer compression showed clear diminishing returns under this particular vocabulary, dataset, and evaluation setup.

The model

The baseline model is a decoder-only Transformer with the following configuration:

mesoGPT model configuration

The implementation uses Pre-LayerNorm blocks, a fused QKV projection, rotary positional embeddings, and tied token-embedding and language-model-head weights.

QKᵀ → scaling → causal mask → softmax → dropout → multiply by V

This makes the operation easy to inspect. However, as the later performance experiment showed that it is not an efficient way to execute attention in this manner.

Planning the training run

I used a tokens/parameter ratio of approximately 20, inspired by the Chinchilla scaling laws

97,655,296 parameters × 20 ≈ 1.953 billion tokens

With a global batch of 512 sequences and a context length of 1,024, each optimizer update processes:

512 × 1,024 = 524,288 tokens

Reaching the planned budget therefore required 3,726 optimizer updates.

The micro-batch contained 16 sequences. 32 gradient-accumulation steps were used to reach the global batch size of 512 without requiring the entire batch to fit in GPU memory at once.

Training used AdamW, BF16 mixed precision, 40 warmup steps, and a cosine learning-rate schedule from a peak 3e-4 to 3e-5.

The baseline run

The full baseline run completed on one NVIDIA L40S in 20 hours, 47 minutes, and 48 seconds. The measured rental cost was $38.70.

The best checkpoint appeared at optimizer step 3,720:

mesoGPT baseline training results

The final step was slightly worse, so the checkpointing code correctly retained the step-3,720 model.

Bits per byte, or BPB, is useful when evaluating models with a custom tokenizer. Cross entropy loss depends on number of tokens generated, and number of tokens generated for a fixed text depends on the compression of a tokenizer. So BPB was used as a metric, making the result independent of one tokenizer’s compression.

Training and validation loss remained very close throughout the run. The validation shard came from the same source distribution as the training shards, so similar curves are reasonable. They were calculated from different data, but their closeness should not be interpreted as evidence of broad generalization.

The run established that the implementation could train a roughly 100M-parameter model stably.

GPU utilization was not the same as efficiency

The baseline reported an average GPU utilization of 96.4%, with the GPU near 100% active utilization for much of the run. At first glance, that sounds efficient.

However, the estimated end-to-end model FLOP utilization was only 4.22%.

These measurements are not contradictory. GPU utilization usually indicates whether the device is busy. It does not indicate how close the arithmetic workload is to the GPU’s theoretical capacity. A GPU can remain busy while spending substantial time on moving intermediate tensors, launching kernels and handling other overhead.

The explicit attention implementation was a likely source of the problem. With a context length of 1,024, it constructs the full attention-score tensor, applies a causal mask, runs softmax, and stores intermediate results needed for backpropagation. That is readable, but expensive in both memory traffic and storage.

Testing PyTorch’s optimized attention path

The next experiment replaced the explicit attention sequence with PyTorch’sscaled_dot_product_attention API:

torch.nn.functional.scaled_dot_product_attention(
    q, k, v, dropout_p=..., is_causal=True,
)

The experiment logs indicated that PyTorch selected its built-in Flash Attention backend for the tested configuration.

To be clear, I did not implement Flash Attention. I changed the model to use PyTorch’s SDPA.

This was a performance benchmark, not another complete pretraining run. The architecture, parameter count, global batch, context length, and precision remained unchanged. Micro-batch sizes of 16, 32, and 64 were tested over short runs.

For an apples-to-apples comparison, the reported throughput and MFU excluded evaluation and used optimizer-update time only:

Optimized attention benchmark results

The best configuration was still the smallest tested micro-batch size: 16 sequences.

Relative to the explicit implementation, it delivered approximately 4.09 times the throughput. Increasing the micro-batch beyond 16 made both MFU and throughput worse in this setup.

That result is a useful reminder that using larger micro-batch sizes does not automatically improve hardware efficiency. Large micro-batch sizes increase memory traffic and automatically reduce MFU and throughput.

The 4.94% baseline MFU in this table differs from the earlier 4.22% figure because the two numbers answer different questions. The former measures training-only optimizer update time for the benchmark comparison. The latter uses the wall-clock duration of the complete baseline run, including evaluation and other overhead. Keeping those definitions separate helps us to make a fair comparison.

What I learned

The project reinforced several lessons.

First, tokenizer design is part of model design. Vocabulary size determines how much text fits into a fixed number of tokens and how much compute is required to process a corpus. It should not be treated as an unrelated preprocessing detail.

Second, end-to-end accounting matters. Parameter count alone does not determine the cost of a run. Evaluation frequency, checkpointing, and data loading all affect runtime.

Third, high device utilization is not proof of high model throughput. GPU telemetry and MFU describe different aspects of the system, and both are needed to understand performance.

What the project does not establish

There are some limitations to keep in view.

The model has not been instruction-tuned or aligned, and it should not be presented as a chat model. Its validation loss demonstrates that next-token pretraining worked, but it does not establish factual reliability, reasoning ability, or usefulness in deployment.

The project also uses existing frameworks, libraries, and data. Its contribution is the implementation and study of a complete training pipeline at a practical scale, not a new optimizer, dataset, attention algorithm, or model architecture.

Closing thoughts

mesoGPT did not produce a frontier model, and that was never its purpose.

What it produced was a concrete understanding of how a language model moves from raw text to a trained checkpoint: how a tokenizer changes the compute budget, how GPU utilization can be misleading, and how a single implementation choice in attention can change MFU and throughput by several times.

For me, that is the value of building at this scale. The model is large enough that the engineering is real, but small enough that the entire system can still be understood.

If you have any questions, I recommend using DeepWiki to get answers in a structured and organized manner.

The code, experiment reports, configurations, logs, and checkpoints are available in the mesoGPT repository.

Leave comment