I Built a Complete LLM Training Pipeline in Pure PyTorch

#ai #python #opensource #machinelearning

If you’ve spent any time working with LLMs recently, you’ve probably noticed that we are increasingly reliant on massive, high-level frameworks like transformers, trl, and peft. They are fantastic for getting a model fine-tuned by Friday afternoon, but they come with a cost: they hide almost all of the actual math and mechanics behind layers of abstraction.

I wanted to actually understand what was happening under the hood. How does a Key/Value cache actually speed up generation? What does the loss mask for Supervised Fine-Tuning look like? How do you implement DPO or GRPO without leaning on trl?

To answer these questions, I built the entire pipeline from the ground up using only PyTorch's tensor and autograd primitives.

The repository is live here: Text-LLM-Training-from-scratch

Here is a breakdown of what the pipeline looks like and how the architecture comes together.

1. Zero-Bloat Data and Tokenization

Before even touching the model, I wanted to ensure the data pipeline was transparent and memory-efficient.

Byte-Level BPE Tokenizer: I wrote a custom tokenizer from scratch. Text is split into chunks by regex, reduced to raw UTF-8 bytes, and merged based on frequency. The training loop repeats until the target vocabulary size is hit. (Though you can easily swap in tiktoken if you prefer).
Memory-Mapped Token Shards: Tokens are stored as a single flat uint16 file. During training, the data is memory-mapped, meaning the trainer draws random windows directly from disk. Your corpus never has to fit into RAM, which is crucial if you are training on consumer hardware.

2. A Modern Decoder-Only Architecture

I didn't want to build a toy model based on the original 2017 formulation; I wanted the codebase to reflect how modern open-weights models are actually built. The architecture includes:

RoPE (Rotary Position Embeddings): Replacing learned position tables so attention scores extrapolate beyond training length.
RMSNorm & SwiGLU: Replacing LayerNorm and ReLU for better stability and performance at equal parameter counts.
Grouped-Query Attention (GQA): Essential for shrinking the cache and speeding up decoding.
A Proper K/V Cache: Without a cache, generation cost grows quadratically. With it, the prompt is processed once, past keys/values are retained, and subsequent tokens require only a single forward pass.

3. The Full Alignment Stack (SFT, DPO, and GRPO)

Pretraining is only half the battle. I implemented the entire alignment stack natively:

Supervised Fine-Tuning (SFT): Conversations are rendered into a ChatML-style stream. I implemented strict prompt masking so only the assistant’s tokens contribute to the loss—the model is trained to respond, not just reproduce the input.
Direct Preference Optimization (DPO): Implemented using the summed response-token log-probabilities of the original formulation against a frozen policy. Because it samples no rollouts, it’s remarkably cheap to run.
GRPO / RLVR (Verifiable Rewards): Instead of a learned reward model, the reward is a program verifying a final answer or format. It normalizes rewards within each group of sampled completions.

Does It Actually Work?

Yes! The repository includes a config for a 17.3M-parameter model (tinystories-15m.json).

Trained on the TinyStories dataset for just 30M tokens (around 3,700 steps at batch size 32), the validation loss drops from 9.0 to 1.90. It successfully generates coherent short narratives, and because I centralized device routing, the exact same code runs on CUDA, Apple Silicon (MPS), and CPU without any modifications.

I also built an interactive terminal UI (the llm command) that walks you through every stage from data prep to alignment, printing the exact bash commands it executes so you can learn the CLI usage as you go.

Final Thoughts

If you’ve been relying on high-level APIs and want to see how modern LLM training actually functions at the tensor level, I hope this repository serves as a useful educational resource.

Feel free to tear the code apart, suggest optimizations (especially around memory management or FlashAttention integration), or use it as a base for your own experiments.

Check out the code and the full documentation on GitHub and let me know what you think in the comments!