From Scratch LLM

Training a Small GPT-Style Language Model from Scratch with PyTorch

---

Overview

Tiny LLM is a from-scratch implementation of a small GPT-style language model, trained using PyTorch. The project explores the process of training a transformer-based language model on diverse datasets, including Wikipedia, stories, Q&A, and instruct datasets. The goal is to understand the mechanics of modern language models, from data preprocessing to model architecture and training dynamics.

The model is trained on an RTX 3060 (12GB VRAM) with a batch size of 16, using mixed precision and gradient accumulation for efficiency.

---

How It Works

Data Pipeline

1. Datasets: The model is trained on a mix of:
   • English Wikipedia
   • Simple Stories
   • FineWeb-Edu
   • OpenWebText2
   • Q&A datasets (e.g., Alpaca, Reddit Instruct)
2. Tokenizer: A custom BPE tokenizer based on GPT-2, extended with special tokens (<|im_start|>, <|im_end|>, <|pad|>).
3. Preprocessing: Data is shuffled, tokenized, and padded dynamically to a context length of 512 tokens.

Model Architecture

• Embedding Dimension: 512
• Attention Heads: 8
• Layers: 8
• Vocabulary Size: 50,260
• Optimizer: AdamW with cosine annealing (lr=3e-4)
• Precision: Mixed (bfloat16)
• Gradient Accumulation: 4 steps

---

Loss evolution ( after 60k steps)

Output evolution ( 40k steps)

First trial training on 25k steps

After the first 25,000 steps, the model shows early signs of learning and loss is decreasing:

Step	Example Output
500	Repetitive, nonsensical output: "The on� to� to� to� ..."
10,000	Begins forming coherent phrases: "The first thing you can do you?"
20,000	Generates contextually relevant (but repetitive) responses: "What is the difference between X and Y?"
25,000	Improved coherence, though still prone to loops: "How do you think about the world?"

Full training logs and generated samples are saved in train_output.txt and second_train_output.txt.

Training on the complete dataset

To be done ...

---

Installation

git clone https://github.com/rantaluca/tiny_llm.git
cd tiny_llm
pip install torch tiktoken datasets tqdm