Post cover image
SpongeBob (For Header Image)

March 5, 2025

Building a Perfect Million Parameter LLM Like ChatGPT in Python

step-by-step guide

By Fareed Khan

32 min read

Read this story for free: link

Quick Note β€” We will first train a tokenizer and then build a 29-million-parameter LLM from scratch. This will give us a model that generates proper sentences. Next, we will fine-tune it using (SFT) to improve its knowledge and response style, making it more like ChatGPT.

I have deployed my trained tiny model on huggingface space. You can chat with it there. Web app ink

Take a look at a few chat conversations between me and our trained LLM.

None None None None

Instead of going through all the theory at once, we will code alongside it to understand everything properly. Everything, from the dataset to the model weights, is replaceable.

GitHub Code

All the code, along with the proper setup, is available in my GitHub repo.

GitHub - FareedKhan-dev/train-tiny-llm: Train a 29M parameter GPT from Scratch Train a 29M parameter GPT from Scratch. Contribute to FareedKhan-dev/train-tiny-llm development by creating an account…

The codebase is organized as follows:

.
β”œβ”€β”€ pretrain.py              # Pre-training script
β”œβ”€β”€ requirement.txt          # Project dependencies
β”œβ”€β”€ train_sft.py             # Supervised fine-tuning script
β”œβ”€β”€ train_tokenizer.py       # Tokenizer training script
β”œβ”€β”€ web_app.py               # Streamlit web application
└── transformer/             # Directory containing core components
    β”œβ”€β”€ dataset.py           # Dataset classes (PretrainDataset, SFTDataset)
    β”œβ”€β”€ LMConfig.py          # Language model configuration class (LMConfig)
    └── model.py             # Transformer model definition (TransformerLM)

GitHub codebase is more flexible, you can perform parallel training and much more!

Table of Contents

Training a Tokenizer

A tokenizer is the first and the important component of Large Language Models (LLMs) that breaks text into smaller units called tokens (words, subwords, or characters).

It is needed because LLMs do not understand raw text but they process numbers.

Tokenizers convert text into numerical representations so the model can analyze and learn language patterns. Every LLM uses tokenization before training to ensure efficient text processing and understanding.

Let's use GPT-2 tokenizer and see how it thinks.

# Load the GPT-2 tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Example text
text = "Hello, world!"

# Tokenization process
tokens = tokenizer.tokenize(text)  # Converts text into tokens
token_ids = tokenizer.convert_tokens_to_ids(tokens)  # Converts tokens into numerical IDs

# Output results
print("Tokens:", tokens)
print("Token IDs:", token_ids)

### OUTPUT ###
Tokens: ['Hello', ',', 'Δ world', '!']
Token IDs: [15496, 11, 995, 0]

It converts text to numerical forms. GPT-2 tokenizer is trained on a large amount of data. If you want to use your own tokenizer, like an already trained one, you can use it.

But we can also build our own Tokenizer, so let's do that.

To train a tokenizer, we need data. Since we are building an English-speaking LLM, Hugging Face is filled with massive datasets that can be used for training a tokenizer.

I am going to use the allenai/c4 Hugging Face dataset, which is created from web scraping. It is very large, but I will use 2 million English rows. You can decide the number on your own. So, let's begin.

# Load dataset in streaming mode
ds = load_dataset("allenai/c4", "en", streaming=True)

# Number of rows to retrieve
num_rows_to_get = 2000000

# Output file
output_file = "training_data.jsonl"

# Write data to file
with open(output_file, "w", encoding="utf-8") as f:
    for i, row in enumerate(iter(ds["train"])):
        if i >= num_rows_to_get:
            break
        json.dump({"text": row["text"]}, f, ensure_ascii=False)
        f.write("\n")  # Ensure each JSON object is on a new line

So, we have converted 2M Hugging Face data into a JSONL file because this format is easy to process.

Let's print a sample of our training dataset to see how it looks.

#this code for displaying purpose only , it not needed to run training
import json

def print_sample_from_jsonl(file_path, num_samples=3):
    """Print a few sample entries from a JSONL file."""
    with open(file_path, "r", encoding="utf-8") as f:
        # Iterate up to num_samples lines
        for _, line in zip(range(num_samples), f):
            print(json.dumps(json.loads(line), indent=4, ensure_ascii=False))

# Assuming your training data is in 'training_data.jsonl'
print_sample_from_jsonl("training_data.jsonl")

### OUTPUT ###
{ "text": "Beginners BBQ Class ..." }
{ "text": "Discussion in 'Mac  ..." }
{ "text": "Foil plaid lycra an ..." }

This will output a few lines from our training_data.jsonl file. Each line is a separate JSON object containing a "text" field.

Now, let's start building our tokenizer! We will use the tokenizers library. We are going to build a Byte-Pair Encoding (BPE) tokenizer.

BPE starts with individual characters and iteratively merges the most frequent pairs, building a vocabulary of subwords. BPE (Byte Pair Encoding) is best because it balances vocabulary size and efficiency, making it great for handling rare and common words in NLP tasks.

First, we need a way to efficiently load our text data from the JSONL file. We'll use a generator for this. Generators are great for large datasets because they don't load everything into memory at once.

def load_texts_from_jsonl(file_path):
    """
    Generator function to read and yield text data from a JSONL file.
    """
    with open(file_path, "r", encoding="utf-8") as file:
        for line in file:
            yield json.loads(line)["text"]

This load_texts_from_jsonl function takes the file path, reads it line by line, extracts the 'text' field, and yields it. This makes it memory-efficient.

Next, let's initialize our tokenizer and set up the pre-tokenizer:

# Initialize a new BPE tokenizer
tokenizer = Tokenizer(models.BPE())

# Set up the pre-tokenizer to handle splitting text into bytes
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)

We create a Tokenizer object using the BPE model.

The pre_tokenizer splits the input text into bytes. add_prefix_space=False is important, it tells the tokenizer not to add a space before the first word of a sequence.

Now, let's define our special tokens and configure the trainer:

# Define special tokens
special_tokens = ["<unk>", "<s>", "</s>"]

# Configure the BPE trainer
trainer = trainers.BpeTrainer(
    vocab_size=6400,  # We'll use a vocabulary size of 6400
    special_tokens=special_tokens,
    show_progress=True,
    initial_alphabet=pre_tokenizers.ByteLevel.alphabet()
)

We have three special tokens (I will tell why they are important when we create our pretraining dataset).

  • <unk> (for unknown words)
  • <s> (start of sequence)
  • </s> (end of sequence).

The BpeTrainer is set up with our desired vocab_size=6400, the special_tokens, and a progress bar (show_progress=True). The initial_alphabet makes sure all individual bytes are in the starting vocabulary.

It's time to train the tokenizer.

# Load the text data using our generator
texts = load_texts_from_jsonl("training_data.jsonl") #your file path

# Train the tokenizer
tokenizer.train_from_iterator(texts, trainer=trainer)

tokenizer.train_from_iterator is where the learning happens. The tokenizer analyzes our text data and figures out the most frequent byte pairs to merge.

Training my tokenizer on Colab free tier (CPU) takes 1.5 hours.

We also need to set the byte_level decoder.

tokenizer.decoder = decoders.ByteLevel()

We set a ByteLevel decoder to ensure that when we decode token IDs back to text, the original formatting is preserved.

Let's make sure our special tokens have the correct IDs:

# Ensure special tokens are assigned correctly
assert tokenizer.token_to_id("<unk>") == 0
assert tokenizer.token_to_id("<s>") == 1
assert tokenizer.token_to_id("</s>") == 2

These assertions check that <unk> is 0, <s> is 1, and </s> is 2. This is crucial for consistency.

Now, let's save our trained tokenizer:

# Create the directory to save the tokenizer (if it doesn't exist)
tokenizer_save_dir = "my_custom_tokenizer"
os.makedirs(tokenizer_save_dir, exist_ok=True)

# Save the tokenizer model and configuration
tokenizer.save(os.path.join(tokenizer_save_dir, "tokenizer.json"))
tokenizer.model.save(tokenizer_save_dir)

We create a directory (my_custom_tokenizer) and save the tokenizer in two parts: tokenizer.json (vocabulary and merge rules) and the model files.

Finally, let's create a tokenizer_config.json file. This file provides extra configuration details so we can easily load the tokenizer later using Hugging Face's AutoTokenizer.

It's like a settings file for our tokenizer. We'll build this configuration step-by-step.

First, let's set some basic options:

config = {
    "add_bos_token": False,  # Don't automatically add a beginning-of-sequence token
    "add_eos_token": False,  # Don't automatically add an end-of-sequence token
    "add_prefix_space": False, # No space before the first word (we already set this)
    "bos_token": "<s>",      # Define the beginning-of-sequence token
    "eos_token": "</s>",     # Define the end-of-sequence token
    "unk_token": "<unk>",     # Define the unknown token
    "pad_token": "<unk>",      # Use <unk> for padding (we'll handle padding later)
}

These options control how the tokenizer handles special tokens and spacing.

We're setting the beginning-of-sequence (bos_token), end-of-sequence (eos_token), unknown (unk_token), and padding (pad_token) tokens explicitly.

Next, we'll add some more general settings:

config.update({
    "model_max_length": 32768,          # Maximum sequence length the model can handle
    "tokenizer_class": "PreTrainedTokenizerFast", # Specify the tokenizer class
    "clean_up_tokenization_spaces": False, # Don't clean up extra spaces
    "additional_special_tokens": [],     # No additional special tokens (for now)
    "spaces_between_special_tokens": False, # Don't add spaces between special tokens
    "sp_model_kwargs": {},                # No special model arguments
})

We're setting the maximum sequence length (model_max_length), specifying the tokenizer class (PreTrainedTokenizerFast for efficiency), and controlling how spaces are handled.

Now, let's define how the special tokens are represented in the decoder:

config["added_tokens_decoder"] = {
    "0": {"content": "<unk>", "lstrip": False, "normalized": False, "rstrip": False, "single_word": False, "special": True},
    "1": {"content": "<s>", "lstrip": False, "normalized": False, "rstrip": False, "single_word": False, "special": True},
    "2": {"content": "</s>", "lstrip": False, "normalized": False, "rstrip": False, "single_word": False, "special": True}
}

This added_tokens_decoder section maps the token IDs (0, 1, 2) of our special tokens to their properties. We're telling the decoder that these tokens are "special" and should not be treated like regular words (no leading/trailing spaces, no normalization).

Finally, and this is very important for conversational models, we'll add a chat template:

config["chat_template"] = """
{% if messages[0]['role'] == 'system' %}
    {% set system_message = messages[0]['content'] %}
    {{ '<s>system\\n' + system_message + '</s>\\n' }}
{% else %}
    {{ '<s>system\\nYou are a helpful AI assistant.</s>\\n' }}
{% endif %}
{% for message in messages %}
    {% set content = message['content'] %}
    {% if message['role'] == 'user' %}
        {{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}
    {% elif message['role'] == 'assistant' %}
        {{ content + '</s>' + '\\n' }}
    {% endif %}
{% endfor %}
"""

This Jinja2 template formats a chat for an LLM using <s> and </s> tokens to separate system, user, and assistant messages. It ensures the model understands the conversation flow.

Now that we've built our configuration, let's save it to a file:

with open(os.path.join(tokenizer_save_dir, "tokenizer_config.json"), "w", encoding="utf-8") as config_file:
    json.dump(config, config_file, ensure_ascii=False, indent=4)

print("Tokenizer training completed and saved successfully.")

We save the config dictionary as a JSON file named tokenizer_config.json in our tokenizer directory. The indent=4 makes the file human-readable.

And that's it! We have trained our custom BPE tokenizer

In the upcoming section, we will build our Transformer architecture from scratch.

Transformer Overview

The Transformer is the heart of modern Large Language Models. Unlike recurrent networks (like RNNs or LSTMs), which process text sequentially, Transformers use a mechanism called attention to process all parts of the input simultaneously.

This way we can have much greater parallelism and scalability, which leads to significantly faster training times and the ability to handle longer range dependencies in text.

Our Transformer model will consist of several key components:

  • Embedding Layer: Converts tokens to vectors
  • Positional Encoding: Adds position info
  • Transformer Blocks: Stack of layers
  • Multi-Head Attention: Focuses on different input parts
  • Feed Forward Network: Processes each token
  • RMSNorm: Normalizes layers
  • Output Layer: Generates token probabilities

The original Transformer architecture diagram is very difficult to understand, So I created a simplified one:

None

This diagram gives a high-level overview. A Transformer converts text into tokens, embeds them, and adds positional encoding. It then passes through stacked self-attention layers to capture context.

The output is normalized, transformed by a linear layer, and processed through softmax to generate probabilities for the next token or class.

Let's build each of its component one by one.

RMSNorm

None

RMSNorm (Root Mean Square Normalization) is a type of layer normalization. Normalization techniques are crucial in deep learning to stabilize training and improve performance.

RMSNorm normalizes the activations of each layer by dividing them by the root mean square of the activations. The addition of a small constant, epsilon (Ξ΅), ensures numerical stability (avoiding division by zero).

This helps prevent the activations from becoming too large or too small, making the training process smoother. It's simpler and faster than traditional LayerNorm because it omits the mean centering step.

Let's create our RMSNorm class:

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float):
        super().__init__()
        self.eps = eps  # Small constant for numerical stability
        self.weight = nn.Parameter(torch.ones(dim)) # Learnable scaling parameter

    def forward(self, x):
        # Calculate the root mean square (RMS) and normalize
        return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)

The __init__ method initializes the epsilon value and a learnable weight parameter (initialized to ones).

The forward method performs the normalization:

  • It calculates the root mean square of the input x.
  • Adds epsilon for numerical stability.
  • Takes the reciprocal of the square root.
  • Multiplies by the learnable weight.
  • The result is cast back to the original data type of x.

Positional Encoding

Since the Transformer processes the input sequence all at once, it doesn't inherently know the position of each token.

We need to add this positional information explicitly. This is where Positional Encoding comes in. We are using rotary positional embeddings.

None

Rotary Positional Embeddings (RoPE) represent token positions using rotations in a high-dimensional space.

The core idea is to encode the position of each token as a series of frequencies. These frequencies are then used to create a complex-valued embedding that rotates the token's representation in a way that depends on its position.

The precompute_pos_cis function calculates these embeddings; it is a helper function.

def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
    """Pre-computes the complex exponentials (cis) for rotary embeddings."""
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)  # Sequence positions
    freqs = torch.outer(t, freqs).float()  # Outer product of positions and frequencies
    pos_cis = torch.polar(torch.ones_like(freqs), freqs)  # Convert to complex exponentials
    return pos_cis

First, it calculates a set of frequencies based on the dimension (dim) and a scaling factor (theta).

Then, it creates a sequence of positions (t) and computes the outer product of t and the frequencies.

Finally, it converts these values to complex exponentials using torch.polar, resulting in the pos_cis tensor.

def apply_rotary_emb(xq, xk, pos_cis):
    """Apply rotary embeddings to query (xq) and key (xk) tensors."""
    def unite_shape(pos_cis, x):
        # Ensure pos_cis has the correct shape and reshape it to match x
        assert pos_cis.shape == (x.shape[1], x.shape[-1])
        shape = [d if i == 1 or i == x.ndim - 1 else 1 for i, d in enumerate(x.shape)]
        return pos_cis.view(*shape)

    # Convert to complex for rotation
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # Align pos_cis shape
    pos_cis = unite_shape(pos_cis, xq_)

    # Apply rotation and convert back
    xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)

The apply_rotary_emb function applies calculated rotary embeddings to queries (xq) and keys (xk) in attention mechanisms, encoding positional information. The unite_shape reshapes embeddings for compatibility.

xq and xk are transformed into complex numbers, multiplied by embeddings, and converted back, preserving the original data type.

Attention

None

Attention is the key mechanism of the Transformer. It allows the model to weigh the importance of different parts of the input sequence when processing each token.

Our implementation uses multi-head attention, which means we perform the attention mechanism multiple times in parallel (with different learned weights), allowing the model to capture different aspects of the relationships between tokens.

We also incorporate rotary positional embeddings (RoPE) to encode positional information.

Key-value caching is used for faster auto-regressive decoding (like in text generation).

class Attention(nn.Module):
    def __init__(self, args):
        super().__init__()
        # Define number of key-value heads (defaults to n_heads if not specified)
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        assert args.n_heads % self.n_kv_heads == 0  # Ensure heads divide evenly

        self.n_local_heads = args.n_heads
        self.n_local_kv_heads = self.n_kv_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads  # Replication factor
        self.head_dim = args.dim // args.n_heads  # Dimension per head

        # Linear projections for query, key, and value
        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)

        # Dropout layers
        self.attn_dropout = nn.Dropout(args.dropout)
        self.resid_dropout = nn.Dropout(args.dropout)
        self.dropout = args.dropout

        # Use Flash Attention if available (PyTorch 2.0+)
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn

        # Create causal mask to prevent attending to future tokens
        mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
        mask = torch.triu(mask, diagonal=1)  # Upper triangular mask
        self.register_buffer("mask", mask, persistent=False)

The __init__ method sets up the parameters:

  • n_kv_heads controls the number of "key-value heads" (which can be different from the number of "query heads" for techniques like grouped-query attention).
  • n_local_heads and n_local_kv_heads are the number of heads per device (relevant for multi-GPU setups).
  • n_rep is the repetition factor for key-value heads.
  • head_dim is the dimensionality of each attention head.

We create linear layers (wq, wk, wv, wo) to project the input into query, key, value, and output spaces.

We create dropout layers, check if optimized attention is implemented, and also have a causal mask that prevents attending to future tokens.

def forward(self,
            x: torch.Tensor,
            pos_cis: torch.Tensor,
            past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
            use_cache=False):
    # Get batch size, sequence length, and hidden dimensions
    bsz, seq_len, _ = x.shape

    # Compute query, key, and value projections
    xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

    # Reshape into (batch, seq_len, num_heads, head_dim)
    xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
    xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
    xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)

    # Apply rotary positional embeddings
    xq, xk = apply_rotary_emb(xq, xk, pos_cis)

The input x is projected into query (xq), key (xk), and value (xv) tensors. These are then reshaped to have separate heads and rotary positional embeddings (pos_cis) are applied to xq and xk.

# Key-Value Cache Implementation
if past_key_value is not None:
    # Concatenate past key/value states with current ones along sequence dimension
    xk = torch.cat([past_key_value[0], xk], dim=1)
    xv = torch.cat([past_key_value[1], xv], dim=1)

# Store updated key-value pairs if caching is enabled
past_kv = (xk, xv) if use_cache else None

If we're using key-value caching (for faster generation), the current xk and xv are concatenated with the cached values (past_key_value).

The past_kv tuple is created for the next iteration if use_cache is enabled.

# Transpose for (batch, heads, seq_len, head_dim)
xq, xk, xv = (
    xq.transpose(1, 2),
    repeat_kv(xk, self.n_rep).transpose(1, 2),  # Repeat KV heads and transpose
    repeat_kv(xv, self.n_rep).transpose(1, 2)  # Repeat KV heads and transpose
)

The query (xq), key (xk), and value (xv) tensors are transposed to have the shape (batch_size, num_heads, seq_len, head_dim).

The key and value tensors are repeated along the head dimension if n_rep is greater than 1 (for grouped-query attention).

if self.flash and seq_len != 1:
    # Use Flash Attention for efficiency if available and not processing a single token
    dropout_p = self.dropout if self.training else 0.0
    output = F.scaled_dot_product_attention(
        xq, xk, xv,
        attn_mask=None,  # Causal masking is handled by `is_causal=True`
        dropout_p=dropout_p,
        is_causal=True
    )
else:
    # Compute scaled dot-product attention manually
    scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)  # Scale by head_dim
    scores += self.mask[:, :, :seq_len, :seq_len]  # Apply causal mask
    scores = F.softmax(scores.float(), dim=-1).type_as(xq)  # Normalize scores
    scores = self.attn_dropout(scores)  # Apply dropout
    output = scores @ xv  # Compute weighted sum of values

If flash attention is available and it is not single token inference time, we can use scaled dot product attention. Otherwise we use causal attention and calculate manually.

# Reshape output back to (batch, seq_len, hidden_dim)
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)

# Apply output projection and residual dropout
output = self.resid_dropout(self.wo(output))

# Return final output and past key-value cache (if caching is enabled)
return output, past_kv

Finally, the output is reshaped, projected by wo, and dropout is applied. The function returns the output and the past_kv tuple (for caching).

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    """Repeat key-value heads `n_rep` times along the head dimension."""
    bs, slen, n_kv_heads, head_dim = x.shape  # Extract tensor dimensions
    if n_rep == 1:
        return x  # No repetition needed if n_rep is 1
    # Expand and repeat along the head dimension, then reshape
    return (
        x[:, :, :, None, :]  # Add a new axis for repetition
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)  # Repeat n_rep times
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)  # Flatten repeated heads
    )

The repeat_kv function is a helper function used to repeat the key and value tensors along the head dimension when using grouped-query attention.

FeedForward

None

The FeedForward network is a simple, fully connected network applied to each token's representation independently after the attention mechanism. It provides additional non-linearity and allows the model to learn more complex representations.

Our implementation uses the SiLU (Sigmoid Linear Unit) activation function.

The hidden dimension of the FFN is calculated based on the paper's suggestions.

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()

        # Set hidden dimension if not provided in config
        if config.hidden_dim is None:
            hidden_dim = 4 * config.dim
            hidden_dim = int(2 * hidden_dim / 3)
            config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)

        # Define the layers: w1 (input to hidden), w2 (hidden to output), w3 (input to hidden interaction)
        self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
        self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
        self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)

        # Dropout layer for regularization
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        # Apply the feedforward layers with SILU activation and dropout
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))

The __init__ method initializes three linear layers: w1, w2, and w3.

  • w1 and w3 project the input to a hidden dimension, and w2 projects it back to the original dimension.
  • The hidden dimension is calculated as described above.
  • A dropout layer is also initialized.

The forward method applies w1 the SiLU activation, w3, multiplies the results applies w2, and finally applies dropout.

Transformer Block

None

The Transformer Block is the fundamental building block of the Transformer model. It combines the Multi-Head Attention and FeedForward components, along with residual connections and layer normalization.

Residual connections (adding the input to the output of a sub-layer) help with training deep networks by mitigating the vanishing gradient problem.

class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, config: LMConfig):
        super().__init__()

        # Set model configuration parameters
        self.n_heads = config.n_heads
        self.dim = config.dim
        self.head_dim = config.dim // config.n_heads

        # Initialize attention and feed-forward components
        self.attention = Attention(config)

        # Layer normalization for attention and feed-forward outputs
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
        self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)

        # Feed-forward neural network
        self.feed_forward = FeedForward(config)

    def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
        # Apply attention with normalization and residual connection
        h_attn, past_kv = self.attention(
            self.attention_norm(x),  # Normalize input before attention
            pos_cis,
            past_key_value=past_key_value,
            use_cache=use_cache
        )

        # Add attention output to the input (residual connection)
        h = x + h_attn

        # Apply feed-forward network with normalization and residual connection
        out = h + self.feed_forward(self.ffn_norm(h))

        return out, past_kv

The __init__ method initializes the Attention and FeedForward modules, along with RMSNorm layers for both the attention and feed-forward components.

The forward method:

  1. First applies attention norm and multi-head self-attention.
  2. After that, a residual connection is applied.
  3. Then, FFN norm and the feed-forward network are applied, followed by another residual connection.

The forward method implements the sequence: Attention Norm -> Attention -> Residual Connection -> FFN Norm -> FeedForward -> Residual Connection.

The method returns the result and key-value cache.

Combining Transformer Components

Let's assemble all the components we have discussed into a complete Transformer model.

We'll start by defining the configuration for our model, storing all the hyperparameters in a dictionary.

# Configuration for our Transformer model
model_config = {
    "vocab_size": 6400,       # Size of the vocabulary
    "dim": 512,               # Dimensionality of the embeddings and hidden states
    "n_heads": 8,             # Number of attention heads
    "n_kv_heads": 2,          # Number of key-value heads (as specified in the LMConfig)
    "norm_eps": 1e-5,         # Epsilon for RMSNorm
    "dropout": 0.0,           # Dropout probability
    "max_seq_len": 1024,      # Maximum sequence length
    "rope_theta": 10000.0,    # Theta parameter for RoPE
    "multiple_of": 64,        # Used for hidden dimension calculation in FFN
    "hidden_dim": None,       # Hidden dimension of the FFN (calculated if None)
    "n_layers": 8,            # Number of Transformer blocks
    "flash_attn": True,       # Use flash attention if available
}

This dictionary, model_config, holds all the hyperparameters that define the structure and behavior of our Transformer. Using a dictionary like this is convenient because it makes it easy to change settings in one place.

Now, let's define our main Transformer model class.

class MyTransformer(PreTrainedModel):
    config_class = dict  # We'll use a dictionary for configuration

    def __init__(self, config: dict):
        super().__init__(config)  # Initialize PreTrainedModel
        self.config = config # store config
        self.vocab_size, self.n_layers = config["vocab_size"], config["n_layers"]

        # Token embeddings
        self.tok_embeddings = nn.Embedding(config["vocab_size"], config["dim"])

        # Dropout layer
        self.dropout = nn.Dropout(config["dropout"])

        # Transformer blocks
        self.layers = nn.ModuleList([TransformerBlock(l, config) for l in range(self.n_layers)])

        # Final normalization layer
        self.norm = RMSNorm(config["dim"], eps=config["norm_eps"])

        # Output layer (linear projection to vocabulary size)
        self.output = nn.Linear(config["dim"], config["vocab_size"], bias=False)

        # Tie the weights of the embedding and output layers
        self.tok_embeddings.weight = self.output.weight

        # Precompute rotary positional embeddings
        self.register_buffer(
            "pos_cis",
            precompute_pos_cis(dim=config["dim"] // config["n_heads"], theta=config["rope_theta"]),
            persistent=False,
        )
        self.OUT = CausalLMOutputWithPast()

The __init__ method initializes the model with the following components:

  • Embedding Layer: self.tok_embeddings converts token IDs to embeddings.
  • Dropout: self.dropout helps prevent overfitting.
  • Transformer Blocks: self.layers contains a stack of Transformer blocks.
  • Normalization: self.norm applies RMSNorm before output.
  • Output Layer: self.output projects the final hidden state to the vocabulary space.
  • Weight Tying: Embedding and output layer weights are tied for better performance.
  • Positional Embeddings: self.pos_cis stores rotary positional embeddings.
  • Output: Returns CausalLMOutputWithPast for logits, loss, and cache.

Now, let's define the forward method, which performs the complete forward pass through the model:

def forward(self,
                input_ids: Optional[torch.Tensor] = None,
                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                use_cache: bool = False,
                **args):

        past_key_values = past_key_values or [None] * len(self.layers)
        start_pos = args.get('start_pos', 0)

        # Apply dropout to the embeddings
        h = self.dropout(self.tok_embeddings(input_ids))

        # Get the appropriate positional embeddings for the current sequence length
        pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]

        # Iterate through the Transformer blocks
        past_kvs = []
        for l, layer in enumerate(self.layers):
            h, past_kv = layer(
                h, pos_cis,
                past_key_value=past_key_values[l],
                use_cache=use_cache
            )
            past_kvs.append(past_kv)

        # Apply the final normalization and output layer
        logits = self.output(self.norm(h))
        # aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward)) #we are not using moe for now.

        self.OUT.__setitem__('logits', logits)
        self.OUT.__setitem__('aux_loss', 0)  # we set 0 because we are not using moe for now
        self.OUT.__setitem__('past_key_values', past_kvs)
        return self.OUT

The forward method performs the following:

  • Embedding and Dropout: Input token IDs (input_ids) are passed through the embedding layer and dropout is applied.
  • Positional Embeddings: The precomputed positional embeddings (self.pos_cis) are sliced to match the sequence length.
  • Transformer Blocks: The model processes the token embeddings through a stack of Transformer blocks (self.layers), passing each block's output along with positional embeddings and optionally cached key-value pairs (past_key_value). Caching is controlled by the use_cache flag.
  • Normalization and Output: The output from the final block is normalized (self.norm) and passed through the output layer (self.output) to generate logits (unnormalized probabilities).
  • Auxiliary Loss and Key-Value Cache: The logits, auxiliary loss (set to 0), and key-value cache are stored in CausalLMOutputWithPast.

We have coded the entire Transformer architecture.

We are very close to training. In the upcoming section, we will create the training loop for the model and the data preprocessing for training.

Handling Pre-Training Dataset

The initial training of our LLM is essentially focused on learning to speak English, meaning it will learn proper grammar, punctuation, and word usage. The presence of meaningful sentences is not a concern at this stage.

You can focus on any dataset for pretraining the LLM, as long as it contains meaningful, end-to-end sentences. I am going to use DeepCtrl SFT Data for pretraining, which includes both Chinese and English datasets. However, we will focus solely on the English part, as our target is to create a perfect English-speaking model.

Let's download and take a look at a sample of our dataset.

!wget https://www.modelscope.cn/datasets/deepctrl/deepctrl-sft-data/resolve/master/sft_data_en.jsonl
file_path = "sft_data_en.jsonl" # Downloaded filepath (23 GB RAM REQUIRED!!!!!!!!!!!!!!!!!!!!!)

# Read the JSONL file with tqdm progress bar
data = []
with open(file_path, "r", encoding="utf-8") as file:
    # Iterate through each line of the JSONL file
    for line in tqdm(file, desc="Loading JSONL file"):
        # Parse each line as a JSON object and append to data
        data.append(json.loads(line))

The downloaded file is quite huge, it will require 23 GB of RAM. Please use Kaggle to avoid issues. Once the file is loaded, let's print a sample of it.

# Printing sample
data[23]

### OUTPUT ###
...
"input" : "what is the PEMDAS Rule when solving math..."
"output" : "PEMDAS is the famous mathematical rule to ..."
### OUTPUT ###

You may see other keys, but the "input" and "output" are the only useful parts for us. The dataset is essentially a natural conversation between a person and an AI chatbot.

The problem is that the data contains long inputs and outputs, which will consume too much memory during training. To avoid this, I will use only those entries where the combined length of the input and output is less than 512 characters. Let's proceed with that.

# Define the output file name
output_file = "petrain_data.jsonl"

# Write data to JSONL format
with open(output_file, "w", encoding="utf-8") as f:
    for item in data:
        # Check input + output length condition
        io_length = len(item['input']) + len(item['output'])
        if io_length < 512:  # Only keep valid entries
            json.dump({"text": f"{item['input']}\n{item['output']}"}, f, ensure_ascii=False)
            f.write("\n")  # Newline for each JSONL entry

print(f"New JSONL file saved as {output_file}")

It has around 660K entries, and we have saved our training data in the same JSONL format we used for training the tokenizer.

Next, we need to create a pretraining dataset class that will format our dataset based on the trained tokenizer.

class PretrainDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        super().__init__()  # Initialize parent class (Dataset)
        self.tokenizer = tokenizer  # Store tokenizer
        self.max_length = max_length  # Set max length for tokenization
        self.samples = self.load_data(data_path)  # Load dataset samples

    def load_data(self, path):
        samples = []  # Initialize list to store samples
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:  # Iterate through each line in the file
                data = json.loads(line.strip())  # Parse JSON data
                samples.append(data)  # Add sample to the list
        return samples  # Return loaded samples

    def __len__(self):
        return len(self.samples)  # Return the number of samples

    def __getitem__(self, index):
        sample = self.samples[index]  # Get the sample at the given index

        # Construct input text with BOS and EOS tokens
        text = f"{self.tokenizer.bos_token}{str(sample['text'])}{self.tokenizer.eos_token}"

        # Tokenize input text with padding and truncation
        encoding = self.tokenizer(
            text,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )

        input_ids = encoding.input_ids.squeeze()  # Get input token IDs
        loss_mask = (input_ids != self.tokenizer.pad_token_id)  # Mask for non-padding tokens

        # Prepare input (X) and target (Y) sequences
        X = torch.tensor(input_ids[:-1], dtype=torch.long)  # Exclude last token for X
        Y = torch.tensor(input_ids[1:], dtype=torch.long)   # Exclude first token for Y

        loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)  # Adjust loss mask

        return X, Y, loss_mask  # Return X, Y, and loss mask

In our pretraining class, we have set the maximum length to 512, since each input during training contains only 512 characters.

We are also wrapping our text with <s>training_input</s>, as I mentioned earlier. This is important because our trained tokenizer knows what the beginning and end tokens are. Wrapping each training data input around those tags will help the model learn where to start and where to end, using those tags as indicators.

Now, it's time to code the training loop.

We have created our pretraining dataset class, and now that we have everything, let's proceed.

Pre-Training Loop (29M LLM)

This is the section where our model will actually learn from the data. The basic idea is to:

  1. Load batches of data from our PretrainDataset.
  2. Feed the input sequences (X) to the model.
  3. Calculate the loss between the model's predictions and the target sequences (Y).
  4. Update the model's weights to minimize the loss.

First, let's define a helper function to calculate the learning_rate. We'll use a cosine learning rate schedule, which gradually decreases the learning_rate during training. This often leads to better results than using a fixed learning_rate.

def get_lr(current_step, total_steps, lr):
    """Calculates the learning rate using a cosine schedule."""
    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))

This function takes the current training_step, the total number of steps, and the initial learning_rate (lr). It calculates the learning_rate for the current step using a cosine function, and a little bit constant learning_rate.

Now, let's set up our model, optimizer, and data_loader:

# Instantiate the model
model = MyTransformer(model_config)
model = model.to("cuda") # Move the model to the GPU if available

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained("my_custom_tokenizer")

# Create the training dataset
train_ds = PretrainDataset("petrain_data.jsonl", tokenizer, max_length=model_config["max_seq_len"])

# Create the data loader
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, drop_last=False)

# Define the optimizer (AdamW is a good choice)
optimizer = optim.AdamW(model.parameters(), lr=5e-4)

We instantiate our MyTransformer model, load our trained tokenizer, create a PretrainDataset instance, and set up a DataLoader.

The DataLoader will handle shuffling and batching the data for us. We also define our optimizer, AdamW, which is a variant of Adam that often works better for training Transformers. We moved our model to GPU to make our training faster.

Next, we define our loss_function. Since we're training a language_model to predict the next token in a sequence, we'll use cross-entropy loss:

# Define the loss function (cross-entropy)
loss_fct = nn.CrossEntropyLoss(reduction='none')

CrossEntropyLoss combines the softmax activation and the negative log-likelihood loss. We set reduction='none' because we want to calculate the loss for each token separately (we'll handle the averaging ourselves).

Now, the main training loop .

epochs = 1 # You can adjust the number of epochs
iter_per_epoch = len(train_loader) #how many batches are there
accumulation_steps = 8 #for gradient accumulation
grad_clip = 1.0 #for gradient clipping

for epoch in range(epochs):
    start_time = time.time()
    for step, (X, Y, loss_mask) in enumerate(train_loader):
        X = X.to("cuda")
        Y = Y.to("cuda")
        loss_mask = loss_mask.to("cuda")

        # Calculate the learning rate for the current step
        lr = get_lr(epoch * iter_per_epoch + step, epochs * iter_per_epoch, 5e-4)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # Forward pass
        res = model(X)

        # Calculate the loss
        loss = loss_fct(
            res.logits.view(-1, res.logits.size(-1)),
            Y.view(-1)
        ).view(Y.size())
        loss = (loss * loss_mask).sum() / loss_mask.sum()
        loss = loss/accumulation_steps # divided by accumulation steps

        # Backward pass
        loss.backward()

        if (step + 1) % accumulation_steps == 0:
            # Gradient Clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            # Update the model's parameters
            optimizer.step()
            # Reset the gradients
            optimizer.zero_grad()

        if step % 100 == 0:
            spend_time = time.time() - start_time
            print(
                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
                    epoch + 1,
                    epochs,
                    step,
                    iter_per_epoch,
                    loss.item() * accumulation_steps,
                    optimizer.param_groups[-1]['lr'],
                    spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))

Let's break this down:

  • Epochs & Batches: Loop through the dataset in batches.
  • Data to GPU: Move data to GPU for faster processing.
  • Forward Pass & Loss: Get model predictions and calculate loss.
  • Backward Pass & Gradients: Compute gradients and accumulate them.
  • Update & Clipping: Clip gradients and update model weights.
  • Zero Gradients & Logging: Reset gradients and print metrics.

Once you run the code, It will start printing training loss and more.

...

Epoch:[2/10](30/100) loss:0.256 lr:0.000012345679 epoch_Time:13min

...

Finally, let's add code to save our trained model:

model.eval() #set to evaluation mode
torch.save(model.state_dict(), "pretrain_model.pth") #save state dict
model.train()# set back to training mode.

Once the training got completed.

we will have a pretrained model with 26–30 million parameters.

You can definitely increase the number of parameters by modifying the configuration settings.

PyTorch Model to SafeTensors

Now that we have trained our model, let's convert it into safetensors format.

The safetensors format is a safer and faster alternative to PyTorch's default pickle-based format for storing tensors.

First, let's define a function to perform the conversion.

def convert_and_save_model(torch_path, transformers_path):
    """Converts a PyTorch model to Transformers format and saves it."""

    # Load the PyTorch model's state_dict
    state_dict = torch.load(torch_path, map_location="cpu")

    # Create an instance of your model (MyTransformer)
    model = MyTransformer(model_config)

    # Load the state_dict into the model
    model.load_state_dict(state_dict, strict=False)
    # Print the number of parameters
    model_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f'Model parameters: {model_params / 1e6:.3f} Million')

    # Save the model in Transformers format
    model.save_pretrained(transformers_path, safe_serialization=False) # set safe_serialization=True to save in safetensors format

    # Export the tokenizer
    export_tokenizer(transformers_path)
    print(f"Model saved in Transformers format to: {transformers_path}")

t loads the state dict using torch.load, instantiates the MyTransformer model, loads the weights (strict=False for flexibility), and prints the total parameters. The model is saved with safe_serialization=False.

Now, let's call this function to perform the conversion:

# Convert and save the model
convert_and_save_model("pretrain_model.pth", "my_transformer_hf")

This will create a directory called my_transformer_hf containing the model files, ready to be used with the Hugging Face transformers library.

Inferencing Pre-Train LLM

Let's run inference on our model and see how it responds based on our input. However, don't expect too much, it is just a pretrained model.

# Load model and tokenizer
model, tokenizer = (AutoModelForCausalLM.from_pretrained("MiniMind2", trust_remote_code=True).to(device).eval(),
                    AutoTokenizer.from_pretrained("MiniMind2", use_fast=False, trust_remote_code=True))

# Generate response from model
def generate_response(prompt, max_new_tokens=512, temperature=0.85, top_p=0.9):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
    output_ids = model.generate(input_ids, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p,
                                eos_token_id=tokenizer.eos_token_id)
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

Let's try different inputs to see how the model responds.

user input: Hi
LLM: I'm sorry, but you haven't provided the information about
the specific data or issues you're using.

user input: how are you?
LLM: Yes, I'm a DAY GPA. However, I can help you with that.

user input: what is 2+3=?
LLM: When you're looking for a 2+3=)

You can see that the LLM responses contain proper words and shorter contexts that make sense. However, we want the LLM to respond more like ChatGPT.

For example, when we ask, What is 2 + 3?, it should correctly respond with 5 .

SFT Dataset Preprocessing

To make our LLM knowledge-aware, we need a high-quality dataset with a conversational format. You can choose any dataset, but it should follow a conversational style like this:

{
    "conversations": [
        {"role": "user", "content": "Hello"},
        {"role": "assistant", "content": "Hello!"},
        {"role": "user", "content": "Goodbye"},
        {"role": "assistant", "content": "Goodbye!"}
    ]
}

I am using the GooAQ dataset provided by Google, which contains 5 million question-answer pairs.

Let's download the data and check a sample of it.

!wget -O gooaq.jsonl "https://github.com/allenai/gooaq/raw/refs/heads/main/data/gooaq.jsonl"
file_path = "gooaq.jsonl"  # Replace with your file path
data = []

# Open file and read line by line with a progress bar
with open(file_path, 'r', encoding='utf-8') as f:
    for line in tqdm(f, desc="Loading JSONL"):
        try:
            data.append(json.loads(line))  # Parse JSON and add to list
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON: {e}")  # Handle decoding errors

# Print sample
print(data[0])

### OUTPUT ###
{'id': 1,
 'question': 'Who is the president of U.S?',
 'short_answer': "Trump",
 'answer': None,
 'answer_type': 'unknown',
 'answer_url': None
}

The format is very easy to understand, but we need to convert it into a conversational format first. Let's do that.

# Process the data
converted_data = []
for item in tqdm(data_list, desc="Processing Data"):
    question = item.get("question")
    answer = item.get("answer") or item.get("short_answer")  # Use short_answer if answer is missing

    if question and answer:  # Skip if both are empty
        conversation = {
            "conversations": [
                {"role": "user", "content": question},
                {"role": "assistant", "content": answer}
            ]
        }
        converted_data.append(conversation)

# Save to a JSONL file
output_file = "sft_data.jsonl"
with open(output_file, "w", encoding="utf-8") as f:
    for entry in tqdm(converted_data, desc="Writing to File"):
        f.write(json.dumps(entry, ensure_ascii=False) + "\n")  # Write each entry as a JSON line

print(f"Converted data saved to {output_file}")

We are checking if a long answer does not exist, if so, we use the short answer instead. The training data is then saved in JSONL format, the same format we used before.

We need to create an SFTDataset class to handle this dataset, similar to how we created the pretraining dataset class. Let's do that first and check how it works.

The goal of this class is to take our conversational data (in JSONL format), format it correctly using the chat template (remember the one we defined in the tokenizer config?), and convert it into numerical inputs that our model can understand.

First, let's define the basic structure of our class and the __init__ method:

class SFTDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_length=1024):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = self.load_data(jsonl_path)
        self.bos_id = tokenizer('<s>assistant\n', add_special_tokens=False).input_ids
        self.eos_id = tokenizer('</s>\n', add_special_tokens=False).input_ids

In __init__:

  • Store tokenizer and max_length.
  • Load data from jsonl_path into self.samples.
  • Store beginning and end of sequence IDs.

Now, let's define the load_data method:

This method simply reads the JSONL file line by line, parses each line as a JSON object, and appends it to the samples list.

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

The __len__ method returns the length of samples.

Now, the most important part: the _create_chat_prompt method. This is where we use the chat_template from our tokenizer to format the conversations correctly.

def _create_chat_prompt(self, conversations):
        """Builds a dialogue in ChatML format."""
        messages = []
        for i, turn in enumerate(conversations):
            role = 'user' if i % 2 == 0 else 'assistant'
            messages.append({"role": role, "content": turn['content']})
        return self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False
        )

This method processes a list of conversation turns (e.g., [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello"}]) and:

  • Iterates through the turns.
  • Formats them into a structured list, ensuring roles alternate between "user" and "assistant."
  • Uses self.tokenizer.apply_chat_template to apply the predefined chat template, inserting special tokens (e.g., <s> and </s>) to separate turns and roles.
  • Sets tokenize=False to return a formatted string instead of token IDs.

Now, let's create _generate_loss_mask method.

def _generate_loss_mask(self, input_ids):
    loss_mask = [0] * len(input_ids)  # Initialize loss mask with zeros
    i = 0
    while i < len(input_ids):
        if input_ids[i:i + len(self.bos_id)] == self.bos_id:  # Check for BOS token
            start = i + len(self.bos_id)
            end = start
            while end < len(input_ids):
                if input_ids[end:end + len(self.eos_id)] == self.eos_id:  # Find EOS token
                    break
                end += 1
            # Mark tokens after BOS until EOS for loss calculation
            for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
                loss_mask[j] = 1
            i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)  # Move index past EOS
        else:
            i += 1
    return loss_mask  # Return the generated loss mask

In this function we creates a mask to indicate which tokens should be considered during loss calculation. The goal is to only calculate the loss on the assistant's responses, not on the user's prompts or the special tokens.

It iterates to find the assistant and eos tokens, then assigns 1s to the assistant and eos tokens and 0s to others.

Finally, the __getitem__ method, which is called when we access an item from the dataset (like dataset[i]):

def __getitem__(self, index):
    sample = self.samples[index]  # Get the sample at the given index
    prompt = self._create_chat_prompt(sample['conversations'])  # Generate chat prompt
    encoding = self.tokenizer(prompt, max_length=self.max_length, truncation=True, return_tensors="pt")
    input_ids = encoding.input_ids.squeeze()  # Extract tokenized input IDs

    loss_mask = self._generate_loss_mask(input_ids.tolist())  # Generate loss mask

    # Append EOS token and update loss mask
    input_ids = input_ids.tolist() + [self.tokenizer.eos_token_id]
    loss_mask = loss_mask + [1]

    # Truncate or pad input to match max_length
    if len(input_ids) > self.max_length:
        input_ids = input_ids[:self.max_length]
        loss_mask = loss_mask[:self.max_length]
    else:
        pad_len = self.max_length - len(input_ids)
        input_ids += [self.tokenizer.pad_token_id] * pad_len
        loss_mask += [0] * pad_len

    # Create input (X) and target (Y) tensors for training
    X = torch.tensor(input_ids[:-1], dtype=torch.long)  # Input sequence
    Y = torch.tensor(input_ids[1:], dtype=torch.long)  # Shifted target sequence
    loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)  # Loss mask for target tokens

    return X, Y, loss_mask  # Return input, target, and loss mask tensors

Here's what's happening in__getitem__

  1. Get Sample β†’ sample = self.samples[index] retrieves the conversation at the given index.
  2. Create Prompt β†’ self._create_chat_prompt(...) formats the conversation.
  3. Tokenize β†’ self.tokenizer(...) tokenizes the prompt, truncating/padding to max_length, returning PyTorch tensors (return_tensors="pt").
  4. Loss Mask & EOS β†’ Generates a loss mask and appends 1 for the EOS token.
  5. Padding/Truncation β†’ Ensures input_ids match max_length by truncating if too long or padding if too short.
  6. Create X & Y β†’ X is the input sequence; Y is X shifted right by one position for next-token prediction.
  7. Return β†’ Outputs X, Y, and loss_mask.

Now that we have correctly preprocessed our SFT training dataset, it's time to code the training loop.

SFT Training Loop

Now that we have our SFTDataset ready, we can train our model using Supervised Fine-Tuning (SFT).

The SFT training loop is very similar to the pre-training loop, but we're using conversational data and focusing the loss on the assistant's responses.

First, let's set up the model, optimizer, and data loader. This is mostly the same as before, but we'll load our pre-trained model weights:

# Load the tokenizer (same as before)
tokenizer = AutoTokenizer.from_pretrained("my_custom_tokenizer")

# Load the PRE-TRAINED model
model = MyTransformer(model_config)
model.load_state_dict(torch.load("pretrain_model.pth", map_location="cpu")) # Load weights!
model = model.to("cuda")

We load our custom tokenizer. Crucially, we instantiate our MyTransformer model and then load the weights from pretrain_model.pth using load_state_dict. We load to the CPU first and then move the model to the GPU.

Now, let's create the SFTDataset and DataLoader:

# Create the SFT dataset
train_ds = SFTDataset("sft_data.jsonl", tokenizer, max_length=model_config["max_seq_len"])

# Create the data loader
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, drop_last=False)

We create an instance of our SFTDataset using the path to our SFT data (sft_data.jsonl), our tokenizer, and the maximum sequence length. The DataLoader is set up similarly to pre-training, with shuffling enabled.

Next, the optimizer and loss function:

# Define the optimizer (AdamW is a good choice)
optimizer = optim.AdamW(model.parameters(), lr=5e-5) # Potentially different LR

# Define the loss function (cross-entropy)
loss_fct = nn.CrossEntropyLoss(reduction='none')

We use AdamW, but the learning rate for SFT is often lower than for pre-training. I've used 5e-5 here, but you should experiment. We're still using CrossEntropyLoss with reduction='none'.

Now, let's define the learning rate scheduler. We'll use the same cosine schedule as before:

def get_lr(current_step, total_steps, lr):
    """Calculates the learning rate using a cosine schedule."""
    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))

This function remains unchanged from pretraining.

Now, the core of the training loop:

for epoch in range(epochs):
    start_time = time.time()
    for step, (X, Y, loss_mask) in enumerate(train_loader):
        X = X.to("cuda")
        Y = Y.to("cuda")
        loss_mask = loss_mask.to("cuda")

        # Calculate learning rate
        lr = get_lr(epoch * iter_per_epoch + step, epochs * iter_per_epoch, 5e-5)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # Forward pass
        res = model(X)

We start the outer loop for epochs and the inner loop for batches. We move our input, target, and loss mask to GPU.

We calculate the learning rate for the current step and set it to the optimizer. The model then makes its predictions (res = model(X)).

Now, the crucial part: the loss calculation. This is where we use the loss_mask:

# Calculate the loss (using the loss_mask!)
        loss = loss_fct(
            res.logits.view(-1, res.logits.size(-1)),
            Y.view(-1)
        ).view(Y.size())
        loss = (loss * loss_mask).sum() / loss_mask.sum() # Key difference: using loss_mask
        loss = loss/accumulation_steps

We compute the cross-entropy loss, mask it to focus on the assistant's tokens, sum the masked losses, and normalize by the mask sum. Gradient accumulation is applied before the backward pass and parameter update.

Next, the backward pass and parameter update:

# Backward pass
        loss.backward()
        if (step + 1) % accumulation_steps == 0:
            # Gradient Clipping
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            # Parameter update
            optimizer.step()
            # Zero gradients
            optimizer.zero_grad()

The loss.backward() computes gradients. If accumulation steps are reached, we clip gradients using clip_grad_norm_, update model parameters with optimizer.step(), and clear gradients with optimizer.zero_grad().

Finally, the logging and saving:

if step % 100 == 0:
            spend_time = time.time() - start_time
            print(
                'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.12f} epoch_Time:{}min:'.format(
                    epoch + 1,
                    epochs,
                    step,
                    iter_per_epoch,
                    loss.item() * accumulation_steps,
                    optimizer.param_groups[-1]['lr'],
                    spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))

model.eval() #set to evaluation mode
torch.save(model.state_dict(), "sft_model.pth") # Save the SFT model
model.train() #set to training mode

We are logging the loss, learning rate, and training time. When we start training, the loss will be printed in the same way as it was during the pretraining stage.

...

Epoch:[5/10](30/100) loss:0.514 lr:0.00002342352 epoch_Time:21min

...

After training, our model weights are saved. We can easily convert them into safetensors using our previous function. Now, let's load the model and see how much it has improved.

EvaluationπŸ˜–

Although it's not ideal to judge our LLM on benchmarks since it only has 30 million parameters, let's do it anyway.

We first created a pretrained model and then fine-tuned it using SFT. Further improvements can be made with RLHF and reasoning, which I will cover in future blogs.

I have also created a Streamlit web app (available in the GitHub repo) that allows you to load and chat with the model for easier communication.

Let's test different inputs and see how it responds.

Addition problem.

user input: 2+3=?

LLM: When you run the same number and the same number,
the number is 2 + 3 = 5.

So, 2 + 3 = 5.

Therefore, 2 + 3 = 5.

Knowledge check.

user input: What is the highest mountain in the world?

LLM: Mount Everest, the world's highest mountain (8,848.86m),
lies in the Himalayas on the Nepal-China border.
user input: Who discovered gravity?

LLM: Newton proposed the law of universal gravitation in 1687
explaining how objects attract each other through gravity.

General Question.

user input: hi
LLM: Hello

user: how are you?
LLM: I'm sorry, but I'm not able to assist you.
However, I'm unable to provide information about the
specific requirements or context, but I can provide
you with a general idea of how you might have:

Some Issues.

user input: 1 KG is how many grams?
LLM: The answer is `3`
user input: 1 KG is how many grams?
LLM: The answer is `3`

Though these answers are cherry picked, the model still requires a lot more pretraining and SFT training to improve further.

Conclusion

If you want to build a serious LLM, you should use a very high-quality synthetic dataset. The same approach applies to domain-specific LLMs.

Make sure to create an LLM with around 100 to 200 million parameters, as they perform well in question-answering tasks.

Happy reading!

None