I Built a Small LLM From Scratch — Here’s What I Finally Understood About Transformer

onlydev — Wed, 21 Oct 2026 05:13:19 +0000

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?

DEV Community: onlydev

I Built a Small LLM From Scratch — Here’s What I Finally Understood About Transformer

What I Built

1. First Problem: How Does Text Become Numbers?

2. RMSNorm — A Simpler Way to Normalize Activations

3. RoPE — How Does a Transformer Know Where a Token Is?

4. GQA — Reducing Attention Memory

5. Causal Attention — Why Can't the Model See the Future?

6. SwiGLU — The Feed-Forward Part Isn't Just ReLU

7. Putting Everything Together

8. The Complete Mini LLM

9. Training the Model

10. Training Loss vs Validation Loss

11. Finally: Generating Text

12. Temperature Controls Randomness

What I Actually Learned From This Project

Why I Built This Instead of Only Using an LLM API

My Biggest Takeaways

1. Build, don't just watch

2. Understand the “why”

3. Follow the tensors

4. AI engineering is more than model training

5. Projects become much more valuable when you can explain them

What's Next?

Final Thought