Post cover image
Reasoning LLM Pipeline (Created by Fareed Khan)

May 19, 2025

Creating a 2M Parameter Thinking LLM (like o3 & DeepSeek-R1) from Scratch Using Python

From Pretraining to SFT to RLHF

By Fareed Khan

28 min read

Read this story for free: link

You might have already seen lot of articles on creating LLMs using Transformer architecture but recent LLMs like OpenAI o3, Gemini 2.5 Pro and many others can now think, So how can we create a reasoning LLM that can think before answer.

A high-level overview of creating a most simplest reasoning LLM involves three steps:

None
  1. Pretraining an LLM using Transformer architecture.
  2. Supervised Fine-Tuning (SFT) on the pretrained LLM.
  3. Reinforcement Learning with Human Feedback (RLHF) on the SFT-trained model to make it thinkable.

In this blog, I will try to make a reasoning LLM by passing through each of these stages and will check its output.

Final Output

This is the output we are getting (On our sample data vs Proper training data)

None

The model trained on a larger dataset performs much better, while the model trained on the sample dataset is not effective but is useful for learning purposes, as it allows us to debug the code quickly.

This notebook with proper comment is available in this GitHub Repo:

GitHub - FareedKhan-dev/train-llm-from-scratch: A straightforward method for training your LLM… A straightforward method for training your LLM, from downloading data to generating text. …

You can test a much better version of this thinking model created by jingyaogong, using the following same approach (pretrain → SFT → RLHF) (0.1 Billion parameters) at this Hugging Face Space and can also download its weights.

MiniMind - a Hugging Face Space by jingyaogong This app allows you to chat with an AI assistant. You can ask questions or have conversations, and the AI will provide…

Table of Contents

Make sure you have a basic understanding of object-oriented programming (OOP) and neural networks (NN). Familiarity with PyTorch will also be helpful in coding.

Setting the Stage

We will be working with a range of Python libraries throughout this project, so let's import them:

# Standard libraries
import os
import json
import math
import time
import random
import warnings
from typing import Optional, Tuple, List, Union, Iterator

# Numerical computation
import numpy as np

# PyTorch libraries for deep learning
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
from torch.utils.data import Dataset, DataLoader

# For automatic mixed precision and context management
from contextlib import nullcontext

# Hugging Face libraries for tokenization and model handling
from transformers import AutoTokenizer, PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,  # Output type for causal language modeling
    BaseModelOutputWithPast  # Base model output type with cached key/values
)
from transformers.activations import ACT2FN  # Activation functions used in transformer models

# Tokenizers from Hugging Face's `tokenizers` library (fast tokenizer library)
from tokenizers import Tokenizer as HFTokenizer  # Renamed to avoid naming conflicts with transformers' tokenizer
from tokenizers import models as hf_models
from tokenizers import trainers as hf_trainers
from tokenizers import pre_tokenizers as hf_pre_tokenizers
from tokenizers import decoders as hf_decoders

Now that we have imported the required modules, which we will later use for training and preprocessing, let's move on to training a tokenizer.

Our Three Stages Dataset

Since we are dealing with three stages of training, each stage requires a different dataset format.

The first dataset is for pretraining, where we create an LLM from scratch using Transformer architecture, enabling it to generate grammatically correct text.

Let's define that sample training dataset.

# --- Pretraining Data ---
sample_pretrain_data = [
    {"text": "The sun shines brightly in the clear blue sky."},
    {"text": "Cats love to chase mice and play with yarn balls."},
    {"text": "Reading books expands your knowledge and vocabulary."},
    {"text": "Artificial intelligence is a rapidly evolving field of study."},
    {"text": "To bake a cake, you need flour, sugar, eggs, and butter."},
    {"text": "Large language models are trained on vast amounts of text data."},
    {"text": "The quick brown fox jumps over the lazy dog."}
]

Once we train a transformer architecture on this data then we use SFT training dataset on top of that transformer trained model, let's define that dataset also.

# --- SFT Data ---
sample_sft_data = [
    {"conversations": [
        {"role": "user", "content": "Hello, how are you?"},
        {"role": "assistant", "content": "I am doing well, thank you! How can I help you today?"}
    ]},
    {"conversations": [
        {"role": "user", "content": "What is the capital of France?"},
        {"role": "assistant", "content": "The capital of France is Paris."}
    ]},
    {"conversations": [
        {"role": "user", "content": "Explain gravity in simple terms."},
        {"role": "assistant", "content": "Gravity is the force that pulls objects towards each other. It's why things fall down to the ground!"}
    ]}
]

SFT training improves LLM knowledge by providing questions and its true answers.

And finally, after pretraining and SFT training, we create an RLHF dataset, which gives our trained LLM reasoning power.

# --- Reasoning Data ---
sample_reasoning_data = [
    {"conversations": [
        {"role": "user", "content": "If I have 3 apples and eat 1, how many are left?"},
        {"role": "assistant", "content": "<think>The user starts with 3 apples. The user eats 1 apple. This means 1 apple is subtracted from the initial amount. So, 3 - 1 = 2.</think><answer>You have 2 apples left.</answer>"}
    ]},
    {"conversations": [
        {"role": "user", "content": "What are the primary colors?"},
        {"role": "assistant", "content": "<think>The user is asking about primary colors. These are colors that cannot be made by mixing other colors. The standard set of primary colors in additive color models (like light) are Red, Green, and Blue (RGB). For subtractive models (like paint), they are often considered Red, Yellow, Blue (RYB) or Cyan, Magenta, Yellow (CMY).</think><answer>The primary colors are typically considered to be red, yellow, and blue. These are colors that can be mixed to create a range of other colors but cannot be created by mixing other colors themselves.</answer>"}
    ]}
]

You can see in our RLHF dataset that when the user asks a question, the answer includes <think> tags that provide thinking tokens for the LLM to generate before giving the answer.

Now that we have defined the training dataset for each stage, it's time to start training our tokenizer. Let's do that.

Training a BPE Tokenizer

There are several types of tokenizers available either pretrained on a large corpus, like those in the tiktoken library, or custom-trained on our dataset. The larger the corpus of English text we use, the larger its vocabulary size will be.

Let's use a sample dataset for training to see what is happening in the code.

# Tokenizer training data
tokenizer_corpus = [
    "Hello world, this is a demonstration of building a thinking LLM.",
    "Language models learn from text data.",
    "Tokenization is a crucial first step.",
    "We will train a BPE tokenizer.",
    "Think before you answer.",
    "The answer is forty-two.",
    "<think>Let's consider the options.</think><answer>Option A seems best.</answer>",
    "<|im_start|>user\nWhat's up?<|im_end|>\n<|im_start|>assistant\nNot much!<|im_end|>"
]

# Save to a temporary file for the tokenizer trainer
tokenizer_corpus_file = os.path.join("YOUR_DIR_PATH", "tokenizer_corpus.txt")
with open(tokenizer_corpus_file, 'w', encoding='utf-8') as f:
    for line in tokenizer_corpus:
        f.write(line + "\n")

I have just defined a bunch of random text data, you can also use the Shakespeare dataset for this.

Let's create a simple BPE function that will train Tokenizer on our same corpus.

def train_demo_tokenizer(
    corpus_files: List[str],
    vocab_size: int,
    save_path: str,
    special_tokens: List[str]
):
    """
    Trains a Byte-Pair Encoding (BPE) tokenizer using Hugging Face Tokenizers.

    Args:
        corpus_files (List[str]): List of file paths containing training corpus.
        vocab_size (int): Desired vocabulary size.
        save_path (str): Path to save the trained tokenizer.
        special_tokens (List[str]): List of special tokens to include in the tokenizer.
    """
    # Initialize BPE tokenizer with unknown token
    tokenizer_bpe = HFTokenizer(hf_models.BPE(unk_token="<unk>"))

    # Set pre-tokenizer to ByteLevel for robust byte-level processing
    tokenizer_bpe.pre_tokenizer = hf_pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True)

    # Set decoder to ByteLevel for correct decoding of byte-level tokens
    tokenizer_bpe.decoder = hf_decoders.ByteLevel()

    # Create BPE trainer with vocab size, special tokens, and full byte alphabet
    trainer = hf_trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=special_tokens,
        show_progress=True,
        initial_alphabet=hf_pre_tokenizers.ByteLevel.alphabet()
    )

    # Ensure corpus_files is a list
    if isinstance(corpus_files, str):
        corpus_files = [corpus_files]

    # Train the tokenizer
    tokenizer_bpe.train(corpus_files, trainer=trainer)
    print(f"Tokenizer training complete. Vocab size: {tokenizer_bpe.get_vocab_size()}")

    # Save the tokenizer as a JSON file
    tokenizer_bpe.save(save_path)
    print(f"Tokenizer saved to {save_path}")

    return tokenizer_bpe

We are defining <unk> to handle unknown tokens. We also need to create special tokens to indicate the start and end of text, as well as a padding token for adding extra spaces. Let's define these in a list.

# Special tokens
SPECIAL_TOKENS_LIST = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "<pad>"]

It's time to call this function to train our tokenizer.

# Train and save the demo BPE tokenizer
# using the specified corpus, vocab size, save path, and special tokens
trained_hf_tokenizer = train_demo_tokenizer(
    corpus_files=[tokenizer_corpus_file],
    vocab_size=1000,
    save_path="NOTEBOOK_TOKENIZER_PATH",
    special_tokens=SPECIAL_TOKENS_LIST
)

Now that we have trained our tokenizer, let's test it on a sample dataset.

# Test encoding and decoding with the trained tokenizer
test_sentence = "Hello <|im_start|> world <think>思考中</think><answer>答案</answer> <|im_end|>"

# Encode the sentence
encoded = trained_hf_tokenizer.encode(test_sentence)

# Display original sentence and its encoded form
print(f"Original: {test_sentence}")
print(f"Encoded IDs: {encoded.ids}")
print(f"Encoded Tokens: {encoded.tokens}")

# Decode back to text
decoded = trained_hf_tokenizer.decode(encoded.ids)
print(f"Decoded: {decoded}")

### OUTPUT ###
Trained Tokenizer Vocab (first 10 and special tokens):
  'ain': 328
  '}': 96
  'D': 39
  'Ġfirst': 428
  'Wh': 324
  ...

You can see how our trained BPE tokenizer breaks down our test sentence. Let's move on to creating our Transformer model, which will be used to perform stage 1 training. Let's do that.

Transformer Overview

Let's have a quick look at how a transformer architecture is used to process and understand text. It works by breaking text into smaller pieces called tokens and predicting the next token in the sequence. A transformer has many layers, called transformer blocks, stacked on top of each other, with a final layer at the end to make the prediction.

Each transformer block has two main components:

  • Self-Attention Heads: These figure out which parts of the input are most important for the model to focus on. For example, when processing a sentence, the attention heads can highlight relationships between words, such as how a pronoun relates to the noun it refers to.
  • MLP (Multi-Layer Perceptron): This is a simple feed-forward neural network. It takes the information emphasized by the attention heads and processes it further. The MLP has an input layer that receives data from the attention heads, a hidden layer that adds complexity to the processing, and an output layer that passes the results to the next transformer block.

Together, the attention heads act as the "what to think about" part, while the MLP is the "how to think about it" part. Stacking many transformer blocks allows the model to understand complex patterns and relationships in the text, but this is not always guaranteed.

let's visualize a simpler and easier architecture diagram that we will be coding.

None

Let's read through the flow of our architecture that we will be coding:

  1. Input tokens are converted to embeddings and combined with position information.
  2. The model has 64 identical transformer blocks that process data sequentially.
  3. Each block first runs multi-head attention to look at relationships between tokens.
  4. Each block then processes data through an MLP that expands and then compresses the data.
  5. Each step uses residual connections (shortcuts) to help information flow.
  6. Layer normalization is used throughout to stabilize training.
  7. The attention mechanism calculates which tokens should pay attention to each other.
  8. The MLP expands the data to 4x size, applies ReLU, and then compresses it back down.
  9. The model uses 16 attention heads to capture different types of relationships.
  10. The final layer converts the processed data into vocabulary-sized predictions.
  11. The model generates text by repeatedly predicting the next most likely token.

Multi Layer Perceptron (MLP)

MLP is a fundamental building block within the transformer's feed-forward network. Its role is to introduce non-linearity and learn complex relationships within the embedded representations. When defining an MLP module, an important parameter is n_embed, which defines the dimensionality of the input embedding.

The MLP typically consists of a hidden linear layer that expands the input dimension by a factor (often 4, which we will use), followed by a non-linear activation function, commonly ReLU. This structure allows our network to learn more complex features. Finally, a projection linear layer maps the expanded representation back to the original embedding dimension. This sequence of transformations enables the MLP to refine the representations learned by the attention mechanism.

None 
# --- MLP (Multi-Layer Perceptron) Class ---

class MLP(nn.Module):
    """
    A simple Multi-Layer Perceptron with one hidden layer.
    This module is used within the Transformer block for feed-forward processing.
    It expands the input embedding size, applies a ReLU activation, and then projects it back
    to the original embedding size.
    """
    def __init__(self, n_embed):
        super().__init__()
        self.hidden = nn.Linear(n_embed, 4 * n_embed)  # Linear layer to expand embedding size
        self.relu = nn.ReLU()                        # ReLU activation function
        self.proj = nn.Linear(4 * n_embed, n_embed)  # Linear layer to project back to original size
    def forward(self, x):
        """
        Forward pass through the MLP.
        Args:
            x (torch.Tensor): Input tensor of shape (B, T, C), where B is batch size,
                              T is sequence length, and C is embedding size.
        Returns:
            torch.Tensor: Output tensor of the same shape as the input.
        """
        x = self.forward_embedding(x)
        x = self.project_embedding(x)
        return x
    def forward_embedding(self, x):
        """
        Applies the hidden linear layer followed by ReLU activation.
        Args:
            x (torch.Tensor): Input tensor.
        Returns:
            torch.Tensor: Output after the hidden layer and ReLU.
        """
        x = self.relu(self.hidden(x))
        return x
    def project_embedding(self, x):
        """
        Applies the projection linear layer.
        Args:
            x (torch.Tensor): Input tensor.
        Returns:
            torch.Tensor: Output after the projection layer.
        """
        x = self.proj(x)
        return x

We just coded our MLP part, where the __init__ method initializes a hidden linear layer that expands the input embedding size (n_embed) and a projection layer that reduces it back. ReLU activation is applied after the hidden layer.

The forward method defines the data flow through these layers, applying the hidden layer and ReLU via forward_embedding, and the projection layer via project_embedding.

Single Head Attention

The attention head is the core part of our model. Its purpose is to focus on relevant parts of the input sequence.

When defining a Head module, some important parameters are head_size, n_embed, and context_length. The head_size parameter determines the dimensionality of the key, query, and value projections, influencing the representational capacity of the attention mechanism.

The input embedding dimension n_embed defines the size of the input to these projection layers. context_length is used to create a causal mask, ensuring that the model only attends to preceding tokens.

Within the Head, linear layers (nn.Linear) for key, query, and value are initialized without bias. A lower triangular matrix (tril) of size context_length x context_length is registered as a buffer to implement causal masking, preventing the attention mechanism from attending to future tokens.

None 
# --- Attention Head Class ---

class Head(nn.Module):
    """
    A single attention head.
    This module calculates attention scores and applies them to the values.
    It includes key, query, and value projections, and uses causal masking
    to prevent attending to future tokens.
    """
    def __init__(self, head_size, n_embed, context_length):
        super().__init__()
        self.key = nn.Linear(n_embed, head_size, bias=False)   # Key projection
        self.query = nn.Linear(n_embed, head_size, bias=False) # Query projection
        self.value = nn.Linear(n_embed, head_size, bias=False) # Value projection
        # Lower triangular matrix for causal masking
        self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))
    def forward(self, x):
        """
        Forward pass through the attention head.
        Args:
            x (torch.Tensor): Input tensor of shape (B, T, C).
        Returns:
            torch.Tensor: Output tensor after applying attention.
        """
        B, T, C = x.shape
        k = self.key(x)     # (B, T, head_size)
        q = self.query(x)   # (B, T, head_size)
        scale_factor = 1 / math.sqrt(C)
        # Calculate attention weights: (B, T, head_size) @ (B, head_size, T) -> (B, T, T)
        attn_weights = q @ k.transpose(-2, -1) * scale_factor
        # Apply causal masking
        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)   # (B, T, head_size)
        # Apply attention weights to values
        out = attn_weights @ v # (B, T, T) @ (B, T, head_size) -> (B, T, head_size)
        return out

Our attention head class __init__ method initializes linear layers for key, query, and value projections, each projecting the input embedding (n_embed) to head_size.

A lower triangular matrix based on context_length is used for causal masking.

The forward method calculates attention weights by scaling the dot product of the query and key, applies the causal mask, normalizes the weights using softmax, and computes the weighted sum of the values to produce the attention output.

Multi Head Attention

To capture diverse relationships within the input sequence, we are going to use the concept of multi-head attention. The MultiHeadAttention module manages multiple independent attention heads operating in parallel.

The key parameter here is n_head, which determines the number of parallel attention heads. The input embedding dimension (n_embed) and context_length are also necessary to instantiate the individual attention heads.

Each head processes the input independently, projecting it into a lower-dimensional subspace of size n_embed // n_head. By having multiple heads, the model can attend to different aspects of the input simultaneously.

None 
# --- Multi-Head Attention Class ---

class MultiHeadAttention(nn.Module):
    """
    Multi-Head Attention module.
    This module combines multiple attention heads in parallel. The outputs of each head
    are concatenated to form the final output.
    """
    def __init__(self, n_head, n_embed, context_length):
        super().__init__()
        self.heads = nn.ModuleList([Head(n_embed // n_head, n_embed, context_length) for _ in range(n_head)])
    def forward(self, x):
        """
        Forward pass through the multi-head attention.
        Args:
            x (torch.Tensor): Input tensor of shape (B, T, C).
        Returns:
            torch.Tensor: Output tensor after concatenating the outputs of all heads.
        """
        # Concatenate the output of each head along the last dimension (C)
        x = torch.cat([h(x) for h in self.heads], dim=-1)
        return x

Now that we have defined the MultiHeadAttention class, which combines multiple attention heads, the __init__ method initializes a list of Head instances (a total of n_head), each with a head_size of n_embed // n_head.

The forward method applies each attention head to the input x and concatenates their outputs along the last dimension, merging the information learned by each head.

Transformer Block

To create a billion-parameter model, we definitely need a deep architecture. For that, we need to code a transformer block and stack them.

The key parameters of a block are n_head, n_embed, and context_length. Each block comprises a multi-head attention layer and a feed-forward network (MLP), with layer normalization applied before each and residual connections after each.

Layer normalization, parameterized by the embedding dimension n_embed, helps stabilize training. The multi-head attention mechanism, as described before, takes n_head, n_embed, and context_length.

The MLP also utilizes the embedding dimension n_embed. These components work together to process the input and learn complex patterns.

None 
# --- Transformer Block Class ---

class Block(nn.Module):
    """
    A single Transformer block.
    This block consists of a multi-head attention layer followed by an MLP,
    with layer normalization and residual connections.
    """
    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):
        """
        Forward pass through the Transformer block.
        Args:
            x (torch.Tensor): Input tensor.
        Returns:
            torch.Tensor: Output tensor after the block.
        """
        # Apply multi-head attention with residual connection
        x = x + self.attn(self.ln1(x))
        # Apply MLP with residual connection
        x = x + self.mlp(self.ln2(x))
        return x
    def forward_embedding(self, x):
        """
        Forward pass focusing on the embedding and attention parts.
        Args:
            x (torch.Tensor): Input tensor.
        Returns:
            tuple: A tuple containing the output after MLP embedding and the residual.
        """
        res = x + self.attn(self.ln1(x))
        x = self.mlp.forward_embedding(self.ln2(res))
        return x, res

Our Block class represents a single transformer block. The __init__ method initializes layer normalization layers (ln1, ln2), a MultiHeadAttention module, and an MLP module, all parameterized by n_head, n_embed, and context_length.

The forward method implements the block's forward pass, applying layer normalization and multi-head attention with a residual connection, followed by another layer normalization and the MLP, again with a residual connection.

The forward_embedding method provides an alternative forward pass focused on the attention and initial MLP embedding stages.

The Final Model

So far, we have coded small components of the transformer model. Next, we integrate token and position embeddings with a series of transformer blocks to perform sequence-to-sequence tasks. To do that, we need to code several key parameters: n_head, n_embed, context_length, vocab_size, and N_BLOCKS.

vocab_size determines the size of the token embedding layer, mapping each token to a dense vector of size n_embed. The context_length parameter is important for the position embedding layer, which encodes the position of each token in the input sequence, also with dimension n_embed. The number of attention heads (n_head) and the number of blocks (N_BLOCKS) dictate the depth and complexity of the network.

These parameters collectively define the architecture and capacity of the transformer model, so let's code it.

None 
# --- Transformer Model Class ---

class Transformer(nn.Module):
    """
    The main Transformer model.
    This class combines token and position embeddings with a sequence of Transformer blocks
    and a final linear layer for language modeling.
    """
    def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
        super().__init__()
        self.context_length = context_length
        self.N_BLOCKS = N_BLOCKS
        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 _pre_attn_pass(self, idx):
        """
        Combines token and position embeddings.
        Args:
            idx (torch.Tensor): Input token indices.
        Returns:
            torch.Tensor: Sum of token and position embeddings.
        """
        B, T = idx.shape
        tok_embedding = self.token_embed(idx)
        pos_embedding = self.position_embed(self.pos_idxs[:T])
        return tok_embedding + pos_embedding
    def forward(self, idx, targets=None):
        """
        Forward pass through the Transformer.
        Args:
            idx (torch.Tensor): Input token indices.
            targets (torch.Tensor, optional): Target token indices for loss calculation. Defaults to None.
        Returns:
            tuple: Logits and loss (if targets are provided).
        """
        x = self._pre_attn_pass(idx)
        for block in self.attn_blocks:
            x = block(x)
        x = self.layer_norm(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            B, T, C = logits.shape
            flat_logits = logits.view(B * T, C)
            targets = targets.view(B * T).long()
            loss = F.cross_entropy(flat_logits, targets)
        return logits, loss
    def forward_embedding(self, idx):
        """
        Forward pass focusing on the embedding and attention blocks.
        Args:
            idx (torch.Tensor): Input token indices.
        Returns:
            tuple: Output after attention blocks and the residual.
        """
        x = self._pre_attn_pass(idx)
        residual = x
        for block in self.attn_blocks:
            x, residual = block.forward_embedding(x)
        return x, residual
    def generate(self, idx, max_new_tokens):
        """
        Generates new tokens given a starting sequence.
        Args:
            idx (torch.Tensor): Initial sequence of token indices.
            max_new_tokens (int): Number of tokens to generate.
        Returns:
            torch.Tensor: The extended sequence of 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

Our Transformer class __init__ method initializes token and position embedding layers (token_embed, position_embed), a sequence of Block modules (attn_blocks), a final layer normalization layer (layer_norm), and a linear layer for language modeling (lm_head).

The _pre_attn_pass method combines token and position embeddings. The forward method processes the input sequence through the embedding layers and the series of transformer blocks, applies final layer normalization, and generates logits.

It also calculates the loss if targets are provided. The forward_embedding method provides an intermediate forward pass up to the output of the attention blocks, and the generate method implements token generation.

Pre-Training Transformer

Before we train, let's set some hyperparameters for our Transformer. For this demonstration, we'll keep them relatively small to ensure we can run this on a standard machine.

# --- Define Model Hyperparameters ---
CONTEXT_LENGTH = 64   # Maximum number of tokens the model will look at in one go
N_EMBED = 256         # Embedding dimension (size of the token vectors)
N_HEAD = 8            # Number of attention heads
N_BLOCKS = 6          # Number of Transformer blocks to stack
VOCAB_SIZE = trained_hf_tokenizer.get_vocab_size() # Get vocab size from our trained tokenizer
LEARNING_RATE = 3e-4  # Learning rate for the optimizer
BATCH_SIZE = 32       # Number of sequences processed at once
N_EPOCHS = 100        # Number of times we'll go through the entire dataset (for demo)

# --- Device Configuration ---
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

# --- Instantiate the Transformer Model ---
model = Transformer(
    n_head=N_HEAD,
    n_embed=N_EMBED,
    context_length=CONTEXT_LENGTH,
    vocab_size=VOCAB_SIZE,
    N_BLOCKS=N_BLOCKS
).to(DEVICE)

# Print the number of parameters in our model
num_params = sum(p.numel() for p in model.parameters())
print(f"Our Transformer model has {num_params/1e6:.2f} million parameters.")

### OUTPUT

Using device: cuda (or cpu if no GPU)
ur Transformer model has 2.15 million parameters.

So, we have got a model with roughly 2 million parameters. This is tiny compared to giants like GPT-3, but it's a great starting point for understanding the mechanics.

We will feed data through our model, calculate how wrong its predictions are (the "loss"), and adjust its internal parameters (weights) using an optimizer to make it better over time.

# --- Optimizer ---
# AdamW is a popular choice for training Transformers
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

# --- Mixed Precision (Optional but Recommended for GPUs) ---
# Autocast enables automatic mixed precision, using float16 for some operations
# to speed up training and reduce memory, while maintaining precision for others.
PTDTYPE = torch.float32 # Default for blog
autocast_ctx = nullcontext() if DEVICE == 'cpu' else torch.amp.autocast(device_type=DEVICE, dtype=PTDTYPE)
# GradScaler helps prevent gradients from becoming too small (underflowing) with float16
scaler = torch.cuda.amp.GradScaler(enabled=(PTDTYPE != torch.float32 and DEVICE == 'cuda'))

# --- Pretraining Loop ---
print(f"Starting Pretraining for {N_EPOCHS} epochs...")
model.train() # Set the model to training mode

for epoch in range(N_EPOCHS):
    total_loss = 0
    batch_count = 0

    for xb, yb in pretrain_dataloader:
        # Move batch to the configured device
        xb, yb = xb.to(DEVICE), yb.to(DEVICE)

        # Zero out gradients from the previous iteration
        optimizer.zero_grad(set_to_none=True)

        with autocast_ctx:
            # Forward pass: get model's predictions (logits) and loss
            # Our Transformer's forward method returns (logits, loss)
            logits, loss = model(xb, targets=yb)

        if loss is not None: # Ensure loss was computed
            # Backward pass: compute gradients of the loss w.r.t. model parameters
            # If using scaler, scale the loss before backward pass
            scaler.scale(loss).backward()
            # Optimizer step: update model parameters
            scaler.step(optimizer)
            # Update the scaler for the next iteration
            scaler.update()

            total_loss += loss.item()
        batch_count += 1

        # Log progress (e.g., every 10 steps for a small dataset)
        if batch_count % max(1, len(pretrain_dataloader)//2) == 0 : # Log twice per epoch
            print(f"Epoch {epoch+1}/{N_EPOCHS}, Batch {batch_count}/{len(pretrain_dataloader)}, Loss: {loss.item():.4f}")

    avg_epoch_loss = total_loss / batch_count if batch_count > 0 else float('nan')
    print(f"--- End of Epoch {epoch+1}, Average Pretraining Loss: {avg_epoch_loss:.4f} ---")

print("Pretraining finished.")

# --- Save the Pretrained Model ---
PRETRAINED_MODEL_PATH = os.path.join(NOTEBOOK_OUT_DIR, "thinking_llm_pretrained.pth")
torch.save(model.state_dict(), PRETRAINED_MODEL_PATH)
print(f"Pretrained model weights saved to: {PRETRAINED_MODEL_PATH}")

When we start the training loop it will print loss for each epoch

### OUTPUT ###
Starting Pretraining for 100 epochs...
Epoch 1/100, Batch 3/7, Loss: 5.8732
Epoch 1/100, Batch 6/7, Loss: 5.1234
--- End of Epoch 1, Average Pretraining Loss: 5.4321 ---
... (many epochs later) ...
Epoch 100/100, Batch 3/7, Loss: 1.5678
Epoch 100/100, Batch 6/7, Loss: 1.4990

--- End of Epoch 100, Average Pretraining Loss: 1.5220 ---
Pretraining finished.
Pretrained model weights saved to: ./out_notebook_scratch_blog/thinking_llm_pretrained.pth

Our model has now learned some basic language structure from the pretraining data. It can probably generate somewhat coherent (but likely repetitive or nonsensical on very small data) text if we were to test its generate method now.

Inferencing Pre-Trained Model

Let's do a quick test to check if our pre-trained model has learned anything, even if minimal.

pt_model.eval()  # Set to evaluation mode
test_prompt_str_pt = "Language models learn"
# Prepend BOS for generation consistency with training if tokenizer doesn't do it automatically
pt_test_input_ids = tokenizer(tokenizer.bos_token + test_prompt_str_pt, return_tensors="pt").input_ids.to(DEVICE)

with torch.no_grad(), autocast_ctx:
    generated_output_pt = pt_model.generate(
        pt_test_input_ids,
        max_new_tokens=15,
        do_sample=False,  # Greedy for this test
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id
    )
decoded_generated_pt = tokenizer.decode(generated_output_pt[0], skip_special_tokens=True)
print(f"Prompt: '{test_prompt_str_pt}' -> Generated: '{decoded_generated_pt}'")

### OUTPUT ###
Language models learn learn learn learn learn

There is a clear hallucination in the output, which is expected since we are training the model for only a few epochs and on a small dataset.

I trained the same architecture using the Pile dataset (~ 5 GB chunk), which is much larger, and for the same number of epochs, we get this output.

# 2 Million Parameter LLM Output (on sample data)
Language models learn learn learn learn learn

# 2 Million Parameter LLM Output (on PILE data)
Language models were directly linked to

But for now, we have a pretrained LLM. Before moving on to the next step of SFT, there is usually a lot of work involved, such as proper hallucination checks, removing swear words, and several trial-and-error iterations.

However, since this blog is a simple implementation of how thinking LLMs work, we can move directly to the SFT stage. Let's do that.

SFT Dataset Loader

Pretraining provides the model with a broad understanding of language by exposing it to large amounts of general text. However, to make the model more useful for specific tasks like answering questions or following instructions, we use Supervised Fine-Tuning (SFT).

In SFT, the model is trained on curated datasets such as our sample_sft_data which include pairs of user prompts and ideal assistant responses.

A critical part of fine-tuning chat models is loss masking. During training, both the user's input and the assistant's reply are included in the model's input sequence.

However, we only want the model to learn from the assistant's reply, not from the user's prompt (which it will receive as input during real interactions).

Loss masking ensures that only the tokens corresponding to the assistant's response contribute to the loss calculation. This focuses the model's learning on generating appropriate and helpful replies, while ignoring the prompt itself during optimization.

First, we need to create a SFT Dataset loader class, which will handle the SFT training for our sample dataset.

class SFTDataset(Dataset):
    """
    Simplified Dataset for Supervised Fine-Tuning using ChatML-style formatting.
    Only the assistant's responses are used for loss computation.
    """
    def __init__(self, file_path, tokenizer, context_length):
        self.tokenizer = tokenizer
        self.context_length = context_length

        with open(file_path, 'r', encoding='utf-8') as f:
            self.conversations = [json.loads(line)['conversations'] for line in f]

        print(f"Loaded {len(self.conversations)} conversations.")

    def __len__(self):
        return len(self.conversations)

    def __getitem__(self, idx):
        conversation = self.conversations[idx]

        tokens = [self.tokenizer.bos_token_id]
        labels = [-100]

        for turn in conversation:
            role, content = turn['role'], turn['content']

            # Encode turn
            prefix = self.tokenizer.encode(f"<|im_start|>{role}\n", add_special_tokens=False).ids
            content_ids = self.tokenizer.encode(content, add_special_tokens=False).ids
            suffix = self.tokenizer.encode("<|im_end|>\n", add_special_tokens=False).ids

            # Append tokens
            tokens.extend(prefix + content_ids + suffix)

            # Create labels
            if role == 'assistant':
                labels.extend([-100] * len(prefix))
                labels.extend(content_ids)
                labels.extend(suffix)
            else:
                labels.extend([-100] * (len(prefix) + len(content_ids) + len(suffix)))

        # Ensure length = context_length + 1 for shifting
        pad_id = self.tokenizer.pad_token_id
        tokens = tokens[:self.context_length + 1] + [pad_id] * max(0, self.context_length + 1 - len(tokens))
        labels = labels[:self.context_length + 1] + [-100] * max(0, self.context_length + 1 - len(labels))

        # Prepare input and target
        input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
        target_ids = torch.tensor(labels[1:], dtype=torch.long)

        return input_ids, target_ids

The SFTDataset prepares the input and label sequences so that only the assistant's responses are used when calculating the loss during training.

This means that tokens from the user are included in the input but ignored in the loss, focusing the learning only on the assistant's part.

In more complex cases, more careful alignment of the tokens might be needed, but this covers the basic approach.

SFT Training Loop

Let's code the SFT training loop, which will use the same pretrained model that we trained earlier using transformer architecture.

# --- SFT Training Loop ---
# Use a smaller learning rate for fine-tuning
SFT_LEARNING_RATE = 1e-4
SFT_N_EPOCHS = 50 # SFT often needs more epochs on smaller datasets

optimizer_sft = torch.optim.AdamW(model.parameters(), lr=SFT_LEARNING_RATE)
# Scaler and autocast_ctx can be reused if GPU settings are the same

print(f"Starting SFT for {SFT_N_EPOCHS} epochs...")
model.train() # Set model to training mode

for epoch in range(SFT_N_EPOCHS):
    total_sft_loss = 0
    sft_batch_count = 0

    for xb_sft, yb_sft_labels in sft_dataloader:
        xb_sft, yb_sft_labels = xb_sft.to(DEVICE), yb_sft_labels.to(DEVICE)

        optimizer_sft.zero_grad(set_to_none=True)

        with autocast_ctx:
            # Forward pass: our model's forward takes `idx` and `targets`.
            # `targets` here should be `yb_sft_labels` which already has -100 for non-assistant parts.
            # The model internally computes CrossEntropyLoss using these labels.
            logits_sft, loss_sft = model(xb_sft, targets=yb_sft_labels)

        if loss_sft is not None:
            scaler.scale(loss_sft).backward()
            scaler.step(optimizer_sft)
            scaler.update()
            total_sft_loss += loss_sft.item()
        sft_batch_count += 1

        if sft_batch_count % max(1, len(sft_dataloader)//1) == 0: # Log once per epoch for small SFT data
             print(f"SFT Epoch {epoch+1}/{SFT_N_EPOCHS}, Batch {sft_batch_count}/{len(sft_dataloader)}, Loss: {loss_sft.item():.4f}")

    avg_epoch_sft_loss = total_sft_loss / sft_batch_count if sft_batch_count > 0 else float('nan')
    print(f"--- End of SFT Epoch {epoch+1}, Average SFT Loss: {avg_epoch_sft_loss:.4f} ---")

print("SFT finished.")

This loop is very simple the only change here is the sft loss, when we start the training, it will start printing each epochs loss, take a look at the output.

### OUTPUT ###
Starting SFT for 50 epochs...
SFT Epoch 1/50, Batch 1/1, Loss: 2.1241
--- End of SFT Epoch 1, Average SFT Loss: 2.1241 ---
...
SFT Epoch 50/50, Batch 1/1, Loss: 0.8721
--- End of SFT Epoch 50, Average SFT Loss: 0.8721 ---
SFT finished.

Our pretrained model, which has previously learned to generate words, has now acquired basic language question-answering abilities due to our SFT training data. It can likely produce somewhat coherent responses, such as replying "Hi, how can I help…" when we say "Hi".

Inferencing SFT Model

Though our SFT training dataset was very small and our pretrained model was quite limited, here is the output of our LLM that was pretrained on the Pile dataset and fine-tuned with SFT on OpenOrca dataset for 50 epochs.

# Prepare a simple test conversation
chat_history = [{"role": "user", "content": "What is the capital of France?"}]
prompt = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

# Generate response
with torch.no_grad(), autocast_ctx_sft:
    outputs = sft_model_demo.generate(
        inputs.input_ids,
        max_new_tokens=200,
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id
    )

# Decode and log response
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
logger(f"SFT Prompt: '{chat_history[0]['content']}' -> Generated: '{response}'")

We are asking a very basic question, let's see the output of our better model.

# 2 Million Parameter SFT LLM Output (on sample SFT data)
unbiased asf gg link

# 2 Million Parameter SFT LLM Output (on ORCA SFT data)
paris is the bigger

A model trained on a larger dataset is starting to provide a more accurate answer to the question, "What is the capital of France?"

But for now, let's assume our SFT model is performing well. Next comes the thinking training, Reinforcement Learning which gives the LLM its reasoning power.

RLHF-style Dataset Loader

We can largely reuse the SFTDataset structure because the core format of conversations is similar. The sample_reasoning_data already contains the assistant's "thinking" and "answer" parts embedded within its content.

The SFTDataset will ensure that only these assistant parts (including the special tags) are used to calculate the loss.

# --- Save Reasoning Data to a File ---
reasoning_file_path = os.path.join(NOTEBOOK_DATA_DIR, "reasoning_data.jsonl")
with open(reasoning_file_path, 'w', encoding='utf-8') as f:
    for item in sample_reasoning_data: # sample_reasoning_data was defined at the start of your blog
        f.write(json.dumps(item) + '\n')
print(f"Reasoning data saved to: {reasoning_file_path}")

# --- Instantiate Reasoning Dataset and DataLoader ---
# We reuse the SFTDataset class as it correctly handles conversation formatting
# and labels only assistant turns for loss calculation.
# The <think> and <answer> tags are within the assistant's content.
reasoning_dataset = SFTDataset(
    file_path=reasoning_file_path,
    tokenizer=trained_hf_tokenizer, # Our trained tokenizer
    context_length=CONTEXT_LENGTH   # Hyperparameter defined earlier
)
# Use a smaller batch size if memory is constrained, especially for larger models or context lengths
REASONING_BATCH_SIZE = max(1, BATCH_SIZE // 2) # Example: half of pretrain/SFT batch size
reasoning_dataloader = DataLoader(reasoning_dataset, batch_size=REASONING_BATCH_SIZE, shuffle=True)

The dataset loader correctly prepares the sequences so the model learns to generate the entire assistant response, including the crucial <think> and <answer> tags.

Training Loop with Weighted Loss for Tags

To really make the model prioritize learning the <think> and <answer> structure, we'll modify the loss calculation slightly.

We will identify these specific tag tokens and assign them a higher weight in the loss, signaling to the model that getting these structural tokens right is especially important.

# SFT_MODEL_PATH = os.path.join(NOTEBOOK_OUT_DIR, "thinking_llm_sft.pth") # Example path
if os.path.exists(SFT_MODEL_PATH):
    model.load_state_dict(torch.load(SFT_MODEL_PATH, map_location=DEVICE))
    print(f"Loaded SFT-trained weights into the model from {SFT_MODEL_PATH}")
else:
    print(f"WARNING: SFT model path {SFT_MODEL_PATH} not found. Reasoning training will use current model state.")

# --- Reasoning Training Hyperparameters ---
REASONING_LEARNING_RATE = 5e-5  # Often smaller for this final tuning stage
REASONING_N_EPOCHS = 75         # More epochs might be needed to learn the new structure

optimizer_rsn = torch.optim.AdamW(model.parameters(), lr=REASONING_LEARNING_RATE)

# --- Define Special Reasoning Tags and Their IDs ---
# These tags should have been added to SPECIAL_TOKENS_LIST during tokenizer training
# to ensure they are treated as single tokens by `trained_hf_tokenizer`.
think_token_id = trained_hf_tokenizer.token_to_id("<think>")
end_think_token_id = trained_hf_tokenizer.token_to_id("</think>")
answer_token_id = trained_hf_tokenizer.token_to_id("<answer>")
end_answer_token_id = trained_hf_tokenizer.token_to_id("</answer>")

Now that we have defined the training parameters lets defien the training loop for RLHF learning.

# --- Reasoning Training Loop ---
print(f"Starting Reasoning Training for {REASONING_N_EPOCHS} epochs...")
model.train() # Set model to training mode

for epoch in range(REASONING_N_EPOCHS):
    total_reasoning_loss = 0
    rsn_batch_count = 0

    for xb_rsn, yb_rsn_labels in reasoning_dataloader:
        xb_rsn, yb_rsn_labels = xb_rsn.to(DEVICE), yb_rsn_labels.to(DEVICE)

        optimizer_rsn.zero_grad(set_to_none=True)

        with autocast_ctx:
            # Get logits from the model. Our model.forward() with `targets` calculates standard CE loss.
            # For weighted loss, we need to get raw logits and compute loss manually.
            logits_rsn, _ = model(xb_rsn) # Pass idx only to get logits

            # Flatten logits and labels for per-token loss calculation
            # logits_rsn shape: (B, T, VocabSize), yb_rsn_labels shape: (B, T)
            flat_logits_rsn = logits_rsn.view(-1, VOCAB_SIZE)  # (B*T, VocabSize)
            flat_labels_rsn = yb_rsn_labels.view(-1)           # (B*T)

            # Calculate raw loss for each token (respecting ignore_index=-100)
            raw_loss_per_token = loss_fct_no_reduction(flat_logits_rsn, flat_labels_rsn)

            # Create a weight tensor for each token.
            # Initialize weights: 1.0 for target tokens (where label is not -100), 0.0 otherwise.
            token_weights = (flat_labels_rsn != -100).float()

            # Identify positions of our special reasoning tags in the target labels
            # is_special_tag_mask will be True where flat_labels_rsn is one of our special_reasoning_tag_ids
            is_special_tag_mask = torch.isin(flat_labels_rsn, special_reasoning_tag_ids)

            # Apply higher weight to special tags *that are also valid target tokens*
            # (i.e., their label is not -100 AND they are one of our special tags)
            target_special_tags_mask = is_special_tag_mask & (flat_labels_rsn != -100)
            token_weights[target_special_tags_mask] *= REASONING_TAG_LOSS_WEIGHT

            # Calculate the final weighted loss
            # Sum of (loss_per_token * weight_per_token)
            # Normalize by the number of original target tokens (where label != -100)
            # This keeps the loss magnitude roughly comparable to unweighted SFT,
            # while still emphasizing the tags.
            num_actual_target_tokens = (flat_labels_rsn != -100).float().sum().clamp(min=1)
            weighted_loss_rsn = (raw_loss_per_token * token_weights).sum() / num_actual_target_tokens

        # Backward pass and optimizer step, checking for valid loss
        if not torch.isnan(weighted_loss_rsn) and not torch.isinf(weighted_loss_rsn):
            scaler.scale(weighted_loss_rsn).backward()
            scaler.step(optimizer_rsn)
            scaler.update()
            total_reasoning_loss += weighted_loss_rsn.item()
        else:
            # This can happen if num_actual_target_tokens is zero for some reason (e.g., all labels are -100 in a batch)
            print(f"Warning: NaN or Inf loss encountered in Reasoning Epoch {epoch+1}, Batch {rsn_batch_count+1}. Skipping update for this batch.")

        rsn_batch_count += 1

        # Log progress
        if rsn_batch_count % max(1, len(reasoning_dataloader)//1) == 0: # Log once per epoch if few batches
            current_loss_item = weighted_loss_rsn.item() if not (torch.isnan(weighted_loss_rsn) or torch.isinf(weighted_loss_rsn)) else float('nan')
            print(f"Reasoning Epoch {epoch+1}/{REASONING_N_EPOCHS}, Batch {rsn_batch_count}/{len(reasoning_dataloader)}, Weighted Loss: {current_loss_item:.4f}")

    avg_epoch_rsn_loss = total_reasoning_loss / rsn_batch_count if rsn_batch_count > 0 else float('nan')
    print(f"--- End of Reasoning Epoch {epoch+1}, Average Reasoning Loss: {avg_epoch_rsn_loss:.4f} ---")

print("Reasoning Training finished.")

In this loop, the model is being trained specifically to handle reasoning tasks over several epochs. For each batch of data, it calculates predictions (logits) and compares them to the true labels.

The key detail is that the loss (which measures how wrong the model's predictions are) is adjusted to give extra importance (higher weight) to certain special reasoning tags, making the model focus more on learning these critical parts.

After calculating the weighted loss, the model updates its internal parameters using backpropagation and an optimizer. We also includes safety checks to skip batches where the loss calculation might go wrong (e.g., producing NaN or Inf).

Throughout the training, it logs the progress and reports the average loss for each epoch.

### OUTPUT ###
Loaded SFT-trained weights into the model from ./blog_output_dir/thinking_llm_sft.pth
Special Reasoning Tag IDs for weighting: [ID_think, ID_endthink, ID_answer, ID_endanswer] (actual token IDs)
Starting Reasoning Training for 75 epochs...
Reasoning Epoch 1/75, Batch 1/1, Weighted Loss: 1.9503
--- End of Reasoning Epoch 1, Average Reasoning Loss: 1.9503 ---
... (many epochs later) ...
Reasoning Epoch 75/75, Batch 1/1, Weighted Loss: 0.3881
--- End of Reasoning Epoch 75, Average Reasoning Loss: 0.3881 ---

It starts the training and prints the loss. After the training is completed, our trained RLHF model gets saved.

We are expecting that our trained RLHF model will generate responses after <think>, meaning it will first generate the thinking tokens and then produce the response based on that thinking process.

Inferencing Our Thinking LLM

As previously observed, SFT or pretraining is only effective when done on larger datasets with more epochs and parameters. In RLHF, I observed the same, model trained on larger RLHF datasets generate much better responses. While not perfect, they are closer to our minimum expectations.

Inferencing RLHF model invovles slight variation in the code we see earlier because the response involves tags in it as compared to previous SFT and pre-train llm responses which are just raw output.

# Load the final reasoning model config matching the saved model
final_model_config = DemoLLMConfig(
    vocab_size=DEMO_VOCAB_SIZE_FINAL,
    hidden_size=DEMO_HIDDEN_SIZE,
    intermediate_size=DEMO_INTERMEDIATE_SIZE,
    num_hidden_layers=DEMO_NUM_LAYERS,
    num_attention_heads=DEMO_NUM_ATTENTION_HEADS,
    num_key_value_heads=DEMO_NUM_KV_HEADS,
    max_position_embeddings=DEMO_MAX_SEQ_LEN,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
)

# Initialize and load the trained model weights
final_thinking_llm = DemoLLMForCausalLM(final_model_config).to(DEVICE)
final_thinking_llm.load_state_dict(torch.load(final_reasoning_model_path, map_location=DEVICE))
final_thinking_llm.eval()

# Prepare the prompt for the model using the tokenizer's chat template
user_query = "If I have 3 apples and eat 1, how many are left?"
chat_history = [{"role": "user", "content": user_query}]
prompt_text = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)

# Tokenize the prompt and move input IDs to the device
input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids.to(DEVICE)

# Generate the model response using sampling parameters
with torch.no_grad(), autocast_ctx_rsn:
    generated_ids = final_thinking_llm.generate(
        input_ids,
        max_new_tokens=DEMO_MAX_SEQ_LEN - 10,
        do_sample=True,
        temperature=0.7,
        top_k=10,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
    )

# Decode the generated tokens, skipping special tokens
assistant_response = tokenizer.decode(generated_ids[0][input_ids.shape[1]:], skip_special_tokens=True)

# Helper function to extract content within tags
def extract_tag(text, tag):
    start_tag = f"<{tag}>"
    end_tag = f"</{tag}>"
    start_idx = text.find(start_tag)
    end_idx = text.find(end_tag)
    if start_idx != -1 and end_idx != -1:
        return text[start_idx + len(start_tag):end_idx].strip()
    return None

# Extract <think> and <answer> sections if they exist
think = extract_tag(assistant_response, "think") or "Not found"
answer = extract_tag(assistant_response, "answer") or assistant_response

# Print the parsed thinking and answer parts
print(f"<think>: {think}")
print(f"<answer>: {answer}")

We have already seen that training a model on a very small sample dataset gives no good result.

However, I also fine-tuned the RLHF model on a better dataset, like the HFHub RL data, and the results were significantly better compared to our model trained on tiny data. Take a look.

# user query
If I have 3 apples and eat 1, how many are left?

# 2 Million Parameter RLHF LLM Output (on sample RLHD data)
<think> ugh far igu when

# 2 Million Parameter SFT LLM Output (on HFHub RLHF data)
<think> User ask apple question how many ge left?</think>2

Our model trained on the sample dataset produced gibberish, but the one trained on a larger dataset with similar parameters started recognizing the thinking tokens, albeit with some errors.

Most importantly, it predicted the correct answer, and while some of the thinking tokens had minor issues, these can definitely be improved with larger training on more extensive data.

What's Next

I suggest starting by developing a model with over 13 million parameters and gradually scaling it by adding increments of 100 parameters, focusing on improving its performance with shorter contexts.

You can decide how many additional parameters to train based on your specific objectives. Once the model approaches 1 billion parameters, consider fine-tuning it on domain-specific datasets, such as email or essay writing, to evaluate how well it generates text in those areas.