If you want to understand how modern Large Language Models work, reading about them is useful — but building a small one yourself makes the concepts much more concrete.
In this project, we'll build a GPT-2-style decoder-only Transformer from scratch in PyTorch, train it on the TinyStories dataset, and use it to generate text.
This isn't an attempt to build a model at the scale of ChatGPT, Llama, or Claude. Instead, the goal is to take the core ideas behind these models and implement a much smaller version that you can actually train, inspect, and experiment with yourself.
We'll go from tokenization and embeddings all the way through self-attention, Transformer blocks, training, and text generation.
You have to build one from the ground up.
In 2017, researchers published the landmark paper "Attention Is All You Need". It introduced the Transformer architecture, originally designed with an Encoder-Decoder structure for tasks like machine translation.
But shortly after, researchers working on the GPT (Generative Pre-trained Transformer) series realized something profound:
If your only goal is to predict the next token in a sequence, you don't need the Encoder.
You just need the right half of the architecture.
You need a Decoder-only Transformer.
Many modern autoregressive LLMs use a decoder-only Transformer architecture.
In this post, we're going to strip away the magic.
We'll build and understand a 124M-parameter GPT-2-style decoder-only Transformer in PyTorch, focusing on:
- Masked self-attention
- Query, Key, and Value projections
- Multi-head attention
- Transformer blocks
- Feed-forward networks
- Positional embeddings
- Context windows
- GPT-2-style weight initialization
- Weight tying
- Autoregressive text generation
- Temperature
- Top-K sampling
- Stop tokens
We'll also connect the settings you see in an LLM UI directly to the Python code and mathematics underneath.
1. The Heart of the Model: Masked Self-Attention
If you take away only one concept from this article, make it this:
Self-attention is how a Transformer gives tokens context.
In a self-attention layer, the model compares every token in a sequence with other tokens to determine which ones are relevant to each other.
Consider:
"The bank of the river."
The word bank has a different meaning from:
"The bank on Main Street."
The surrounding tokens provide the context necessary to understand which meaning is relevant.
Attention allows the model to mathematically capture these relationships.
Query, Key, and Value
For every token, the input is projected into three different representations:
Query ( )
The Query represents what information the current token is looking for.
Key ( )
The Key represents what information a token contains and can be matched against.
Value ( )
The Value represents the actual information that gets passed forward when a token is considered relevant.
Conceptually:
Query → What am I looking for?
Key → What information do I contain?
Value → What information should I pass along?
The model computes similarities between Queries and Keys.
For example, if the Query associated with "bank" strongly matches the Key associated with "river", the corresponding Value from "river" contributes more strongly to the representation of "bank".
The Attention Equation
The core attention operation is:
There are several steps hidden inside this equation.
Step 1 — Compute
The Query matrix is multiplied by the transpose of the Key matrix.
This produces attention scores describing how strongly tokens relate to one another.
Step 2 — Scale the scores
The scores are divided by:
where is the dimensionality of each attention head.
Without this scaling, the dot products can become very large as the dimensionality increases.
That can cause the softmax function to produce extremely sharp probability distributions and make optimization more difficult.
Step 3 — Apply Softmax
Softmax converts the attention scores into normalized weights.
These weights determine how much attention each token gives to other tokens.
Step 4 — Multiply by
Finally, the attention weights are multiplied by the Value matrix.
The result is a context-aware representation of every token.
2. Why Decoder-Only Models Need a Causal Mask
A decoder-only LLM is autoregressive.
Its job is to predict the next token.
For example:
The capital of
should be used to predict:
India
During training, the model must not be allowed to see the answer ahead of time.
If the sequence is:
The capital of India is
then the model should only be allowed to attend to tokens at or before the current position.
Conceptually, the attention matrix looks like this:
The capital of India is
The ✓
capital ✓ ✓
of ✓ ✓ ✓
India ✓ ✓ ✓ ✓
is ✓ ✓ ✓ ✓ ✓
Future tokens are masked.
This is called causal self-attention.
The token at position can only attend to positions:
3. Implementing Causal Self-Attention in PyTorch
Modern versions of PyTorch provide:
F.scaled_dot_product_attention()
This function implements the core scaled dot-product attention operation and can use highly optimized attention kernels depending on the hardware and PyTorch configuration.
Here's the implementation:
import torch
import torch.nn as nn
from torch.nn import functional as F
class CausalSelfAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
# Project inputs to Q, K, and V
# in a single linear transformation.
self.c_attn = nn.Linear(
embed_dim,
3 * embed_dim,
bias=False
)
# Final projection after attention
self.c_proj = nn.Linear(
embed_dim,
embed_dim,
bias=False
)
self.num_heads = num_heads
self.embed_dim = embed_dim
def forward(self, x):
B, T, C = x.size()
# Calculate Q, K, and V
qkv = self.c_attn(x)
q, k, v = qkv.split(
self.embed_dim,
dim=2
)
# Dimension of each attention head
head_dim = C // self.num_heads
# Reshape:
#
# (B, T, C)
#
# into:
#
# (B, num_heads, T, head_dim)
q = q.view(
B,
T,
self.num_heads,
head_dim
).transpose(1, 2)
k = k.view(
B,
T,
self.num_heads,
head_dim
).transpose(1, 2)
v = v.view(
B,
T,
self.num_heads,
head_dim
).transpose(1, 2)
# Causal self-attention
#
# is_causal=True prevents the model
# from attending to future tokens.
y = F.scaled_dot_product_attention(
q,
k,
v,
is_causal=True
)
# Restore original shape:
#
# (B, num_heads, T, head_dim)
#
# → (B, T, C)
y = (
y.transpose(1, 2)
.contiguous()
.view(B, T, C)
)
return self.c_proj(y)
4. What Happens Inside Multi-Head Attention?
Instead of performing one large attention operation, Transformers split the embedding dimension into multiple smaller attention heads.
For example:
Embedding dimension = 768
Number of heads = 12
Therefore:
Each head operates on a 64-dimensional representation.
Conceptually:
Input
|
+-----------+-----------+
| | |
Head 1 Head 2 Head 3 ... Head 12
| | |
Attention Attention Attention
| | |
+-----------+-----------+
|
Concatenate
|
Linear Projection
|
Output
Different heads can learn different relationships.
One head might focus heavily on nearby syntactic relationships while another can learn longer-range dependencies.
The important idea is that multi-head attention is still just multiple independent Q/K/V transformations followed by attention, concatenation, and a final projection.
5. The Transformer Block
Attention alone isn't the entire Transformer.
A Transformer block typically combines:
- Layer Normalization
- Multi-head causal self-attention
- Residual connection
- Layer Normalization
- Feed-forward network
- Residual connection
A simplified representation is:
Input
|
LayerNorm
|
Causal Self-Attention
|
+<-------------+
| |
v |
Residual |
| |
LayerNorm |
| |
Feed-Forward |
Network |
| |
+<-------------+
|
Output
Here's the implementation:
class TransformerBlock(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
self.ln_1 = nn.LayerNorm(embed_dim)
self.attn = CausalSelfAttention(
embed_dim,
num_heads
)
self.ln_2 = nn.LayerNorm(embed_dim)
self.ffwd = nn.Sequential(
nn.Linear(
embed_dim,
4 * embed_dim,
bias=False
),
nn.GELU(),
nn.Linear(
4 * embed_dim,
embed_dim,
bias=False
),
nn.Dropout(0.1)
)
def forward(self, x):
# Attention + residual connection
x = x + self.attn(
self.ln_1(x)
)
# Feed-forward network + residual connection
x = x + self.ffwd(
self.ln_2(x)
)
return x
6. What Does the Feed-Forward Network Do?
The attention mechanism allows tokens to exchange information.
The feed-forward network then performs additional nonlinear transformations independently for each token position.
The basic structure is:
In our implementation, the hidden dimension is four times the embedding dimension:
768 → 3072 → 768
So if:
embed_dim = 768
the feed-forward network becomes:
nn.Linear(768, 3072)
nn.GELU()
nn.Linear(3072, 768)
This gives the model additional capacity to transform the contextual representations produced by attention.
7. Building the Full MiniLLM
Now we can assemble the individual pieces.
A GPT-2-style 124M-parameter configuration uses approximately:
Vocabulary size = 50,257
Embedding size = 768
Attention heads = 12
Transformer layers = 12
Our model will contain:
- Token embeddings
- Positional embeddings
- 12 Transformer blocks
- Final LayerNorm
- Language-model head
Token Embeddings
The input to the Transformer is not raw text.
Text is first converted into tokens.
Each token is mapped to an integer:
"The capital of India"
might become something conceptually similar to:
[464, 3139, 286, 3182]
The embedding layer then converts each token ID into a dense vector.
self.token_embedding = nn.Embedding(
vocab_size,
embed_dim
)
8. Context Window and Positional Embeddings
When a model says it supports a:
128K Context Window
it means the model can process up to approximately:
128,000 tokens
within a sequence.
But how does a Transformer know where each token occurs?
Attention itself does not inherently encode sequential order.
Therefore, we need positional information.
With learned absolute positional embeddings, the model has a position embedding for every possible position:
Position 0
Position 1
Position 2
...
Position 511
If:
max_seq_len = 512
then the positional embedding matrix contains 512 positions.
Conceptually:
Token embedding
+
Position embedding
|
v
Transformer blocks
For example:
Token: "capital"
Position: 2
Token Embedding
+
Position Embedding
=
Contextual Input Representation
9. Full MiniLLM Implementation
import math
import torch
import torch.nn as nn
class MiniLLM(nn.Module):
def __init__(
self,
vocab_size=50257,
max_seq_len=512,
embed_dim=768,
num_heads=12,
num_layers=12
):
super().__init__()
self.max_seq_len = max_seq_len
# Token embeddings
self.token_embedding = nn.Embedding(
vocab_size,
embed_dim
)
# Positional embeddings
self.position_embedding = nn.Embedding(
max_seq_len,
embed_dim
)
# Transformer blocks
self.blocks = nn.Sequential(
*[
TransformerBlock(
embed_dim,
num_heads
)
for _ in range(num_layers)
]
)
# Final LayerNorm
self.ln_f = nn.LayerNorm(
embed_dim
)
# Language-model head
self.lm_head = nn.Linear(
embed_dim,
vocab_size,
bias=False
)
# Weight tying
#
# Use the same weights for token embeddings
# and the final vocabulary projection.
self.token_embedding.weight = (
self.lm_head.weight
)
self.apply(self._init_weights)
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(
module.weight,
mean=0.0,
std=0.02
)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(
module.weight,
mean=0.0,
std=0.02
)
# GPT-2-style residual projection scaling
for pn, p in self.named_parameters():
if pn.endswith("c_proj.weight"):
torch.nn.init.normal_(
p,
mean=0.0,
std=0.02 / math.sqrt(
2 * 12
)
)
def forward(self, idx):
B, T = idx.size()
# Position indices:
#
# [0, 1, 2, ..., T-1]
pos = torch.arange(
0,
T,
dtype=torch.long,
device=idx.device
)
# Combine token and positional information
x = (
self.token_embedding(idx)
+ self.position_embedding(pos)
)
# Transformer blocks
x = self.blocks(x)
# Final normalization + vocabulary projection
logits = self.lm_head(
self.ln_f(x)
)
return logits
The most important line here is:
x = self.token_embedding(idx) + self.position_embedding(pos)
This combines:
What is the token?
with:
Where is the token?
10. Weight Tying
Notice this:
self.token_embedding.weight = self.lm_head.weight
The token embedding layer and language-model head use the same weight matrix.
Without weight tying, we'd have two separate matrices:
Token Embedding
50257 × 768
Language Model Head
768 × 50257
These matrices contain a huge number of parameters.
Weight tying allows both operations to share the same parameters.
Conceptually:
Token IDs
|
v
Embedding Matrix
|
Transformer
|
v
Hidden Representation
|
Same Matrix
|
v
Vocabulary Logits
11. GPT-2-Style Weight Initialization
Initialization becomes increasingly important as the network gets deeper.
If the activations become too large as they pass through the Transformer blocks, the model can become unstable.
The code therefore scales certain projection weights by:
where represents the number of Transformer layers.
For 12 layers:
This helps control the variance of the residual stream as information passes through many Transformer blocks.
12. Training the 124M-Parameter Model
Up to this point, we have built the architecture.
But there is an important distinction:
A Transformer architecture is not automatically an intelligent language model.
At initialization, the weights are essentially random.
The model has the machinery to perform:
Token → Embedding → Attention → Feed-Forward → Logits
but it has not yet learned what those transformations should represent.
Hardware Used
The model was trained using Google Colab with an NVIDIA A100 80 GB GPU.
| Component | Configuration |
|---|---|
| Platform | Google Colab |
| GPU | NVIDIA A100 80 GB |
| System RAM | ~167 GB |
| Model | ~124M parameters |
| Batch Size | 64 |
| Sequence Length | 512 |
The A100 was used because the training configuration combines a relatively large batch size, a 512-token context window, and a 124M-parameter Transformer.
The architecture itself is not specific to the A100. It can be trained on other CUDA-capable GPUs by adjusting the batch size and other training parameters according to the available memory.
Training is the process of repeatedly showing the model sequences of text and updating its parameters so that its prediction of the next token becomes better.
The complete training pipeline looks like this:
TinyStories Dataset
|
v
Text
|
v
GPT-2 BPE
Tokenization
|
v
Token IDs
|
v
Training Sequences
|
v
+---------------------+
| Decoder Transformer |
+---------------------+
|
v
Logits
|
v
Cross-Entropy Loss
|
v
Backpropagation
|
v
AdamW
|
v
Updated Parameters
|
+------------------+
|
v
Next Training Step
Let's walk through every stage.
13. The Dataset: TinyStories
For this experiment, we use Microsoft's TinyStories dataset.
The dataset contains short synthetic stories written using relatively simple vocabulary and grammar.
That makes it particularly useful for a small language model experiment because the model does not need to learn the enormous vocabulary and linguistic complexity required for a general-purpose web-scale LLM.
We load the dataset using Hugging Face Datasets:
from datasets import load_dataset
dataset = load_dataset(
"roneneldan/TinyStories",
split="train"
)
Instead of training on the entire dataset, the implementation uses a subset:
full_text = "\n<|endoftext|>\n".join(
dataset["text"][:250000]
)
The special token:
<|endoftext|>
is inserted between stories.
This is important.
Without an explicit boundary, the model could see:
Story A
Story B
as if they were one continuous story.
With an end-of-text token, the model gets an explicit signal:
Story A
<|endoftext|>
Story B
<|endoftext|>
Story C
During generation, this token can also tell us that the model believes the current story is finished.
14. Tokenization with GPT-2 BPE
The model does not directly receive strings.
It receives integers.
We use the GPT-2 tokenizer provided by tiktoken:
import tiktoken
enc = tiktoken.get_encoding("gpt2")
vocab_size = enc.n_vocab
The vocabulary contains:
50,257 tokens
A sentence such as:
Once upon a time
is converted into token IDs.
Conceptually:
"Once upon a time"
|
v
[7454, 2402, 257, 640]
The exact token IDs depend on the tokenizer.
The important point is that the neural network operates on those integer IDs rather than raw characters.
15. Turning the Dataset into One Long Token Stream
The implementation combines the selected stories into one long sequence:
full_text = "\n<|endoftext|>\n".join(
dataset["text"][:250000]
)
data = torch.tensor(
enc.encode(
full_text,
allowed_special={"<|endoftext|>"}
),
dtype=torch.long
)
The result is conceptually:
[token_1, token_2, token_3, ..., token_N]
We then split the token stream into training and validation data:
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
So approximately:
90% → training
10% → validation
The training set is used to update the weights.
The validation set is used to measure how well the model performs on tokens it did not directly train on.
16. Creating Training Batches
The Transformer expects sequences of a fixed maximum length.
Our configuration uses:
max_seq_len = 512
Therefore, each training example contains 512 input tokens.
The key idea is that the target is simply the input shifted by one token.
Suppose the sequence is:
The cat sat on the mat
The training pair becomes:
Input:
The cat sat on the mat
Target:
cat sat on the mat <next>
More formally:
Input: [t0, t1, t2, t3, t4]
Target: [t1, t2, t3, t4, t5]
The model is therefore trained to answer:
Given everything up to this position, what is the next token?
The batch function implements exactly this:
def get_batch(split):
data_split = (
train_data
if split == "train"
else val_data
)
ix = torch.randint(
len(data_split) - max_seq_len,
(batch_size,)
)
x = torch.stack([
data_split[i : i + max_seq_len]
for i in ix
])
y = torch.stack([
data_split[
i + 1 :
i + max_seq_len + 1
]
for i in ix
])
return x.to(device), y.to(device)
If:
batch_size = 64
max_seq_len = 512
then one batch contains:
input tokens.
That means every training step performs next-token prediction over 32,768 token positions.
17. The Training Objective: Next-Token Prediction
This is the central idea behind decoder-only language-model pre-training.
Suppose the input is:
Once upon a time
The model predicts:
there
Then:
Once upon a time there
predicts:
was
Then:
Once upon a time there was
predicts:
a
And so on.
During training, however, we do not have to generate one token at a time.
The causal mask allows us to compute all of these predictions in parallel.
For:
The cat sat
the model effectively learns:
The → cat
The cat → sat
The cat sat → ...
in a single forward pass.
This is one of the most important differences between training and generation.
Training
Many next-token predictions
↓
Computed in parallel
Generation
Predict one token
↓
Append it
↓
Predict another token
↓
Append it
↓
Repeat
18. Cross-Entropy Loss
The model outputs logits for every vocabulary token.
If:
vocab_size = 50,257
then for every token position the model produces:
50,257 logits
The correct next token has a target ID.
Cross-entropy measures how much probability the model assigned to the correct answer.
The loss is conceptually:
where:
- is the input context
- is the correct next token
- is the probability assigned to that token
If the correct token receives high probability, the loss is small.
If the correct token receives very low probability, the loss is large.
The implementation is:
loss = F.cross_entropy(
logits.view(B * T, -1),
targets.view(B * T)
)
Suppose:
B = 64
T = 512
Then the logits are reshaped from:
(64, 512, 50257)
into:
(32768, 50257)
The targets are reshaped from:
(64, 512)
into:
(32768)
So the loss is calculated over all 32,768 token predictions in the batch.
19. Backpropagation
Once the loss has been calculated, we need to determine how every parameter contributed to that error.
That's where backpropagation comes in.
The code is:
optimizer.zero_grad(
set_to_none=True
)
loss.backward()
loss.backward() computes gradients for the model parameters.
Conceptually:
Loss
|
v
∂Loss/∂W
|
v
Gradients
|
v
Every trainable parameter
The gradient tells the optimizer which direction each parameter should move to reduce the loss.
20. AdamW Updates the Parameters
We use AdamW:
optimizer = torch.optim.AdamW(
model.parameters(),
lr=learning_rate,
fused=True
)
The optimizer receives the gradients and updates the model's parameters.
Conceptually:
where:
- = model parameters
- = learning rate
- = gradient of the loss
AdamW adds adaptive moment estimates and decoupled weight decay to this basic optimization process.
21. Gradient Clipping
Large gradients can sometimes destabilize training.
The training loop therefore clips the gradient norm:
torch.nn.utils.clip_grad_norm_(
model.parameters(),
1.0
)
The maximum gradient norm is:
1.0
If the gradient norm exceeds that threshold, the gradients are scaled down.
The goal is not to remove the learning signal, but to prevent unusually large updates from destabilizing training.
22. Learning-Rate Warmup and Cosine Decay
The learning rate isn't kept constant throughout training.
The implementation starts with a warmup:
warmup_iters = 500
During warmup, the learning rate increases gradually:
if it < warmup_iters:
return learning_rate * (
it + 1
) / warmup_iters
The maximum learning rate is:
6e-4
After warmup, the learning rate follows cosine decay.
The schedule is:
where represents the normalized position within the decay period.
The implementation is:
def get_lr(it):
# 1. Linear warmup
if it < warmup_iters:
return (
learning_rate
* (it + 1)
/ warmup_iters
)
# 2. Minimum learning rate
if it > lr_decay_iters:
return min_lr
# 3. Cosine decay
decay_ratio = (
it - warmup_iters
) / (
lr_decay_iters
- warmup_iters
)
coeff = 0.5 * (
1.0
+ math.cos(
math.pi * decay_ratio
)
)
return (
min_lr
+ coeff
* (
learning_rate
- min_lr
)
)
23. The Complete Training Script
Here is the training script used to train the model:
import torch
import torch.nn as nn
from torch.nn import functional as F
import tiktoken
from datasets import load_dataset
import time
import math
# ==========================================
# 1. Hyperparameters
# ==========================================
batch_size = 64
max_seq_len = 512
embed_dim = 768
num_heads = 12
num_layers = 12
learning_rate = 6e-4
min_lr = 6e-5
max_iters = 10000
warmup_iters = 500
lr_decay_iters = 10000
eval_interval = 500
device = "cuda"
torch.set_float32_matmul_precision("high")
# ==========================================
# 2. Tokenizer and Dataset
# ==========================================
print("Loading tokenizer and TinyStories...")
enc = tiktoken.get_encoding("gpt2")
vocab_size = enc.n_vocab
dataset = load_dataset(
"roneneldan/TinyStories",
split="train"
)
print("Tokenizing data...")
full_text = "\n<|endoftext|>\n".join(
dataset["text"][:250000]
)
data = torch.tensor(
enc.encode(
full_text,
allowed_special={"<|endoftext|>"}
),
dtype=torch.long
)
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]
def get_batch(split):
data_split = (
train_data
if split == "train"
else val_data
)
ix = torch.randint(
len(data_split) - max_seq_len,
(batch_size,)
)
x = torch.stack([
data_split[
i : i + max_seq_len
]
for i in ix
])
y = torch.stack([
data_split[
i + 1 :
i + max_seq_len + 1
]
for i in ix
])
return x.to(device), y.to(device)
# ==========================================
# 3. Model Architecture
# ==========================================
class CausalSelfAttention(nn.Module):
def __init__(self):
super().__init__()
self.c_attn = nn.Linear(
embed_dim,
3 * embed_dim,
bias=False
)
self.c_proj = nn.Linear(
embed_dim,
embed_dim,
bias=False
)
self.num_heads = num_heads
def forward(self, x):
B, T, C = x.size()
qkv = self.c_attn(x)
q, k, v = qkv.split(
embed_dim,
dim=2
)
head_dim = (
C // self.num_heads
)
q = q.view(
B,
T,
self.num_heads,
head_dim
).transpose(1, 2)
k = k.view(
B,
T,
self.num_heads,
head_dim
).transpose(1, 2)
v = v.view(
B,
T,
self.num_heads,
head_dim
).transpose(1, 2)
y = F.scaled_dot_product_attention(
q,
k,
v,
is_causal=True
)
y = (
y.transpose(1, 2)
.contiguous()
.view(B, T, C)
)
return self.c_proj(y)
class TransformerBlock(nn.Module):
def __init__(self):
super().__init__()
self.ln_1 = nn.LayerNorm(
embed_dim
)
self.attn = CausalSelfAttention()
self.ln_2 = nn.LayerNorm(
embed_dim
)
self.ffwd = nn.Sequential(
nn.Linear(
embed_dim,
4 * embed_dim,
bias=False
),
nn.GELU(),
nn.Linear(
4 * embed_dim,
embed_dim,
bias=False
),
nn.Dropout(0.1)
)
def forward(self, x):
x = x + self.attn(
self.ln_1(x)
)
x = x + self.ffwd(
self.ln_2(x)
)
return x
class MiniLLM(nn.Module):
def __init__(self):
super().__init__()
self.token_embedding = nn.Embedding(
vocab_size,
embed_dim
)
self.position_embedding = nn.Embedding(
max_seq_len,
embed_dim
)
self.blocks = nn.Sequential(
*[
TransformerBlock()
for _ in range(num_layers)
]
)
self.ln_f = nn.LayerNorm(
embed_dim
)
self.lm_head = nn.Linear(
embed_dim,
vocab_size,
bias=False
)
# Weight tying
self.token_embedding.weight = (
self.lm_head.weight
)
def forward(
self,
idx,
targets=None
):
B, T = idx.size()
pos = torch.arange(
0,
T,
dtype=torch.long,
device=idx.device
)
x = (
self.token_embedding(idx)
+ self.position_embedding(pos)
)
x = self.blocks(x)
logits = self.lm_head(
self.ln_f(x)
)
loss = None
if targets is not None:
loss = F.cross_entropy(
logits.view(
B * T,
-1
),
targets.view(
B * T
)
)
return logits, loss
# ==========================================
# 4. Generation
# ==========================================
@torch.no_grad()
def generate(
model,
idx,
max_new_tokens,
temperature=0.7,
top_k=40
):
was_training = model.training
model.eval()
for _ in range(max_new_tokens):
idx_cond = idx[
:,
-max_seq_len:
]
logits, _ = model(
idx_cond
)
logits = (
logits[:, -1, :]
/ temperature
)
if top_k is not None:
v, _ = torch.topk(
logits,
min(
top_k,
logits.size(-1)
)
)
logits[
logits < v[:, [-1]]
] = -float("Inf")
probs = F.softmax(
logits,
dim=-1
)
idx_next = torch.multinomial(
probs,
num_samples=1
)
idx = torch.cat(
(idx, idx_next),
dim=1
)
if was_training:
model.train()
return idx
# ==========================================
# 5. Learning-Rate Schedule
# ==========================================
def get_lr(it):
if it < warmup_iters:
return (
learning_rate
* (it + 1)
/ warmup_iters
)
if it > lr_decay_iters:
return min_lr
decay_ratio = (
it - warmup_iters
) / (
lr_decay_iters
- warmup_iters
)
coeff = 0.5 * (
1.0
+ math.cos(
math.pi * decay_ratio
)
)
return (
min_lr
+ coeff
* (
learning_rate
- min_lr
)
)
# ==========================================
# 6. Create Model and Optimizer
# ==========================================
model = MiniLLM().to(device)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=learning_rate,
fused=True
)
print(
"Model parameters:",
f"{sum(p.numel() for p in model.parameters()) / 1e6:.2f}M"
)
print("Starting training...")
# ==========================================
# 7. Training Loop
# ==========================================
t0 = time.time()
for iter in range(max_iters):
lr = get_lr(iter)
for param_group in optimizer.param_groups:
param_group["lr"] = lr
# --------------------------
# Validation
# --------------------------
if (
iter % eval_interval == 0
or iter == max_iters - 1
):
model.eval()
with torch.no_grad():
with torch.autocast(
device_type="cuda",
dtype=torch.bfloat16
):
x_val, y_val = get_batch(
"val"
)
_, val_loss = model(
x_val,
y_val
)
print(
f"Step {iter}: "
f"Validation Loss = "
f"{val_loss.item():.4f} | "
f"LR = {lr:.4e}"
)
model.train()
# --------------------------
# Training batch
# --------------------------
xb, yb = get_batch(
"train"
)
with torch.autocast(
device_type="cuda",
dtype=torch.bfloat16
):
logits, loss = model(
xb,
yb
)
# --------------------------
# Backpropagation
# --------------------------
optimizer.zero_grad(
set_to_none=True
)
loss.backward()
# --------------------------
# Gradient clipping
# --------------------------
torch.nn.utils.clip_grad_norm_(
model.parameters(),
1.0
)
# --------------------------
# Parameter update
# --------------------------
optimizer.step()
# --------------------------
# Logging
# --------------------------
if iter % 100 == 0:
t1 = time.time()
print(
f"Iter {iter} | "
f"Loss: {loss.item():.4f} | "
f"Time/100 steps: "
f"{t1 - t0:.2f}s | "
f"LR: {lr:.4e}"
)
t0 = time.time()
# ==========================================
# 8. Save the Model
# ==========================================
torch.save(
model.state_dict(),
"tiny_stories_gpt2.pth"
)
print(
"\nModel saved to "
"'tiny_stories_gpt2.pth'"
)
# ==========================================
# 9. Generate a Sample
# ==========================================
prompt = (
"Once upon a time, "
"there was a little dog named"
)
context = torch.tensor(
[
enc.encode(prompt)
],
dtype=torch.long,
device=device
)
generated_indices = generate(
model,
context,
max_new_tokens=250,
temperature=0.7,
top_k=40
)
print(
enc.decode(
generated_indices[0].tolist()
)
)
25. The Training Loop in One Picture
The entire training process can now be visualized as:
TinyStories
|
v
Tokenization
|
v
Token ID Stream
|
v
Random 512-token
windows
|
v
+---------------+
| Input Tokens |
+---------------+
|
v
+-----------------------+
| Decoder Transformer |
| |
| Embeddings |
| ↓ |
| Attention |
| ↓ |
| Feed Forward |
| ↓ |
| ... |
+-----------------------+
|
v
Logits
|
v
Cross Entropy Loss
|
v
Backpropagation
|
v
AdamW
|
v
Updated Parameters
|
+------------------+
|
v
Next Batch
This process runs thousands of times.
26. Saving the Trained Model
After training, the model's learned parameters are saved:
torch.save(
model.state_dict(),
"tiny_stories_gpt2.pth"
)
This file contains the learned weights.
We can later create the same architecture and load those parameters without retraining the model.
That is exactly what the separate inference script does.
24. Loading the Trained Checkpoint for Inference
For inference, we recreate the architecture:
model = MiniLLM().to(device)
Then load the checkpoint:
model.load_state_dict(
torch.load(
"tiny_stories_gpt2.pth",
map_location=device
),
strict=False
)
The important rule is:
The architecture used during inference must match the architecture used during training.
That includes:
vocab_size
embed_dim
num_heads
num_layers
max_seq_len
If the training model used:
768-dimensional embeddings
12 heads
12 layers
512 positions
the inference model must use the same configuration.
25. A Better Inference Script
The first generation script used two useful restrictions:
logits[0, eot_token] = -float("Inf")
and:
logits[0, last_token] = -float("Inf")
The first one prevents the model from ending.
The second one prevents the model from immediately repeating the exact previous token.
These can help produce longer output, but there is an important downside.
If we always ban <|endoftext|>, the model has no way to naturally finish a story.
And if the model is already somewhat undertrained, forcing it to continue can make repetition worse.
A better generation function is:
@torch.no_grad()
def generate_text(
model,
prompt,
max_new_tokens=250,
temperature=0.7,
top_k=40
):
model.eval()
idx = torch.tensor(
[enc.encode(prompt)],
dtype=torch.long,
device=device
)
for _ in range(max_new_tokens):
idx_cond = idx[
:,
-max_seq_len:
]
if device == "cuda":
with torch.autocast(
device_type="cuda",
dtype=torch.bfloat16
):
logits, _ = model(
idx_cond
)
else:
logits, _ = model(
idx_cond
)
# Only the final position predicts
# the next token.
logits = logits[:, -1, :]
# Temperature
logits = logits / temperature
# Top-K
if top_k is not None:
v, _ = torch.topk(
logits,
min(
top_k,
logits.size(-1)
)
)
logits[
logits < v[:, [-1]]
] = -float("Inf")
probs = F.softmax(
logits,
dim=-1
)
idx_next = torch.multinomial(
probs,
num_samples=1
)
# Stop naturally when the model predicts EOT.
if idx_next.item() == eot_token:
break
idx = torch.cat(
(idx, idx_next),
dim=1
)
return enc.decode(
idx[0].tolist()
)
Then:
print(
generate_text(
"Once upon a time, "
"there was a clever little fox named",
max_new_tokens=250,
temperature=0.7,
top_k=40
)
)
26. Why Did My Model Generate Repetitive Stories?
The output from the experiment was:
Once upon a time, there was a clever little fox named Jack. He lived in a garden with his friends and friends. One day, Max was playing with his owner and Max. Suddenly, a loud dog came and said, "Hi, Max! You have a mean joke!"
The old rabbit said, "Wow, cat! You are a very kind dog. You have so much a big one." So, Spot went to sleep and found his dog Tom. The other cat ran down and said, "Look, little mouse! You are very brave!"
Benny and his friends played in the grass and had fun. They laughed and played games together. They played together and the rest of the day.
"Let's go back here and see what a new adventure.
There are several interesting things happening here.
The model has clearly learned some characteristics of the TinyStories distribution:
- Story-like openings
- Simple sentences
- Children's vocabulary
- Names such as Jack, Max, Tom, Benny, etc.
- Animals
- Dialogue
- Repeated narrative patterns
- Basic grammatical structure
But it has not learned long-range story consistency particularly well.
For example:
fox named Jack
is followed by:
Max
and later:
Spot
and:
Tom
This suggests that the model has learned local statistical patterns from the training data, but its generated narrative state is not consistently maintained over a long sequence.
27. Why Does Repetition Happen?
Repetition can arise from several factors.
1. Limited Training
The model was trained for:
10,000 iterations
A 124M-parameter model still needs substantial optimization to learn strong language and story structure.
More training does not automatically guarantee better results, but insufficient training can leave the model with weak representations.
2. Temperature
The experiment uses:
temperature=0.7
Lower temperatures make the distribution sharper.
If the model strongly prefers a small set of tokens, repeated patterns can become more likely.
Increasing the temperature somewhat can increase diversity, although excessive temperature can make the text incoherent.
For example, you can experiment with:
temperature=0.8
or:
temperature=0.9
3. Top-K
The experiment uses:
top_k=40
This means only the 40 highest-scoring tokens remain candidates.
Top-K can reduce extremely unlikely tokens, but an overly restrictive sampling configuration can also make the model repeatedly choose from a relatively narrow set of tokens.
Try comparing:
top_k=20
top_k=40
and:
top_k=80
and observe how the generated text changes.
4. The Model Is Forced to Continue
The original inference code explicitly does:
logits[0, eot_token] = -float("Inf")
This means:
Never generate
<|endoftext|>.
That is useful if you want to force a fixed-length sample, but it removes the model's ability to decide when the story is finished.
For a story generator, allowing the model to generate its EOT token is generally more natural.
5. Immediate Repetition Is Not the Same as Long-Range Repetition
The original code prevents:
the the
by banning the immediately previous token:
last_token = idx[0, -1].item()
logits[0, last_token] = -float("Inf")
But this does not prevent:
They played together.
...
They played together.
because the repeated token or phrase is not necessarily immediately adjacent.
This is why a simple "don't repeat the previous token" rule is not a complete repetition-control mechanism.
28. The Difference Between Training and Sampling Problems
This distinction is important.
If the model produces:
The dog was happy.
The dog was happy.
The dog was happy.
there are two possible sources.
Training problem
The model may not have learned strong enough representations because it needs:
- More training
- Better data
- Better hyperparameters
- More diverse examples
- Better regularization
Sampling problem
The model may have learned useful representations but the sampling strategy is causing it to repeatedly select the same high-probability patterns.
Therefore, it is useful to test the same checkpoint with multiple sampling configurations.
For example:
generate_text(
prompt,
temperature=0.6,
top_k=40
)
versus:
generate_text(
prompt,
temperature=0.8,
top_k=40
)
versus:
generate_text(
prompt,
temperature=0.9,
top_k=80
)
This helps separate model-quality problems from decoding behavior.
29. The Generation Loop Is Autoregressive
This is one of the most important concepts in a decoder-only LLM.
Suppose the prompt is:
Once upon a time
The model predicts:
there
The sequence becomes:
Once upon a time there
The model predicts:
was
Now:
Once upon a time there was
Then:
a
Then:
little
Then:
fox
and so on.
Conceptually:
Prompt
|
v
Predict token 1
|
v
Append token 1
|
v
Predict token 2
|
v
Append token 2
|
v
Predict token 3
|
v
...
This is why generation is inherently sequential even though Transformer training is highly parallelizable.
30. What the Model Actually Predicts
It is important to remember that the model does not directly predict:
"A complete story."
It predicts:
The probability distribution of the next token given the previous tokens.
For example:
Input:
Once upon a time, there was a clever little fox named
The model might internally produce something like:
Jack → 0.31
Max → 0.14
Sam → 0.08
Tom → 0.06
...
The actual distribution is over all 50,257 vocabulary tokens.
Sampling then chooses one token.
If:
Jack
is selected, the model sees:
Once upon a time, there was a clever little fox named Jack
and computes another distribution.
This happens repeatedly.
The coherent story is therefore an emergent result of thousands of consecutive next-token decisions, not a story that the model writes internally in one operation.
31. From Random Weights to Story Generation
The entire experiment can now be summarized as:
TRAINING
|
v
TinyStories Dataset
|
v
Tokenizer
|
v
Training Batches
|
v
Decoder-Only Transformer
|
v
Logits
|
v
Cross-Entropy Loss
|
v
Backpropagation
|
v
AdamW
|
v
Learned Weights
|
v
tiny_stories_gpt2.pth
|
|
INFERENCE
|
v
Prompt
|
v
Tokenization
|
v
Decoder Transformer
|
v
Logits
|
+-----+-----+
| |
v v
Temperature Top-K
| |
+-----+-----+
|
v
Softmax
|
v
Sampling
|
v
Next Token
|
+----------------+
|
v
Append to Context
|
v
Repeat Generation
This is the complete lifecycle of the model.
32. The Code for This Project
The complete implementation for the project can be published on GitHub.
GitHub Repository: MiniLLM
The model.py file can contain the architecture.
The train.py file can contain the TinyStories training pipeline.
The generate.py file can contain the inference and sampling logic.
33. What I Learned from the First Generation
The generated sample is actually a useful diagnostic.
It demonstrates that the model has learned more than random text.
The output contains:
Once upon a time...
which matches the style of the training corpus.
It produces:
animals
names
dialogue
simple actions
and generally resembles a children's story.
However, the output also contains inconsistencies such as:
fox → Jack → Max → Spot → Tom
and repeated phrases.
That tells us something important:
Training a language model is not simply about making the loss smaller.
We also care about what the model has learned to represent and how those representations behave during generation.
This is exactly why evaluating generated samples is valuable alongside validation loss.
34. What Would I Improve Next?
There are several directions this project can take.
1. Train for longer
Increase:
max_iters = 10000
to a larger number and compare:
- Training loss
- Validation loss
- Generated samples
2. Save checkpoints
Instead of saving only at the end:
torch.save(
model.state_dict(),
"tiny_stories_gpt2.pth"
)
save periodically:
checkpoint_1000.pth
checkpoint_2000.pth
checkpoint_3000.pth
...
This makes it possible to compare model quality throughout training.
3. Track validation loss
The training script already evaluates every:
eval_interval = 500
steps.
Plotting training and validation loss would make it easier to identify:
- Undertraining
- Overfitting
- Optimization instability
4. Experiment with decoding
Compare:
temperature = 0.6
temperature = 0.7
temperature = 0.8
temperature = 0.9
and:
top_k = 20
top_k = 40
top_k = 80
Then compare the resulting samples.
5. Add stronger repetition controls
A future version could experiment with:
- Repetition penalties
- No-repeat n-gram constraints
- Top-P / nucleus sampling
- Frequency penalties
These operate at generation time rather than changing the trained model itself.
6. Add a proper evaluation set
Instead of evaluating only loss, create a fixed collection of prompts such as:
Once upon a time, there was a little dog named
A little girl went to the park and
The small rabbit was scared because
Then generate samples from the same prompts after different checkpoints.
This makes improvements easier to observe.
35. Final Takeaway
This experiment takes us all the way from raw stories to a functioning decoder-only language model.
We started with:
TinyStories
and converted the text into token IDs.
Those tokens were fed into:
Token Embeddings
+
Positional Embeddings
and then passed through:
12 Transformer Blocks
Each block performs:
LayerNorm
↓
Causal Self-Attention
↓
Residual Connection
↓
LayerNorm
↓
Feed-Forward Network
↓
Residual Connection
The final representation is projected into:
50,257 vocabulary logits
Those logits are compared with the correct next token using:
Cross-Entropy Loss
The error is propagated backward through:
Backpropagation
and the weights are updated using:
AdamW
After thousands of iterations, the model learns statistical patterns from TinyStories.
We save those learned parameters:
tiny_stories_gpt2.pth
and load them later for inference.
At inference time, the model receives a prompt such as:
Once upon a time, there was a clever little fox named
and repeatedly performs:
Prompt
↓
Transformer
↓
Logits
↓
Temperature
↓
Top-K
↓
Softmax
↓
Sample next token
↓
Append token
↓
Repeat
That is how the story is generated.
The model doesn't generate the entire story at once.
It generates one token at a time.
And that simple loop, repeated over and over, is at the heart of autoregressive decoder-only language models.
What's Next?
We've now gone from:
"What is a Transformer?"
to:
"How do I actually train one?"
The next step is to go deeper into what happens after pre-training.
That means exploring:
- Instruction fine-tuning
- Supervised fine-tuning
- Chat templates
- Preference optimization
- FlashAttention
- Rotary positional embeddings
- Grouped-query attention
Because a base language model that predicts the next token is only the beginning.
Top comments (0)