Building LLM, Part 1 — Training: Predicting the Next Token

Abstract. If you finished the Vision Transformer series, you already understand about 70% of a language model — the embeddings, the attention, the residual stack, the training loop are all the same machinery. This opening part marks the delta. We build a small GPT from first principles and train it, from scratch on a single laptop GPU, on a corpus of Indian folk tales — Jataka, Panchatantra-style fables, the tales of the Punjab and Bengal. Along the way we hit a real problem: the standard GPT-2 tokenizer shatters words like Brahman and Rajah, because it never learned them. So we build our own. By the end you’ll have watched a 9-million-parameter model go from pure noise to a model whose weights are ready to write — every number traced to a model you could train yourself in ten minutes. Setting it loose to actually write is Part 2 .

In this 8-part series (you are reading Part 1):

1. Part 1 — Training: predicting the next token — you are here
2. Part 2 — Inference: how the model writes
3. Part 3 — ViT vs LLM: what carries over, what’s new
4. Part 4 — Fine-tuning: teaching the model to follow instructions
5. Part 5 — RLHF: learning from human preference
6. Part 6 — PPO: the policy-gradient engine
7. Part 7 — GRPO: reinforcement without a value network
8. Part 8 — DPO: preference learning with no RL loop

This is Series 2. Series 1 built a Vision Transformer from pixels up; here we reach a language model by comparison. Wherever something is identical to the ViT, we’ll say so and move on — the point is to spend our attention on what’s genuinely new.

Why a Vision Transformer first?

It can feel backwards to study images before language models — after all, ChatGPT reads text, not pictures. But building the Vision Transformer first was a deliberate choice, and it’s the whole premise of this course.

A transformer doesn’t care whether you feed it patches of an image or words of a sentence. Its engine — multi-head attention, LayerNorm, residual connections, the MLP, and the training loop — is exactly the same either way. So the real question is never “vision or language?” It’s where that shared engine is easiest to learn. And that’s vision:

• A “token” is just an image patch — a fixed grid cut straight from the picture. No vocabulary, no tokenizer, no “unknown word” problem to solve before you even begin.
• No causal mask. A ViT sees the whole image at once, so you meet plain attention before adding language’s extra “don’t peek at the future” rule.
• You can literally see it work — the patches, the attention maps, what the model looks at. Pixels are concrete; text is abstract.
• It trains from scratch on a single GPU in minutes (CIFAR-10), so every figure is a number you produced yourself, not one you took on faith.

Here’s the payoff. By the end of Series 1 you had built the entire transformer — it was just wearing its “vision” clothes. That means you already understand roughly 70% of a language model. This series only has to teach the delta: how text becomes tokens (a lookup, not a computation), the causal mask, and predicting the next token. The genuinely hard, abstract part — language — is far easier to reason about once the machine itself is second nature.

The one idea: predict the next token

Everything a language model does at training time reduces to a single game: read the tokens so far, and predict the one that comes next.

Figure 1: Prediction of next token, given a earlier tokens

That’s it. Give the model “In a certain village there lived a poor Brahman” and it returns a probability for every possible next token — here leaning toward a full stop, a comma, or a continuing word like who or named, all of which are perfectly grammatical. No labels, no annotation: the “right answer” is simply the next token that was already there in the text. This is why language models can learn from raw text alone, and it is the whole training signal. The rest of this article is just the machinery that makes this one prediction possible — and most of that machinery you already met in the ViT.

In code, that idea is two lines: run the model, then normalise its scores into a distribution over the next token.

logits, _ = model(ids)                     # ids = the tokens so far, shape (1, T)
probs = F.softmax(logits[0, -1], dim=-1)   # one probability per vocabulary token

Everything below is how model is built and trained so those probabilities become good ones.

Text → tokens: why we had to build our own tokenizer

The ViT cut an image into patches. A language model cuts text into tokens — but where a patch was computed from pixels, a text token is looked up from a fixed vocabulary. The catch: that vocabulary has to be built in advance, and who builds it decides what your model finds easy.

Our first attempt just reused r50k_base, GPT-2's off-the-shelf tokenizer, whose vocabulary was frozen from ~2019 Western web text. Watch what it does to the words in our tales:

Figure 2: Limitation of Tokenizer and Custom Tokenizer

Brahman becomes Brah + man; Rajah becomes Raj + ah; Panchatantra shatters into five pieces. These aren't obscure words — Rajah appears 709 times in our corpus — but the tokenizer never saw them often enough in its training data to give them their own token. Fragmentation means longer sequences and a model that has to reassemble Brahman from Brah + man (where man is shared with woman and human). With only ~1M tokens of training data, that's a real handicap.

So we built our own tokenizer from the tales themselves, using byte-pair encoding (BPE) — the same algorithm GPT uses, just trained on the right corpus.

Those tales are twelve public-domain collections of Indian folk tales from Project Gutenberg: Jataka Tales and More Jataka Tales, Indian Fairy Tales (Jacobs) and Indian Fairy Tales (Stokes), Old Deccan Days, Folk-Tales of Bengal, The Talking Thrush, Tales of the Punjab, Twenty-Two Goblins, Folklore of the Santal Parganas, Santal Folk Tales, and Tales of the Sun (Southern India) — about 809,000 words in all.

Figure 3: Byte-pair Merging in Tokenizer training

BPE starts from the 256 raw bytes — so no text is ever “unknown” — and repeatedly merges the most frequent adjacent pair into a new token. Run it 7,935 times on our tales and the vocabulary reaches 8,192 tokens, now including Rajah, Brahman and jackal as whole units. One design choice worth being explicit about: we replaced GPT-2's tokenizer completely rather than extending it.

Figure 4: Custom Tokenizer with 8192 words

Extending r50k (bolting new tokens onto its 50,257) makes sense when you want to preserve a big pretrained model’s existing English ability — that’s what projects like OpenHathi do. But we’re training from scratch on a tiny narrow corpus, so r50k’s 50k Western-web vocabulary is pure dead weight (our corpus used only 37% of it). A fresh 8k vocab is smaller, 96.6% utilised, and tuned to our own words — a clean swap.

In code, training the tokenizer is a handful of lines with HuggingFace’s tokenizers:

import glob
from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from llm_transformer.prep_tales import strip_gutenberg, RAW_DIR

# read every raw .txt book and strip its Project Gutenberg header/footer,
# leaving just the story text - one big string per book
tales_texts = [strip_gutenberg(open(p, encoding="utf-8", errors="ignore").read())
               for p in sorted(glob.glob(f"{RAW_DIR}/*.txt"))]   # 12 books, ~809k words
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)   # start from bytes
trainer = trainers.BpeTrainer(vocab_size=8192, special_tokens=["<|endoftext|>"],
                              initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(tales_texts, trainer)   # learn the 7,935 merges
tok.save("tales_bpe.json")

Full script: llm_transformer/train_tokenizer.py.

The .json file is now our finished tokenizer — a frozen table of merges. From here on we never re-train it; we just load it and use it to turn text into ids (and back). That .encode() call is the bridge between raw text and the neural network:

tok = Tokenizer.from_file("tales_bpe.json")   # reload the trained tokenizer
ids = tok.encode("Once upon a time").ids      # -> [412, 987, 15, 733]   (integers)
idx = torch.tensor(ids).unsqueeze(0)          # (1, 4) — a batch of one sequence

vocab_size = tok.get_vocab_size()             # 8192 - sets the NUMBER OF ROWS below

Those integers are the only thing the model ever sees — the words are gone. Each id is just a row number, and vocab_size tells the embedding layer how many rows to keep.

Token → vector: a lookup, not a computation

Once text is a list of token ids, each id becomes a vector. This is the one spot where an LLM is simpler than the ViT:

Figure 5: Tokenizer Index mapping to embedding vector

In the ViT, each patch’s token was computed from pixels by a small matrix multiply. In an LLM the token is discrete, so its vector is literally the id-th row of a learned table — a pure lookup. Token 789 → row 789 → a 256-dimensional vector. Add a positional embedding (exactly as in the ViT) and we're done.

In code, the “table” is a single nn.Embedding layer and the lookup is just calling it — indexing, not multiplying:

self.token_embed = nn.Embedding(vocab_size, n_embed)   # the 8,192 x 256 table
# .weight IS a learnable (8192, 256) parameter tensor — 2.1M weights, trained by backprop;
# each id's row is nudged only on steps where that token appears in the batch
...
tok_emb = self.token_embed(idx)     # idx (B, T)  ->  vectors (B, T, 256), one row per id

Everything after this point — the attention, the residual stream, the MLP blocks — is identical to the Vision Transformer. With one exception.

The exception: the causal mask

The ViT let every patch look at every other patch. A language model can’t: it’s predicting the next token, so if position t were allowed to look at position t+1, it would simply read off the answer. The fix is a single line — mask attention so each token sees only itself and the past:

Figure 6: Triangular Causal mask

Before the softmax, every “future” position is set to −∞; softmax turns −∞ into a weight of exactly zero, so a blocked token contributes nothing. That lower-triangular mask is the entire architectural difference between the ViT’s attention and an LLM’s. The dot-products, the scaling by √d, the softmax, the weighted sum of values — all exactly as we built them in Series 1.

Get Chandrashekhar Padole’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

In code, it is one extra line on top of ordinary attention:

# tril = torch.tril(torch.ones(T, T))  — a lower-triangular matrix of 1s, made once
attn = q @ k.transpose(-2, -1) / math.sqrt(head_size)          # scores, (B, T, T)
attn = attn.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # blank out the future
attn = F.softmax(attn, dim=-1)                                  # -inf -> weight 0
out  = attn @ v

From src/models/attention.py.

One pass, every context at once

That one-line mask hides something worth pausing on. It’s tempting to picture a language model working the way it generates — one token at a time, appending the next to a growing context and re-running attention each step. That picture is correct for inference, but it is not how the model is trained. During training the whole sentence is already known, so — exactly as in the ViT — we compute Q, K, V for all T tokens in a single pass and form the full T×T attention matrix in one shot.

The causal mask is what makes that single pass legal. Look at row t of the masked matrix: everything to the right of the diagonal is −∞, so after softmax it’s zero. When we do the final out = attn @ V, row t's output is a weighted sum of only V₁…V_t — the first t tokens. So every context length is computed simultaneously:

Figure 7: Diffrent runs and shape of mask in training

• Row 1 attends to token 1 only → its output predicts token 2 from a 1-token context.
• Row 3 attends to tokens 1–3 → predicts token 4 from a 3-token context.
• Row 6 attends to all six → predicts token 7 from the full context.

Another way to see it: each row’s context is a growing lower-triangular matrix — a 1×1 for the first token, 2×2 for the first two, 3×3 for the first three, and so on. Those are exactly the nested top-left blocks of the big matrix (outlined in orange), so the one 6×6 pass already contains every smaller run inside it — the strip along the bottom of the figure just unrolls them. All six happen in the same attn @ V — there is no Python loop over positions. This is the whole reason transformers train efficiently: one sequence of length T is really T training examples (contexts of length 1, 2, …, T) scored in parallel, and the very same triangular mask that gives us those T examples is what guarantees, at generation time, that position t never peeks at the future. (Generation then walks these rows one at a time, reusing a KV-cache of past keys and values so each step stays cheap.)

Inside one run: what attn @ V actually computes

So far the cells have just been “attend / don’t attend.” Let’s open up the multiply and watch real numbers flow, for the same three runs — using the actual attention weights and value vectors from our trained model (block 0, head 0) on “Once there lived a clever jackal”:

Figure 8: Example of Q,K,V calculations across runs

Read one row at a time. The attention weights in that row are positive and sum to 1, so the output is a convex blend — a weighted average — of the value vectors the row is allowed to see:

• Run 1 — only one token is in view, its weight is 1.00, so out1 = V1. The very first position simply copies its own value; there is nothing yet to mix.
• Run 3 — out3 = 0.20·V1 + 0.06·V2 + 0.74·V3. The row leans hardest on its own value (0.74 on lived) but still pulls in a fifth of Once — a genuine mixture.
• Run 6 — out6 spreads weight across all six values (0.26 each on there and jackal). That single blended 32-dimensional vector is the head's summary of the whole prefix, and it's what the MLP and later layers turn into a prediction.

This is the entire content of attention: the query–key scores decide the weights, and those weights blend the values. The masking we just saw guarantees the blend never reaches past the current token — so every row is a legitimate “given what I’ve read so far” summary. (A real head has 32 value dimensions and the model has eight heads per layer across six layers; the strips show the first 8 dimensions of one head, but the mechanism is identical everywhere.)

One pass, many lessons: T predictions, T losses

Here’s a payoff of next-token prediction that the ViT never had. The ViT produced one prediction per image (its class). A language model, thanks to the causal mask, predicts the next token at every position simultaneously — so a single sequence yields as many training signals as it has tokens:

Figure 9: Average of Multiple T predictions and losses calculation for T tokens input for all runs

Each position gets a cross-entropy loss — how surprised the model was by the token that actually came next. Some are easy (predicting a after "lived" costs just 0.58); some are genuinely hard (predicting the adjective clever costs 7.17 — there were many plausible words). Average the six and you get 3.59, and that single number is what backpropagation — the exact chain-rule machinery from Series 1, Part 6 — pushes down.

In code, all T losses are a single cross_entropy call: flatten the (B, T) grid of positions into one long list and score every one against its true next token.

logits = self.lm_head(x)                       # (B, T, vocab)
B, T, C = logits.shape
loss = F.cross_entropy(logits.reshape(B*T, C), # every position ...
                       targets.reshape(B*T))   # ... vs the token that actually followed

From src/models/transformer.py.

Each loss is a vocabulary classification (just like the ViT, but 8,192 classes)

It’s worth being precise about what that cross_entropy compares, because the prediction and the target are not single words — they're vectors. At each position the lm_head produces a logits vector of length 8,192 (one score per vocabulary token); softmax turns it into a probability distribution over all 8,192 tokens. The target is a one-hot vector, also length 8,192 — a 1 at the word that actually came next, 0 everywhere else. Cross-entropy compares those two vectors. In other words, next-token prediction is a classification, and the classes are the vocabulary words:

Figure 10: Diffrence in Logits for LLM and ViT

One subtlety the figure makes explicit: those bars are sorted by probability for readability, but the raw output is a fixed-layout vector — slot i is always token id i. “son” lives at index 498 in every forward pass, training and inference alike; softmax only changes the heights on those fixed slots, never their order. That’s why the one-hot target’s single 1 sits at index 498, not at some "first" position — the position is the token's identity, assigned once by the tokenizer.

This is precisely the ViT’s classifier head — a softmax over class scores, scored by cross-entropy against a one-hot label. The only differences are the number of classes (8,192 vocabulary tokens vs. the ViT’s ~10 image labels) and how often it runs (at every one of the T positions vs. once per image). And there’s a neat simplification: against a one-hot target, cross-entropy collapses to −log(the probability the model gave the correct word). That’s exactly why every per-position loss above is a “how surprised was I?” number — son at 9.8% costs 2.32; a word the model was sure of would cost far less, a genuine surprise far more.

How a batch is built — and how its one loss is formed

That reshape(B*T) quietly does two things worth unpacking: it says a batch has B sequences, and each contributes T position-losses. So where do those B sequences come from, and are they all the same length?

Here’s a detail that trips people up: in pretraining, the sequences are not padded — they’re all exactly the context length. We concatenate the entire corpus into one long token stream (documents separated by an <|endoftext|> token) and cut fixed-length windows out of it at random offsets. Because every window is a slice of the same continuous stream, none is ever short, so there is no padding token at all:

Figure 11: Training Corpus to batches and Loss calculation

The only special token in the stream is <|endoftext|> (EOT), and it's a separator, not padding — it marks where one folk tale ends and the next begins. A sampled window can land across a boundary and contain an EOT in the middle, which is exactly how the model learns "a story just ended, start fresh." (Padding does appear later, in Part 4's fine-tuning, where examples are genuinely different lengths and the padded positions are masked out of the loss — a different setting from this packed pretraining.)

Now the loss. This is the second place the LLM differs from the ViT, and it’s a happy one. The ViT made one prediction per image, so it had one loss. Our LLM makes a prediction at every position, so a single 128-token sequence yields 128 losses — one per position — which are averaged into that sample’s loss (3.78 for sample 1 above). Then the 32 sample losses in the batch are averaged into the single batch loss (3.92) that backprop actually minimizes. That's two levels of averaging, and it means one training step pushes down the average of 32 × 128 = 4,096 next-token losses at once — a lot more signal per step than the ViT's one-loss-per-image.

Training: watching the loss come down

Wire that loss to an optimizer and repeat. Here is the real training run of our model on the tales corpus:

Figure 12: Training & Validation (held-out) Loss with respect to epochs

It starts at 9.2 — which is exactly ln(8192), the loss of pure guessing across an 8,192-token vocabulary — and falls to a best held-out loss of 4.73. Notice the gap between the two curves stays small (~1.0): because the domain is narrow, the model generalises rather than just memorising. And the model doing this learning is tiny:

Figure 13: Training Flow and corresponding count/extent of learning parameters (weights)

9.0 million parameters — small enough to train in about ten minutes on a 6 GB laptop GPU. Just over half live in the six transformer blocks; the rest are the embedding table and the output head. “Training” is nothing more than nudging these 9 million numbers, over and over, so that the next-token probabilities put more weight on the token that actually came next.

In code, the loop is exactly what you’d guess — and identical in shape to the ViT’s training loop from Series 1:

for step in range(train_steps):
    xb, yb = get_batch(train_data)     # random (context, next-token) windows
    _, loss = model(xb, yb)            # forward: T predictions -> one mean loss
    opt.zero_grad(); loss.backward()   # backprop the loss into all 9M parameters
    opt.step()                         # nudge them downhill

Full script (with the cosine schedule and eval): llm_transformer/train_tales.py.

What we built — and what’s next

We now have a trained language model. Tally the delta from the Vision Transformer of Series 1 and it’s remarkably small:

• New in the LLM: a lookup-table token embedding, a causal mask, and next-token loss at every position — a classification over the whole vocabulary. Plus a practical lesson: the tokenizer is part of the model, and building your own for your domain is worth it.
• Unchanged from Series 1: attention, positional embeddings, the residual stream, LayerNorm, the MLP blocks, and the entire backprop-and-optimizer training loop.

Along the way we opened up the mechanics with real numbers: how the causal mask lets one forward pass score every context length at once; what attn @ V actually computes — each output row is a weighted blend of the value vectors a token can see; how a training batch is really built by packing, not padding; and how a single position's loss is a vocabulary-sized classification whose per-position numbers average into one.

The weights are set — but a trained model just sits there until you run it. So in Part 2 we turn to inference: how the model actually writes, the generation loop, the KV-cache that keeps it fast, and what our little model produces when we hand it the opening of a tale.