Train Your Own LLM From Scratch: A Complete End-to-End Guide
A hands-on tutorial for building, pretraining, and aligning a Transformer LLM from scratch using PyTorch, covering SFT, DPO, PPO, and GRPO.
Have you ever wanted to train your own large language model (LLM) from scratch, not just fine-tune an existing one? The open-source project train-llm-from-scratch by FareedKhan-dev provides a complete, transparent pipeline that takes you from raw text to an aligned, reasoning-capable model. This isn't a black-box wrapper around existing libraries; it's a hands-on educational journey where every algorithm—from multi-head attention to GRPO—is implemented in plain PyTorch.
This guide will walk you through the entire process, explaining the why and how behind each step, so you can understand the fundamentals and run the code yourself.
Why Build an LLM From Scratch?
Most developers interact with LLMs through high-level APIs like Hugging Face's transformers or trl. While these are powerful, they abstract away the core mechanics. Building from scratch gives you:
- Deep Understanding: You learn exactly how tokenization, attention, and training loops work.
- Full Control: You can modify any component—loss functions, architectures, or training strategies—without library constraints.
- Educational Value: It's the best way to grasp concepts like causal masking, reward modeling, and policy gradient methods.
This project is designed for students, developers, and researchers who want to see the entire pipeline in one place, from data download to text generation.
The Pipeline at a Glance
The journey follows a clear, sequential path:
- Data Preparation: Raw text → Token IDs → Stored arrays.
- Model Construction: Building a Transformer from small, reusable components (MLP, Attention, Blocks).
- Pretraining: Next-token prediction on a large corpus (The Pile).
- Text Generation: Sampling from the trained model.
- Post-Training (Alignment): Turning a base model into a helpful assistant via SFT, Reward Modeling, DPO, PPO, and GRPO.
- Evaluation: Measuring performance on a benchmark like GSM8K.
- Inference & Chat: Interacting with the final model.
Step 1: Preparing the Data
A model only understands integers. The first job is to convert text into token IDs and store them efficiently.
Tokenization
The project uses OpenAI's r50k_base tokenizer (from tiktoken), the same one used by GPT-3. Text becomes a list of integers, and a special <|endoftext|> token (ID 50256) is appended to each document so the model learns boundaries.
# Download and tokenize data from The Pile
python scripts/data_download.py
python scripts/data_preprocess.py
For the newer, faster path:
python scripts/prepare_pretrain_data.py --split val --out data/pile_dev.h5
python scripts/prepare_pretrain_data.py --split train --num_shards 1 --out data/pile_train.h5
Chat Format and Loss Masking
For post-training, the model needs to understand roles (user vs. assistant). The project uses plain text markers:
<|user|>
What is 13 + 29?<|endoftext|><|assistant|>
<think>13 + 29 = 42</think><answer>42</answer><|endoftext|>
A loss mask ensures the model only learns from the assistant's tokens, not the user's prompt. This is critical for SFT and RL.
def encode_chat(messages, add_generation_prompt=False):
ids, mask = [], []
for m in messages:
role = m["role"]
header_ids = _encode_ordinary(_header_for(role))
ids.extend(header_ids)
mask.extend([0] * len(header_ids))
content_ids = _encode_ordinary(m["content"])
is_completion = role == "assistant"
ids.extend(content_ids)
mask.extend([1 if is_completion else 0] * len(content_ids))
ids.append(EOT_ID)
mask.append(1 if is_completion else 0)
return ids, mask
Step 2: Building the Transformer
The Transformer is built from four small, reusable PyTorch modules. This modular design makes the code easy to understand and modify.
Multi-Layer Perceptron (MLP)
The "thinking" part of each block. It expands the token vector, applies a ReLU, and projects it back.
class MLP(nn.Module):
def __init__(self, n_embed):
super().__init__()
self.hidden = nn.Linear(n_embed, 4 * n_embed)
self.relu = nn.ReLU()
self.proj = nn.Linear(4 * n_embed, n_embed)
def forward(self, x):
x = self.relu(self.hidden(x))
x = self.proj(x)
return x
Single Head Attention
Allows a token to look at other tokens. It computes query, key, and value vectors, applies a causal mask (so a token can't see future tokens), and takes a weighted sum of values.
class Head(nn.Module):
def __init__(self, head_size, n_embed, context_length):
super().__init__()
self.key = nn.Linear(n_embed, head_size, bias=False)
self.query = nn.Linear(n_embed, head_size, bias=False)
self.value = nn.Linear(n_embed, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))
def forward(self, x):
B, T, C = x.shape
k = self.key(x)
q = self.query(x)
scale_factor = 1 / math.sqrt(C)
attn_weights = q @ k.transpose(-2, -1) * scale_factor
attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
attn_weights = F.softmax(attn_weights, dim=-1)
v = self.value(x)
out = attn_weights @ v
return out
Multi-Head Attention
Runs multiple attention heads in parallel, each learning different relationships, then concatenates and projects the results.
The Transformer Block
Combines attention and MLP with pre-norm residual connections. This is the core repeating unit.
class Block(nn.Module):
def __init__(self, n_head, n_embed, context_length):
super().__init__()
self.ln1 = nn.LayerNorm(n_embed)
self.attn = MultiHeadAttention(n_head, n_embed, context_length)
self.ln2 = nn.LayerNorm(n_embed)
self.mlp = MLP(n_embed)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
The Full Transformer
Wraps everything: token embeddings, position embeddings, a stack of blocks, final layer norm, and a linear projection to vocabulary-sized logits.
class Transformer(nn.Module):
def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
super().__init__()
self.token_embed = nn.Embedding(vocab_size, n_embed)
self.position_embed = nn.Embedding(context_length, n_embed)
self.attn_blocks = nn.ModuleList([Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)])
self.layer_norm = nn.LayerNorm(n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
self.register_buffer('pos_idxs', torch.arange(context_length))
def forward(self, idx, targets=None):
x = self.forward_hidden(idx)
logits = self.lm_head(x)
loss = None
if targets is not None:
B, T, C = logits.shape
flat_logits = logits.reshape(B * T, C)
targets = targets.reshape(B * T).long()
loss = F.cross_entropy(flat_logits, targets)
return logits, loss
Step 3: Pretraining the Base Model
Pretraining is next-token prediction on a large corpus. The model reads random windows of tokens, predicts the next token at every position, and updates its weights via cross-entropy loss.
# Train a 13M parameter model
python scripts/train_transformer.py
# For larger models with memory optimizations
python scripts/train_transformer.py --amp --grad-checkpointing --grad-accum 8
The loss starts near ln(vocab_size) (~10.8) and drops as the model learns language statistics. A 77M parameter model trained on 2x L40 GPUs reached a loss of ~3.73 after 2000 steps.
Step 4: Generating Text
Once trained, you can generate text by sampling from the model's output distribution.
def generate(self, idx, max_new_tokens):
for _ in range(max_new_tokens):
idx_cond = idx[:, -self.context_length:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
python scripts/generate_text.py --model_path models/transformer_B.pt --input_text "The" --max_new_tokens 100
Step 5: Post-Training (Alignment)
A base model can continue text, but it can't follow instructions. Post-training aligns it using supervised fine-tuning (SFT) and reinforcement learning (RL).
SFT (Supervised Fine-Tuning)
Teaches the model to answer in a chat format. The loss is only computed on the assistant's tokens.
python scripts/prepare_sft_data.py --context_length 1024
torchrun --standalone --nproc_per_node=2 scripts/train_sft.py
The Reward Model
Trains a small linear head on top of the SFT model to score answers. Uses the Bradley-Terry loss to prefer chosen responses over rejected ones.
def bradley_terry_loss(chosen_rewards, rejected_rewards):
return -F.logsigmoid(chosen_rewards - rejected_rewards).mean()
DPO, ORPO, and KTO
Direct Preference Optimization (DPO) skips the reward model and RL loop, working directly on preference pairs. ORPO and KTO are variants.
torchrun --standalone --nproc_per_node=2 scripts/train_dpo.py --loss_type dpo
PPO (Proximal Policy Optimization)
The classic RLHF loop. The model generates answers, scores them with a reward model, and updates via a clipped policy loss.
def ppo_policy_loss(new_logp, old_logp, advantages, mask, clip=0.2):
ratio = torch.exp(new_logp - old_logp)
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1.0 - clip, 1.0 + clip) * advantages
loss = -masked_mean(torch.min(surr1, surr2), mask)
return loss
GRPO (Group Relative Policy Optimization)
The DeepSeek-R1 style method. It samples a group of answers for each prompt, scores them, and uses the group's mean and std as a baseline.
def group_advantages(rewards, group_size, eps=1e-4):
r = rewards.view(-1, group_size)
mean = r.mean(dim=1, keepdim=True)
std = r.std(dim=1, keepdim=True)
adv = (r - mean) / (std + eps)
return adv.reshape(-1)
Step 6: Evaluation
The key metric is greedy GSM8K accuracy: the model must output the correct answer inside <answer> tags.
for s in base_pretrained sft dpo ppo grpo; do
python scripts/eval_post_training.py --ckpt models/$s.pt --label $s --limit 200 --append logs/table.jsonl
done
Step 7: Talking to the Model
Use the chat script to interact with any checkpoint.
python scripts/chat.py --ckpt models/sft.pt --prompt "What is 13 + 29?"
Getting Started
Clone the repo:
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git cd train-llm-from-scratch pip install -e .Start small: Train the 13M parameter model first. It's fast and will give you a feel for the pipeline.
Scale up: Increase
n_embedandn_blocksuntil you hit your GPU's memory limit.Walk the post-training chain: Run SFT, then DPO, then PPO or GRPO, and watch the GSM8K accuracy improve.
This project is a rare gem: it's both a practical tool and a comprehensive tutorial. Whether you're a student wanting to understand Transformers deeply or a developer looking to experiment with alignment techniques, train-llm-from-scratch is the perfect starting point.