Demystifying LLMs: Building a 124M-Parameter Decoder-Only Transformer in PyTorch

S Aaron Dennis — Wed, 30 Sep 2026 19:19:06 +0000

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 ( $Q$ )

The Query represents what information the current token is looking for.

Key ( $K$ )

The Key represents what information a token contains and can be matched against.

Value ( $V$ )

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:

Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V

There are several steps hidden inside this equation.

Step 1 — Compute $Q K^{T}$

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:

d_{k}

where $d_{k}$ 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 $V$

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 $t$ can only attend to positions:

0, 1, 2, \dots, t

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:

Head Dimension = \frac{768}{12} = 64

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:

FFN (x) = W_{2} GELU (W_{1} x + b_{1}) + b_{2}

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:

\frac{1}{2 N}

where $N$ represents the number of Transformer layers.

For 12 layers:

\frac{1}{2 ( 12 )} = \frac{1}{24}

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:

64 \times 512 = 32, 768

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:

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:

L = - lo g P (y ∣ x)

where:

$x$ is the input context
$y$ is the correct next token
$P (y ∣ x)$ 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:

θ_{t + 1} = θ_{t} - η \nabla_{θ} L

where:

$θ$ = model parameters
$η$ = learning rate
$\nabla_{θ} L$ = 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:

η_{t} = η_{m i n} + \frac{1}{2} (1 + cos (π r)) (η_{m a x} - η_{m i n})

where $r$ 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:

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.

DEV Community: S Aaron Dennis

Demystifying LLMs: Building a 124M-Parameter Decoder-Only Transformer in PyTorch

1. The Heart of the Model: Masked Self-Attention

Query, Key, and Value

Query ( Q )

Key ( K )

Value ( V )

The Attention Equation

Step 1 — Compute QKT

Step 2 — Scale the scores

Step 3 — Apply Softmax

Step 4 — Multiply by V

2. Why Decoder-Only Models Need a Causal Mask

3. Implementing Causal Self-Attention in PyTorch

4. What Happens Inside Multi-Head Attention?

5. The Transformer Block

6. What Does the Feed-Forward Network Do?

7. Building the Full MiniLLM

Token Embeddings

8. Context Window and Positional Embeddings

9. Full MiniLLM Implementation

10. Weight Tying

11. GPT-2-Style Weight Initialization

12. Training the 124M-Parameter Model

Hardware Used

13. The Dataset: TinyStories

14. Tokenization with GPT-2 BPE

15. Turning the Dataset into One Long Token Stream

16. Creating Training Batches

17. The Training Objective: Next-Token Prediction

Training

Generation

18. Cross-Entropy Loss

19. Backpropagation

20. AdamW Updates the Parameters

21. Gradient Clipping

22. Learning-Rate Warmup and Cosine Decay

23. The Complete Training Script

25. The Training Loop in One Picture

26. Saving the Trained Model

24. Loading the Trained Checkpoint for Inference

25. A Better Inference Script

26. Why Did My Model Generate Repetitive Stories?

27. Why Does Repetition Happen?

1. Limited Training

2. Temperature

3. Top-K

4. The Model Is Forced to Continue

5. Immediate Repetition Is Not the Same as Long-Range Repetition

28. The Difference Between Training and Sampling Problems

Training problem

Sampling problem

29. The Generation Loop Is Autoregressive

30. What the Model Actually Predicts

31. From Random Weights to Story Generation

32. The Code for This Project

33. What I Learned from the First Generation

34. What Would I Improve Next?

1. Train for longer

2. Save checkpoints

3. Track validation loss

4. Experiment with decoding

5. Add stronger repetition controls

6. Add a proper evaluation set

35. Final Takeaway

What's Next?

Query ( $Q$ )

Key ( $K$ )

Value ( $V$ )

Step 1 — Compute $Q K^{T}$

Step 4 — Multiply by $V$