Building NanoChatLM: How I Built a 28M Parameter Llama-3 Style LLM From Scratch for Free

There is a massive difference between calling an API like OpenAI and actually understanding how a Large Language Model works under the hood. I wanted to see if I could build a mini-LLM completely from scratch implementing the exact same architectural principles used by state-of-the-art models like Llama 3 without spending a single dollar.

The result is NanoChatLM, a 28-million parameter language model built from the ground up using Python and PyTorch.

If you are an aspiring AI engineer or a software developer curious about the inner workings of deep learning, here is the end-to-end breakdown of how I built it, how I structured the code blocks for readability, and the crucial architectural flaws I had to navigate along the way.

The Stack: Engineering on a Zero-Dollar Budget

Training an AI model usually requires thousands of dollars in high-end enterprise hardware. Since I am working on a thin laptop that does not have a massive discrete GPU, I had to get creative with free cloud infrastructure.

Here is the exact zero-cost ecosystem I used:

• Google Colab (Free T4 GPU Tier): Colab provided me with a remote virtual machine hooked up to an Nvidia Tesla T4 graphics card. By routing my PyTorch code to run over Nvidia’s CUDA software layer, I was able to offload millions of matrix math operations to thousands of parallel GPU cores instead of relying on my local CPU.
• Hugging Face Datasets: I pulled the OpenAssistant (OASST1) conversational dataset entirely for free using Hugging Face’s open-source Python library.
• Tiktoken (OpenAI’s Tokenizer): To turn raw human text strings into digital integer IDs that a machine can digest, I used OpenAI’s byte-pair encoding library.

Breaking Down the Implementation Step by Step

To make the codebase highly readable for deployment and open-source contributions, I structured the entire Colab notebook into 6 clean, modular blocks. Breaking it down this way ensures that anyone reviewing the project can understand the chronological flow of data and compute.

Block 1: Environment Setup & Hardware Verification

This initial block checks the Google Colab physical GPU attachment using standard Nvidia system management tools. It mounts Google Drive to permanently save your progress and installs the required data processing frameworks.

import os
import math
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import tiktoken
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
from google.colab import drive

# Verify physical hardware acceleration
!nvidia-smi

# Mount persistent storage
drive.mount('/content/drive')
os.makedirs('/content/drive/MyDrive/NanoChat', exist_ok=True)

# Install required frameworks
!pip install datasets tiktoken -q

# Initialize compute device targets
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Target Execution Device: {device}")

Block 2: Architectural Configuration & Data Pipeline

This segment structures the model configuration parameters, pulls the OpenAssistant conversational text data from Hugging Face, flattens the conversational dataset sequences, and tokenizes strings into integer arrays.

class Config:
    vocab_size   = 50257  # GPT-2 Tokenizer vocabulary count
    block_size   = 256   # Context window length (scaled down for T4 VRAM stability)
    n_layer      = 6      # Total Transformer Blocks stacked vertically
    n_head       = 6      # Attention Query heads
    n_embd       = 384    # Channel hidden dimensions size
    n_kv_heads   = 2      # Grouped-Query Attention key/value head pairs
    dropout      = 0.1    # Overfitting regularizer drop probability
    batch_size   = 32     # Sequence layout count per optimization phase
    max_iters    = 5000   # Optimization loops
    eval_every   = 500    # Progress evaluation validation checks frequency
    lr           = 3e-4   # Peak optimization learning rate
    weight_decay = 0.1    # L2 structural regularizer factor
    grad_clip    = 1.0    # Exploding gradient safety threshold

cfg = Config()

# Extract and Flatten Conversational Dataset Sequences
ds = load_dataset("OpenAssistant/oasst1")

def flatten_oasst(split):
    text = ""
    for row in split:
        if row["text"]:
            text += row["text"].strip() + "\n"
    return text

train_text = flatten_oasst(ds["train"])
val_text   = flatten_oasst(ds["validation"])

# Apply Tokenization Vector Processing
enc = tiktoken.get_encoding("gpt2")
train_ids = enc.encode(train_text, allowed_special={"<|endoftext|>"})
val_ids   = enc.encode(val_text,   allowed_special={"<|endoftext|>"})

# Construct Optimized Long-Integer Tensors
train_data = torch.tensor(train_ids, dtype=torch.long)
val_data   = torch.tensor(val_ids,   dtype=torch.long)
print(f"Dataset Loaded -> Train: {len(train_data):,} | Val: {len(val_data):,}")

Block 3: Custom Transformer Architecture Modules

This block houses the specific mathematical mechanics that power today’s leading commercial models. Rather than using legacy vanilla Transformers, I coded the model to replicate Llama-style design patterns:

1. Grouped-Query Attention (GQA): Instead of giving every Query head its own unique memory buffer, GQA groups multiple Query heads to share a single Key/Value head. This drastically speeds up text generation and saves memory.
2. Rotary Position Embedding (RoPE): Instead of adding fixed numbers to tracks sequence order, RoPE rotates the mathematical vectors in geometric space based on their position index, naturally capturing relative word distances.
3. RMSNorm: Traditional normalization layers calculate both mathematical mean and variance. RMSNorm drops the mean calculation entirely, saving massive computational overhead without hurting structural stability.
4. SwiGLU Activation: It uses gated multiplication driven by the SiLU (Swish) function to help the neural network map out sharp, non-linear logical paths.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps    = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm * self.weight

def precompute_rope(head_dim, seq_len, base=10000):
    theta = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    pos   = torch.arange(seq_len).float()
    freqs = torch.outer(pos, theta)
    return torch.cos(freqs).to(device), torch.sin(freqs).to(device)

def apply_rope(x, cos, sin):
    B, H, T, D = x.shape
    x1, x2 = x[..., :D//2], x[..., D//2:]
    return torch.cat([x1*cos[:T] - x2*sin[:T],
                      x1*sin[:T] + x2*cos[:T]], dim=-1)

class GQAttention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_head   = cfg.n_head
        self.n_kv     = cfg.n_kv_heads
        self.head_dim = cfg.n_embd // cfg.n_head
        self.groups   = cfg.n_head // cfg.n_kv_heads

        self.q_proj  = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
        self.k_proj  = nn.Linear(cfg.n_embd, cfg.n_kv_heads * self.head_dim, bias=False)
        self.v_proj  = nn.Linear(cfg.n_embd, cfg.n_kv_heads * self.head_dim, bias=False)
        self.o_proj  = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
        self.drop    = nn.Dropout(cfg.dropout)

        cos, sin = precompute_rope(self.head_dim, cfg.block_size)
        self.register_buffer("cos", cos)
        self.register_buffer("sin", sin)
        self.register_buffer("mask",
            torch.tril(torch.ones(cfg.block_size, cfg.block_size))
            .view(1, 1, cfg.block_size, cfg.block_size))

    def forward(self, x):
        B, T, C = x.shape
        H, Hkv, D = self.n_head, self.n_kv, self.head_dim

        q = self.q_proj(x).view(B, T, H,   D).transpose(1, 2)
        k = self.k_proj(x).view(B, T, Hkv, D).transpose(1, 2)
        v = self.v_proj(x).view(B, T, Hkv, D).transpose(1, 2)

        q = apply_rope(q, self.cos, self.sin)
        k = apply_rope(k, self.cos, self.sin)

        k = k.repeat_interleave(self.groups, dim=1)
        v = v.repeat_interleave(self.groups, dim=1)

        att = (q @ k.transpose(-2, -1)) * (D ** -0.5)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
        att = self.drop(F.softmax(att, dim=-1))

        out = (att @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)

class SwiGLU(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        hidden  = int(n_embd * 8 / 3)
        self.w1 = nn.Linear(n_embd, hidden, bias=False)
        self.w2 = nn.Linear(hidden, n_embd, bias=False)
        self.w3 = nn.Linear(n_embd, hidden, bias=False)
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.n_embd)
        self.norm2 = RMSNorm(cfg.n_embd)
        self.attn  = GQAttention(cfg)
        self.mlp   = SwiGLU(cfg.n_embd)
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

class NanoChatLM(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.drop    = nn.Dropout(cfg.dropout)
        self.blocks  = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.norm    = RMSNorm(cfg.n_embd)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear) or isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, 0.0, 0.02)

    def forward(self, idx, targets=None):
        x = self.drop(self.tok_emb(idx))
        for block in self.blocks:
            x = block(x)
        x = self.norm(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens=100, temperature=0.8, top_k=40):
        for _ in range(max_new_tokens):
            idx_cond  = idx[:, -cfg.block_size:]
            logits, _ = self(idx_cond)
            logits    = logits[:, -1, :] / temperature
            v, _      = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = float('-inf')
            probs     = F.softmax(logits, dim=-1)
            next_tok  = torch.multinomial(probs, num_samples=1)
            idx       = torch.cat([idx, next_tok], dim=1)
        return idx

model = NanoChatLM(cfg).to(device)
params = sum(p.numel() for p in model.parameters())
print(f"Model Initialized. Parameters: {params/1e6:.2f}M")

Block 4: Training Engine & Optimization Loop

This block coordinates the actual training phase. It handles data slicing into random mini-batches, handles gradient clipping to prevent exploding math states, and runs backpropagation using the AdamW optimizer and a Cosine Annealing learning rate scheduler. It saves the optimized parameters into a checkpoint file.

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
torch.cuda.empty_cache()

optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay, betas=(0.9, 0.95))
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.max_iters, eta_min=cfg.lr/10)

def get_batch(split):
    data = train_data if split == "train" else val_data
    ix   = torch.randint(len(data) - cfg.block_size, (cfg.batch_size,))
    x    = torch.stack([data[i:i+cfg.block_size] for i in ix])
    y    = torch.stack([data[i+1:i+cfg.block_size+1] for i in ix])
    return x.to(device), y.to(device)

@torch.no_grad()
def eval_loss():
    model.eval()
    losses = [model(*get_batch("val"))[1].item() for _ in range(20)]
    model.train()
    return np.mean(losses)

print("Starting 5,000 Step Training Loop...")
model.train()
for step in range(cfg.max_iters):
    x, y    = get_batch("train")
    _, loss = model(x, y)

    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
    optimizer.step()
    scheduler.step()

    if step % cfg.eval_every == 0 or step == cfg.max_iters - 1:
        vl = eval_loss()
        print(f"Step {step:5d} / {cfg.max_iters} | Train Loss: {loss.item():.4f} | Val Loss: {vl:.4f}")

torch.save(model.state_dict(), '/content/drive/MyDrive/your_project_folder/model.pt')
with open('/content/drive/MyDrive/your_project_folder/config.json', 'w') as f:
    json.dump({
        "vocab_size": cfg.vocab_size, "block_size": cfg.block_size, 
        "n_layer": cfg.n_layer, "n_head": cfg.n_head, 
        "n_embd": cfg.n_embd, "n_kv_heads": cfg.n_kv_heads, "dropout": cfg.dropout
    }, f)
print("Model weights exported safely to Google Drive.")

Block 5: Staging & Git Synchronization

This step automates pushing the code cleanly to GitHub. Note that instead of exposing private security strings in plain text, I updated the logic to pull from Colab’s internal Secrets vault securely.

from google.colab import userdata

GITHUB_USER  = "YOUR_GITHUB_USERNAME"
REPO_NAME    = "YOUR_REPOSITORY_NAME"
GITHUB_TOKEN = userdata.get('GH_TOKEN') 

!git config --global credential.helper store
!git config --global url."@github.com/">https://{GITHUB_TOKEN}@github.com/".insteadOf "https://github.com/"
!git clone https://github.com/{GITHUB_USER}/{REPO_NAME}.git

os.chdir(f"/content/{REPO_NAME}")

# Write standard runtime configurations files
with open("requirements.txt", "w") as f:
    f.write("torch\ntiktoken\ndatasets\n")

with open(".gitignore", "w") as f:
    f.write("model.pt\n__pycache__/\n*.pyc\n.env\n")

!git add .
!git commit -m "Deployment optimization — structured codeblocks layout implementation"
!git push origin main
print("Files pushed to production GitHub repository successfully!")

Block 6: Inference Runner

This script restores the model architecture configuration from a dry structure state, loads the specialized saved checkpoint matrix parameters back into the GPU, and uses top-k sampling with temperature scaling to generate new context text

inference_model = NanoChatLM(cfg).to(device)
inference_model.load_state_dict(torch.load('/content/drive/MyDrive/NanoChat/model.pt', map_location=device))
inference_model.eval()

prompt = "What is machine learning?"
tokens = enc.encode(prompt)
idx    = torch.tensor([tokens], dtype=torch.long).to(device)

output_sequence = inference_model.generate(idx, max_new_tokens=120, temperature=0.6, top_k=30)
print(enc.decode(output_sequence[0].tolist()))

The Catch: Building an “Empty Brain”

One of the coolest epiphanies I had during this project was realizing what actually happens when you initialize a model layer file.

When you instantiate NanoChatLM, you are not creating an intelligent machine. You are creating a structurally perfect, but completely empty network shell. To prevent the math from collapsing, the system automatically loops through every linear projection and embedding layer to fill blank matrix rows with fractional random numbers.

At that stage, the model has its physical brain structure, but if you give it a prompt like "What is machine learning?", it outputs complete gibberish. It requires the training phase in Block 4 to read the dataset and progressively modify those random variables into text representations.

Technical Flaws & Lessons Learned

Building a neural network from absolute scratch means you are going to hit walls. My model did not compile flawlessly on the first try. If you are building your own, watch out for these major pitfalls I encountered:

1. The Dataset Formatting Trap

Initially, my code tried to read lines from the OpenAssistant dataset as standard, flat text chunks. However, conversational data tracks explicit historical parameters like user queries and assistant turn boundaries. Mashing them together directly meant the model struggled to differentiate context shifts, occasionally trying to answer its own questions during inference.

2. Missing Configuration Map Handlers

When decoupling modules from standalone script cells into clean individual deployment objects, it is easy to leave behind context bindings. I hit multiple runtime crashes where the optimization loops expected specific model settings variables that had not been cleanly imported or passed through class definitions.

3. The Lack of System Prompt Wrappers

Right now, the architecture treats text generation as a raw next-token completion engine. Because I have not yet implemented a dedicated System Prompt Wrapper (the special conditioning sequence tokens that tell a model: "<|im_start|>system\nYou are a helpful assistant..."), the model does not truly understand how to stay in character or follow strict user instructions. It simply completes text statistically.

Final Review & Path Forward

I ran the training engine for 5,000 continuous iterations on my free Colab T4 GPU instance. Over those 5,000 steps, the cross-entropy loss dropped dramatically from an initial 10.88 down to a stable 3.29.

Building a 28-million parameter model from scratch taught me more about AI infrastructure than any basic online tutorial ever could. Now that the core architecture is live and working on GitHub, my next step is to tackle fine-tuning and implement proper conversation template formatting so I can turn this raw text completer into an interactive assistant.

If you want to view the full repository or fork the code layout to run on your own free Colab instance, check out the implementation here: [https://github.com/kadamsumit335-ctrl/NanoChatLM](https://github.com/kadamsumit335-ctrl/NanoChatLM)

Have you tried building or training custom architectures from scratch? Let’s connect and talk about it in the comments below!