I used to understand Large Language Models mostly from the outside.
I knew terms like:
- Transformer
- Attention
- Embeddings
- Tokens
- Fine-tuning
- GPT
- LLM
But knowing the terminology is very different from understanding what actually happens inside a model.
So instead of only reading about LLMs, I decided to build a small language model from scratch using PyTorch.
The goal wasn't to build the next ChatGPT.
The goal was much more important for me as a learner:
What I Built
The model starts with Shakespeare text and learns to predict the next character.
The basic pipeline is:
Text
↓
Character Tokenization
↓
Token Embeddings
↓
Transformer Blocks
↓
RMSNorm
↓
Grouped Query Attention + RoPE
↓
SwiGLU Feed Forward Network
↓
Final Normalization
↓
Language Model Head
↓
Next-token prediction
The model configuration I worked with was:
Vocabulary size: 65
Embedding dimension: 256
Transformer blocks: 4
Query heads: 8
KV heads: 2
SwiGLU hidden size: 680
Context length: 256
The important thing is that these aren't just names in a diagram. I implemented the pieces and followed the tensors through the model.
1. First Problem: How Does Text Become Numbers?
Neural networks don't understand strings directly.
So the first step was tokenization.
For this project, I used a character-level tokenizer.
For example:
"Hello"
can become something like:
[20, 43, 50, 50, 53]
The tokenizer maintains two mappings:
char_to_idx = {
character: integer
}
idx_to_char = {
integer: character
}
Then:
def encode(s):
return [char_to_idx[c] for c in s]
def decode(ids):
return "".join([idx_to_char[i] for i in ids])
The important lesson for me was simple:
Before a model can learn language, language has to become numerical data.
The training examples are created by shifting the sequence by one position.
For example:
Input: To be o
Target: o be or
So the model is continuously learning:
Given previous characters → predict the next character
This turns one sequence of length N into N next-token prediction examples.
2. RMSNorm — A Simpler Way to Normalize Activations
One component I learned about was RMSNorm (Root Mean Square Normalization).
Traditional LayerNorm performs operations involving the mean and standard deviation.
RMSNorm simplifies this.
Conceptually:
x
↓
calculate RMS
↓
normalize x
↓
learnable scale
The implementation is roughly:
rms = torch.sqrt(
x.pow(2).mean(dim=-1, keepdim=True) + eps
)
output = (x / rms) * weight
The interesting part is that RMSNorm only needs a learnable scale vector.
My takeaway:
Normalization isn't just a mathematical trick. It helps keep the values flowing through a deep network under control.
The material compares RMSNorm with LayerNorm and notes its use in modern architectures such as LLaMA, Mistral, and Gemma.
3. RoPE — How Does a Transformer Know Where a Token Is?
This was one of the concepts I found most interesting.
A transformer needs positional information.
One approach is to add positional embeddings.
RoPE — Rotary Positional Embeddings — takes a different approach.
Instead of simply adding a position vector, RoPE rotates the Query and Key vectors.
The important intuition is:
Position 3 → Position 5
distance = 2
Position 7 → Position 9
distance = 2
RoPE allows the attention mechanism to capture this kind of relative positional relationship.
At a high level:
Query + position
Key + position
↓
rotations
↓
attention
The rotation is applied to pairs of dimensions:
(d0, d1)
(d2, d3)
(d4, d5)
...
and each pair gets its own rotation frequency.
The biggest lesson I took from this:
Position isn't just “where is this token?” — attention can benefit from understanding how far apart tokens are.
4. GQA — Reducing Attention Memory
Next came Grouped Query Attention (GQA).
This was especially useful because it showed me that modern LLM architecture isn't only about making models bigger.
It's also about making them more efficient.
Instead of giving every Query head its own Key and Value heads, GQA allows multiple Query heads to share KV heads.
In my model:
Query heads = 8
KV heads = 2
So four Query heads share each KV head.
Conceptually:
Q0 ─┐
Q1 ─┤
Q2 ─┤──→ K0, V0
Q3 ─┘
Q4 ─┐
Q5 ─┤
Q6 ─┤──→ K1, V1
Q7 ─┘
This significantly reduces the size of the KV cache.
The implementation first projects Q, K and V, reshapes them into heads, applies RoPE to Q and K, and then expands the KV heads to match the Query heads.
This taught me an important engineering lesson:
AI engineering isn't only about model accuracy. Memory usage and inference efficiency matter too.
5. Causal Attention — Why Can't the Model See the Future?
During training, the model shouldn't be allowed to cheat.
Suppose we're predicting:
The cat sat
When predicting the next token after "The", the model shouldn't be able to look at "cat".
That's why causal masking is used.
Conceptually:
T0 T1 T2 T3
T0 ✓ ✗ ✗ ✗
T1 ✓ ✓ ✗ ✗
T2 ✓ ✓ ✓ ✗
T3 ✓ ✓ ✓ ✓
Each token can attend to itself and previous tokens, but not future tokens.
The attention computation looks roughly like:
scores = (q @ k.transpose(-2, -1)) * scale
scores = scores.masked_fill(
mask,
float("-inf")
)
weights = F.softmax(scores, dim=-1)
out = weights @ v
This made the phrase “next-token prediction” much more concrete for me.
6. SwiGLU — The Feed-Forward Part Isn't Just ReLU
Another component I implemented was SwiGLU.
A traditional feed-forward network can look like:
Input
↓
Linear
↓
ReLU
↓
Linear
↓
Output
SwiGLU introduces a gating mechanism.
It has two paths:
┌──→ SiLU ──┐
Input ───────┤ × ──→ Linear ──→ Output
└──→ Up ────┘
The core operation is:
gate = F.silu(self.w_gate(x))
up = self.w_up(x)
output = self.w_down(gate * up)
The gate learns which information should be emphasized or suppressed.
The training material describes this as a gated dual-path feed-forward network, with SiLU on the gate path and element-wise multiplication controlling information flow.
7. Putting Everything Together
A transformer block in my implementation follows this pattern:
Input
│
├── RMSNorm
│
├── GQA + RoPE
│
└── Residual Add
↓
x'
│
├── RMSNorm
│
├── SwiGLU
│
└── Residual Add
↓
Output
In code:
x = x + self.attention(
self.attn_norm(x),
rope_cos,
rope_sin
)
x = x + self.ffn(
self.ffn_norm(x)
)
This is a pre-norm + residual architecture.
The residual connection means the block learns a change to the existing representation instead of having to completely rebuild it at every layer.
8. The Complete Mini LLM
The complete model then becomes:
Character IDs
↓
Token Embedding
↓
Transformer Block × 4
│
├── RMSNorm
├── GQA
│ └── RoPE
├── Residual
├── RMSNorm
├── SwiGLU
└── Residual
↓
Final RMSNorm
↓
Language Model Head
↓
Logits
↓
Softmax
↓
Next Character
The model uses weight tying between the token embedding and language-model head.
The complete architecture in the project uses 4 transformer blocks, 8 Query heads, 2 KV heads, and a 256-dimensional representation.
9. Training the Model
Once the architecture was ready, the next question was:
How does it actually learn?
The training loop follows the familiar deep-learning pattern:
Get batch
↓
Forward pass
↓
Calculate loss
↓
Zero gradients
↓
Backward pass
↓
Gradient clipping
↓
Optimizer step
The project used:
Batch size: 64
Learning rate: 3e-4
Optimizer: AdamW
Training steps: 3000
The training code follows:
xb, yb = get_batch(...)
logits, loss = model(xb, yb)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0
)
optimizer.step()
Gradient clipping is used to reduce the risk of exploding gradients during training.
10. Training Loss vs Validation Loss
One thing I don't want to overlook when building ML models is evaluation.
A model getting better on training data doesn't automatically mean it's generalizing well.
So we can compare:
Training loss
vs
Validation loss
A growing gap between them can be a signal of overfitting.
During evaluation, dropout is disabled and gradients aren't needed:
@torch.no_grad()
def estimate_loss():
model.eval()
...
model.train()
The training material uses the train/validation loss gap as an indicator for overfitting.
11. Finally: Generating Text
After training, the model can generate new text.
The process is autoregressive:
Prompt
↓
Model
↓
Predict next token
↓
Append token
↓
Model again
↓
Predict next token
↓
Repeat...
For example:
"ROMEO:"
gets encoded into token IDs.
The model predicts the next character, adds it to the sequence, and repeats the process.
The generation code takes the logits from the last position, converts them into probabilities using softmax, samples the next token, and appends it to the sequence.
12. Temperature Controls Randomness
Another concept that became much clearer after implementing generation was temperature.
Lower temperature:
More predictable
Less random
Higher temperature:
More diverse
More random
For example:
T = 0.3 → conservative
T = 0.8 → balanced
T = 1.5 → creative
So temperature doesn't make the model “smarter”.
It changes how we sample from its probability distribution.
What I Actually Learned From This Project
The biggest lesson wasn't memorizing what RMSNorm, RoPE, GQA and SwiGLU stand for.
It was understanding how the pieces connect.
Before this project, I could say:
“Transformers use attention.”
Now I can explain:
Text
↓
Tokens
↓
Embeddings
↓
Q / K / V
↓
RoPE
↓
Causal Attention
↓
GQA
↓
Residual Connection
↓
SwiGLU
↓
Repeated Transformer Blocks
↓
Logits
↓
Probability Distribution
↓
Next Token
That difference is important.
Knowing the vocabulary is not the same as understanding the system.
Why I Built This Instead of Only Using an LLM API
As someone preparing for an AI/ML engineering role, I think there is value in going one layer deeper.
Using:
response = client.chat(...)
is useful.
But understanding what happens underneath gives me a stronger foundation for:
- debugging ML systems
- reading research papers
- understanding model architectures
- optimizing inference
- experimenting with transformers
- explaining technical decisions during interviews
I don't expect a small Shakespeare model to qualify me for an LLM engineering role by itself.
But building it gave me something more valuable:
A mental model of how modern language models work.
My Biggest Takeaways
1. Build, don't just watch
Reading about transformers is useful.
Implementing one forces you to understand the tensor shapes and data flow.
2. Understand the “why”
Don't memorize:
RoPE
GQA
RMSNorm
SwiGLU
Ask:
Why is it there?
What problem does it solve?
What happens if I remove it?
3. Follow the tensors
One of the best ways I found to understand deep learning code is to constantly ask:
What is the shape here?
For example:
[b, seq, 256]
↓
[b, 8, seq, 32]
↓
attention
↓
[b, seq, 256]
4. AI engineering is more than model training
A good ML engineer needs to think about:
Data
↓
Architecture
↓
Training
↓
Evaluation
↓
Inference
↓
Efficiency
5. Projects become much more valuable when you can explain them
This is probably the biggest career lesson I took from this.
Don't just put:
“Built an LLM using PyTorch”
on a resume.
Be able to explain:
“I implemented character tokenization, causal next-token prediction, RMSNorm, RoPE, GQA with shared KV heads, SwiGLU, residual transformer blocks, AdamW training, validation loss evaluation, and autoregressive generation.”
That demonstrates actual understanding.
What's Next?
This project gave me a foundation.
My next step is to move from:
“I understand how the architecture works.”
to:
“I can build useful AI systems with it.”
That means exploring things like:
- fine-tuning
- embeddings
- RAG
- vector databases
- evaluation
- inference optimization
- model deployment
- AI agents
- production ML pipelines
The goal isn't to learn every AI buzzword.
The goal is to build things, understand them, and document what I learn.
Final Thought
I'm learning that the fastest way to improve as an AI/ML engineer isn't to wait until I know everything.
It's to:
Learn → Build → Break → Debug → Explain → Share → Repeat.
This small LLM project was one step in that process.
And I'm going to keep documenting the journey.
Tech Stack: Python · PyTorch · Deep Learning · Transformers · LLMs
What would you recommend I build next — a RAG system, fine-tuned model, or an AI agent?
Top comments (0)