Building a 350M Transformer From Scratch: The Complete Guide

In a world where training a language model usually means renting a cluster of GPUs for thousands of dollars, one developer proved you can build and train a 350M-parameter transformer from scratch for just $101. No frameworks, no shortcuts — just raw PyTorch, a single H100, and a deep commitment to understanding every line of code.

This is a comprehensive breakdown of John's 4-part series documenting the entire journey: from writing the first attention layer to serving a model that speaks fluent English. It's one of the most transparent, well-documented from-scratch LLM training projects I've seen, and it covers everything most tutorials skip: the architecture decisions, the training gotchas, and the expensive hardware mistakes.

The project spans four parts, each building on the last: V1 established the baseline architecture, V2 modernized it with cutting-edge techniques, V3 scaled to 672M parameters across two GPUs, and Part 3 covered post-training with SFT and GRPO to turn the base model into something actually useful. Every model is open source on HuggingFace, and the code is on GitHub [Link].

What makes this series stand out isn't just the technical content — it's the honesty about what went wrong. The H100 NVL vs SXM mistake that cost $35 extra. The checkpoint bugs that would have killed a 20-hour run. The subtle causal mask leak that silently let the model peek at future tokens. These are the details that separate a tutorial from a real engineering journal.

The V1 Baseline Architecture

The first model, V1, was designed as a clean baseline: modern architecture without any experimental additions. The goal wasn't to push the state of the art — it was to code the building blocks of a modern LLM from scratch, prove the model could learn, and establish numbers to measure future improvements against.

Here's the configuration that defined V1:

• d_model: 1024 — Width of the residual stream, the model's core capacity
• n_layers: 24 — Network depth, each layer applying attention then feed-forward
• n_heads: 16 — Parallel attention heads per layer
• seq_len: 1024 — Context window, how far back the model can look
• Vocabulary: 50,257 — GPT-2 BPE tokenizer via tiktoken, padded to 50,304 for hardware efficiency

The architecture uses Multi-Head Attention with RoPE for positional encoding, RMSNorm for layer normalization, SwiGLU in the feed-forward blocks, tied embeddings, and AdamW as the optimizer. It's an upgrade on GPT-2's design, without reaching the complexity of something like Llama 3.

What V1 deliberately does not have is equally important: no GQA, no QK-Norm, no Muon, no Differential Attention, no XSA. Those come in V2 and V3. The series is about evolution, not pretending the first version was already perfect.

Architecture Deep Dive

The Transformer Block

At its core, a large language model does one thing: it takes a sequence of tokens as input and predicts the next one. Everything else — conversations, code generation, reasoning — emerges from doing that one thing extremely well, billions of times over.

The forward pass through each transformer block follows a clean pattern:

x = x + Attention(RMSNorm(x))
x = x + FFN(RMSNorm(x))

Two things worth unpacking: pre-norm and the residual stream.

Pre-Norm vs Post-Norm

The original transformer used post-norm: normalize after each sublayer. This causes problems in deep networks. During backpropagation, gradients travel back through all 24 layers. If they shrink at each step (vanishing gradients) or blow up (exploding gradients), the early layers stop learning. Post-norm makes this worse by disrupting the residual stream at every layer.

Pre-norm fixes this by normalizing before the sublayer, then adding the result back to the original input. The residual stream stays clean throughout, and that's what lets gradients flow freely. Because each block adds its output to x rather than replacing it, the residual path contributes a derivative of 1 during backpropagation. That constant 1 means the gradient always has an unobstructed route back to earlier layers.

RMSNorm

The original transformer used LayerNorm: for each token, compute the mean and variance across the model dimension, then normalize so the result is centered around 0 with variance 1. Two learned parameters — a scale γ and a shift β — are applied on top.

RMSNorm removes the mean-centering step entirely. After experimentation, researchers found that mean-centering wasn't pulling its weight: later layers can re-center a representation themselves if needed. What actually matters for stability is controlling the spread of values, not their center. The formula squares every value, takes the mean, adds a small constant ε, takes the square root, then divides every value by this RMS before scaling by γ.

The result: simpler, cheaper, faster. Each RMSNorm layer holds just 1,024 values (γ only). Across 24 blocks with 2 norm layers each, that's roughly 50K parameters — about 0.015% of the model. But the bigger win is compute: skipping mean-centering removes a full reduction from every forward and backward pass, speeding up training by roughly 10–15% compared to LayerNorm.

Multi-Head Attention

Attention is the mechanism that lets every token look at every previous token in the sequence and decide which ones matter most. Take the sentence "The animal didn't cross the street because it was too tired." What does "it" refer to? Attention is how the model figures that out: when processing "it", it learns to look back at "animal" and weigh it heavily.

At a high level, attention relies on three matrices — Query (Q), Key (K), and Value (V) — each with a distinct role. A clean analogy is a library: the Query is your question ("what am I looking for?"), the Key is each book's title ("I'm a book about mammals"), and the Value is the book's actual content. You match your query against every key to decide which books are relevant, then pull information from their values, weighted by how well they matched.

The scaled dot-product attention formula:

Attention(Q, K, V) = softmax(Q · K^T / √d_head) · V

The scaling by √d_head isn't arbitrary. Weights are initialized around N(0, 1), so the dot product sums d_head independent products, ending up distributed like N(0, d_head). Dividing by √d_head pulls it back to N(0, 1). Without this, the softmax saturates — concentrating almost all weight on one or two tokens and starving everything else.

Multi-head attention computes this 16 times in parallel. Each head can specialize: one might focus on grammar, another on long-range dependencies, another on temporal relationships. Random initialization pushes heads apart, and gradient descent's "laziness" means once head A captures syntax, the cheapest way to reduce remaining error is for head B to capture something else.

Causal Masking

Decoder-only models must never see the future. This is enforced with a causal mask: a matrix added to attention scores before softmax, where all future positions are set to -∞. After softmax, those become exactly 0. The model is structurally incapable of attending to anything ahead of the current position.

One of the most critical tests: assert_no_future_leak. Change a token at some future position and verify that outputs at earlier positions don't change at all. If an earlier output shifts when you edit a later token, information is leaking backward through time, and the causal mask is broken. This kind of behavioral test catches what a code review won't.

Flash Attention

Standard attention scales quadratically with sequence length. Q and K each have shape [seq_len, d_head], so QK^T produces a [seq_len, seq_len] matrix. For 1024 tokens, that's 1.05M values per head. Flash Attention solves this by computing attention block by block, keeping working chunks in fast on-chip SRAM instead of writing the full matrix to slower HBM. Same math, fraction of the memory traffic, big speedup.

RoPE: Rotary Position Embeddings

Attention computes all scores in parallel, which means it has no inherent sense of position. If you shuffled the tokens, you'd get the exact same attention scores. "Dog bites man" is not "man bites dog." RoPE solves this by rotating query and key vectors in 2D planes by angles proportional to their position.

Imagine our model dimension is 2 instead of 1024. Each token has its own query and key vectors. RoPE rotates these in the plane by an angle proportional to the token's position, like the hands of a clock. When it's time to compare two vectors via dot product, what matters is the angle between them — which comes down to the difference in their positions. The further apart two words are, the larger the angle.

Technically, Q and K vectors are split into pairs of dimensions (32 pairs per head at d_head=64), and each pair gets its own 2D plane to rotate in. Each pair rotates at a different frequency: earlier pairs rotate faster (sensitive to small position differences), later pairs rotate slower (stay distinguishable across large gaps). A single shared frequency would force the model to trade off short-range precision against long-range reach.

Crucially, only Q and K are rotated — never V. V carries the token's actual content, for which position is irrelevant. Q and K determine attention weights, and that matching is where position matters.

The beauty of RoPE: it has no learned parameters. The rotation frequencies are fixed, precomputed once, and cached. It encodes relative position, extrapolates beyond training context length, and is the default across virtually every modern LLM.

SwiGLU: The Feed-Forward Network

After attention lets tokens share information (the meeting), the feed-forward network lets each token process what it learned individually (going back to your desk to think). SwiGLU is a modern upgrade on the vanilla transformer's MLP.

Instead of a single expansion path, SwiGLU uses two parallel projections:

• One path expands x with W_up, then applies Swish activation
• The other path expands x with W_gate (linear)

The two are multiplied elementwise. Because both are learned projections, this lets the model decide how much of each feature to let through — amplifying some, suppressing others. The gate gives the network fine-grained control over which features pass. The result is projected back down with W_down.

The Swish activation — x · sigmoid(x) — is smooth and non-monotonic around 0, giving the network more expressive flexibility than ReLU. For negative inputs, it pushes output close to 0 (but not exactly 0). For positive inputs, it approaches x.

To keep parameter count comparable to a standard FFN, SwiGLU uses an 8/3× hidden dimension per matrix instead of 4×. Three matrices at 8/3 gives 8× total — the same as two matrices at 4× each.

The FFN is where most of the model's parameters live. In V1's case, over 57% of the 353M total parameters reside in the feed-forward layers. Attention gets the attention, but the FFN holds the weight.

Training: From Code to Loss Curve

Data

FineWeb-Edu: web text filtered for educational quality, from HuggingFace's FineWeb family [Link]. Why FineWeb-Edu over raw Common Crawl? At small scales, quality matters more than raw volume.

10B tokens total — a ratio of ~29 tokens per parameter, deliberately past the Chinchilla-optimal ~20. The data was sharded into blocks of 100M tokens (100 shards total), with 2% held out for evaluation.

Training Configuration

• Optimizer: AdamW with decoupled weight decay
• Learning rate: Cosine decay with 1000-step warmup, max 3e-4, min 3e-5 (10× ratio)
• Batch size: ~524K tokens/step via gradient accumulation (16 sequences × 32 accumulation steps)
• Precision: bf16 mixed precision (same exponent range as fp32, less mantissa precision)
• Gradient clipping: Threshold of 1.0
• Steps: ~19,000 total across 10B tokens

Checkpointing: The Unseen Bugs

Checkpointing seems simple until something kills your run at hour 19. The author tested resume logic by deliberately killing the process mid-training, surfacing three bugs:

1. Atomic save: If killed mid-torch.save, you get a corrupt checkpoint. Fix: write to .tmp first, then os.replace (atomic rename).
2. Skip-corrupt fallback: find_latest_checkpoint goes newest-first and skips unreadable files, so a corrupt latest checkpoint falls back to the previous good one. You lose ~1000 steps, not the whole run.
3. RNG state restore: Restoring RNG state with map_location="cuda" threw a ByteTensor error — RNG state has to stay on CPU.

None of these show up until something interrupts your run. Finding them before real training is exactly why the resume test was worth writing.

The Hardware Mistake

Here's the expensive lesson: the author started on an H100 NVL ($3.2/hr), hitting ~83K tokens/sec. After pausing and resuming, they got an H100 SXM ($3.3/hr): nearly the same price, but 140K tokens/sec — a 1.7× speedup. They'd done ~85% of training on the slower GPU.

The NVL choice cost about $35 extra — roughly the price of the entire run over again, for a card that was only 3 cents/hr cheaper. Lesson: for H100 training, always take the SXM.

Results: Did It Work?

The training loss shows a steep early drop from ~11 (roughly log(50,304) — random guessing across the vocabulary), then settles gradually around 2.8. This is a textbook loss curve: exactly the shape you want to see. It signals healthy training.

The final validation loss lands around 2.86. Running the math in reverse: e^2.86 ≈ 17.5. So the model's perplexity is about 17 — at each step, it's effectively choosing among ~17 plausible tokens rather than 50,304. Given that English has synonyms and many ways to say the same thing, that's a reasonable place for a 350M model to land.

Validation loss follows the same curve as training loss and never turns back up. That's what you want: an increasing val loss while train loss keeps dropping would signal overfitting. Here, the two stay locked together.

The gradient norm spikes during warmup, then settles to a stable ~0.2 for the rest of training. No explosions, no instability — gradient clipping at 1.0 was rarely needed after warmup.

Qualitative Progress

Every 500 steps, the model generated completions from fixed prompts. Watching these evolve is the most intuitive way to see learning happen:

• Step 100: Near-random word salad. Shape of text, no coherence.
• Step 3,500: Grammatical, locally coherent, topically on-track.
• Step 8,500: More fluent, holding a thread across longer spans.
• Step 17,500: Fluent, structured, even attempting attribution and narrative.

You can watch grammar emerge first, then local coherence, then longer-range consistency. The Fibonacci prompt is telling: the model recognizes it's a math context and reaches for formula-like notation, but has no real grasp of the algorithm — it's pattern-matching the shape of an answer, not reasoning.

V1 is a base model — pretraining only. It understands English but isn't useful yet. No concept of a user, assistant, or prompt-to-answer format. Instruction-following comes through SFT and reinforcement learning, covered in Part 3 of the series.

What Comes Next: V2, V3, and Post-Training

V1 was the baseline. The series continues with increasingly sophisticated iterations:

• Part 2 — V2: Modernizing the architecture with GQA (Grouped-Query Attention), QK-Norm, EMA (Exponential Moving Average), Muon optimizer, Differential Attention, and mHC. Some changes earned their place, others didn't improve on the baseline and were dropped.
• Part 3 — Post-training: SFT (Supervised Fine-Tuning) and GRPO (Group Relative Policy Optimization), turning a base model into something actually useful. This is where the evaluation harness is introduced, putting the model against standardized benchmarks.
• Part 4 — V3: Scaling to 672M parameters, switching to XSA (Exclusive Self-Attention), training across 2 GPUs with DDP (Distributed Data Parallel), and serving the model so people can try it.

The full series is available on John's Substack [Link]. The code is on GitHub [Link], and the models are on HuggingFace: modern-llm-v1-base [Link], modern-llm-v1-sft [Link], modern-llm-v1-grpo [Link].

Key Takeaways

This series is a masterclass in transparent, from-scratch ML engineering. Here are the lessons that apply regardless of whether you're building a 350M model or a 350B one:

1. Build to understand. Reading about RoPE, attention, RMSNorm, and SwiGLU is not the same as making them work together and watching a loss curve come down. The bugs you fix teach you more than the papers you read.
2. Test before you train. Writing behavioral tests for your causal mask, shape assertions for every module, and loss sanity checks before spending money on GPUs saves you from silent failures that cost far more to debug.
3. Checkpoint properly. Atomic saves, corrupt-check fallbacks, and RNG state restoration are not edge cases — they're the difference between losing 1000 steps and losing 31 hours of training.
4. Hardware matters as much as code. The H100 SXM vs NVL mistake cost $35 for a 3-cent-per-hour difference. Always benchmark your hardware choices before committing to a long run.
5. Start cheap, scale up. Smoke-test on a cheap A100 before launching the expensive GPU. The pattern of "validate on cheap hardware, train on expensive hardware" applies at every scale.
6. Quality over volume. At small scales, data quality matters enormously. FineWeb-Edu over Common Crawl, even if it means fewer total tokens.

The point was never to compete with open-source models. The point was to build a language model end to end and understand how these things actually work. Measured against that goal: a clean architecture, a stable training run, a coherent base model, and a stack built from scratch — V1 did exactly what it was meant to.

References & Further Reading

• Original series: "Building a 350M Transformer From Scratch" by John [Link]
• GitHub repository: github.com/JohnEnev/modern-llm [Link]
• HuggingFace models: modern-llm-v1-base [Link], modern-llm-v1-sft [Link], modern-llm-v1-grpo [Link]
• Vaswani et al. (2017): "Attention Is All You Need" — the paper that started it all [Link]
• Su et al. (2021): "RoFormer: Enhanced Transformer with Rotary Position Embedding" [Link]
• Zhang & Sennrich (2019): "Root Mean Square Layer Normalization" [Link]
• Shazeer (2020): "GLU Variants Improve Transformer" [Link]
• Hoffmann et al. (2022): "Training Compute-Optimal Large Language Models" (Chinchilla) [Link]
• Dao et al. (2022): "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" [Link]
• Karpathy: "Let's build GPT" / nanoGPT series [Link]
• Raschka: "Build a Large Language Model (From Scratch)" [Link]
• Georgia Tech: Transformer Explainer — interactive attention visualization [Link]
• Penedo et al. (2024): "The FineWeb Datasets" [Link]