Pretrained Mini LLM
MyLLM v1 — Train your own decoder-only GPT language model from randomly initialized weights using your custom BPE tokenizer, memory-mapped dataset pipeline, and full pretraining engine
Hardware-Matched Educational Model Scales (MyLLM v1)
Choose a model scale matched to your hardware budget: (1) Tiny Educational Model (100K–1M parameters, e.g., 4 layers, d_model=128) for CPU or integrated laptops; (2) Small Single-GPU Model (5M–20M parameters, e.g., 6–8 layers, d_model=384) for consumer GPUs or free cloud notebooks; or (3) Optional Extended Model (30M–100M parameters, e.g., 12 layers, d_model=768) for dedicated GPU workstations. The exact same pipeline scales across all three tiers.
What You Will Build
This is the capstone engineering milestone of Phase 6 (`MyLLM v1`). You will bring together your custom BPE tokenizer (Project 8) and your Mini GPT architecture (Project 10) to train a complete decoder-only language model from randomly initialized weights. You will curate and clean a text dataset, perform document-level deduplication, split into leak-free train/validation/test sets, encode and pack documents into memory-mapped binary shards (`train.bin`, `val.bin`), run a mixed-precision pretraining loop with gradient accumulation, linear warmup + cosine learning rate decay, and checkpoint recovery, evaluate validation perplexity and sample generation across training steps, run a controlled ablation study, and package `MyLLM v1` with a Model Card and standalone CLI inference script.
Why This Project Matters in the AI/LLM Stack
Most developers only call hosted APIs or fine-tune existing checkpoints. Training `MyLLM v1` from random weights—watching cross-entropy loss drop from `ln(V)` (~9.0) down to coherent domain prose (~2.1), debugging learning rate warmup spikes, and verifying zero train-validation contamination—gives you end-to-end ownership of how foundation models are born.
Learning Objectives
- Curate, clean, deduplicate, and split a raw text corpus at the document level without train-validation leakage.
- Tokenize and pack documents with `<|eos|>` separators into contiguous `np.memmap` (`uint16`) binary shards.
- Engineer a fault-tolerant PyTorch pretraining loop with gradient accumulation, mixed precision, gradient clipping, and checkpoint resumption.
- Evaluate pretrained checkpoints quantitatively via validation perplexity, baseline comparisons, and controlled ablation experiments.
- Package `MyLLM v1` with weights, tokenizer artifacts, `config.json`, a transparent Model Card, and a streaming CLI inference script.
Prerequisites
- Phase 4 Project: BPE Tokenizer (`milestone-08-bpe-tokenizer`)
- Phase 5 Project: Mini GPT Architecture (`milestone-10-mini-gpt-architecture`)
- Phase 6 Modules 6.1–6.5: Dataset Engineering, Pretraining Pipeline, and Packaging
Required Mathematical Concepts
- Autoregressive next-token cross-entropy loss and validation perplexity: PPL = exp(L_val)
- Initial uniform random loss baseline: L_0 = ln(V)
- Effective batch size in tokens: B_tokens = B_micro * GradAccumSteps * T
- Linear warmup + cosine learning-rate decay schedule
- Decoupled AdamW weight decay (2D matrix weights vs. 1D biases/norms)
Step-by-Step Implementation Tasks (19 Steps)
Check off each engineering task as your code passes its acceptance criteria. Progress is saved automatically in your browser.
Dataset Curation, Cleaning, Deduplication & Binary Sharding
Select an open text corpus, clean extraction artifacts, deduplicate documents, create document-level splits, and pack tokenized IDs into memory-mapped binary files.
A cohesive domain corpus (such as TinyStories, open educational textbooks, Project Gutenberg literature, or permissive Python files) allows a small model to learn grammar and domain structure rapidly.
Corrupted UTF-8 bytes, HTML boilerplate, and repetitive control characters waste vocabulary capacity and distort loss.
Splitting must occur at the document level after deduplication so identical passages never appear in both `train` and `val`.
Fitting the tokenizer strictly on the training split prevents vocabulary leakage from held-out validation documents.
Concatenating tokenized documents separated by `<|eos|>` and writing them to disk as `uint16` binary arrays allows `np.memmap` to stream any dataset size with near-zero RAM.
Each training step samples `batch_size` random starting offsets `i` in `data` and slices `x = data[i : i + T]` and `y = data[i + 1 : i + T + 1]`.
Model Configuration & Optimizer Setup
Define hardware-matched model configurations, separate 2D vs. 1D parameter groups for AdamW, and implement the warmup-cosine learning rate scheduler.
Matching parameter count (`100K–1M`, `5M–20M`, or `30M–100M`) and context length (`T=256` or `512`) to your hardware ensures fast iteration.
Weight decay should regularize 2D projection/embedding matrices (`weight_decay=0.1`) but never 1D layer-norm scales or biases (`weight_decay=0.0`).
Warming up the learning rate linearly from `0` to `max_lr` over the first 2–5% of steps prevents early Adam variance spikes from diverging random weights.
Before launching a multi-hour pretraining run, verifying that the pipeline can memorize a single fixed batch down to `loss < 0.05` catches bugs in loss shifting, optimizer stepping, or masking.
Pretraining Loop, Checkpointing & Live Telemetry
Implement the main training loop with gradient accumulation, mixed precision, gradient norm clipping, periodic validation evaluation, and resumable checkpoints.
Gradient accumulation simulates a large effective token batch size (`B_micro * grad_accum_steps * T`) on limited VRAM, while `torch.amp.autocast` (`bfloat16`/`float16`) accelerates matmuls.
Clipping global gradient norm at `1.0` protects against unlucky data batches, and logging pre-clip `grad_norm` diagnoses training stability.
A single validation batch is noisy; averaging loss over `eval_iters` (e.g., 50 batches) in `model.eval()` + `torch.no_grad()` provides a reliable validation perplexity curve.
Saving model weights, optimizer states, step counter, RNG states, and config allows training to resume seamlessly after preemption.
Sampling from a fixed set of 5 benchmark prompts at step 0, 25%, 50%, 75%, and 100% visually demonstrates the progression from random token salad to coherent prose.
Evaluation, Ablation Study & MyLLM v1 Packaging
Compare validation perplexity against baselines, run a controlled ablation study, and export the self-contained `MyLLM v1` release bundle.
Comparing your trained model's test perplexity against a random uniform baseline (`V`) and a unigram/bigram frequency baseline proves how much structure the Transformer learned.
Running two controlled training runs where only one hyperparameter changes (e.g., `block_size=128 vs 256` or `RoPE vs learned positional embeddings`) teaches scientific LLM experimentation.
Anyone cloning your repository should be able to generate text from your trained model with a single terminal command.
A complete LLM release includes inference weights (stripped of optimizer state), `config.json`, `tokenizer.json`, evaluation report, and `MODEL_CARD.md`.
Starter & Verification Code Examples
AdamW 2D vs. 1D Parameter Group Partitioning & Warmup-Cosine LR
Reference implementation for separating weight-decayed matrix parameters from non-decayed norms/biases and computing cosine learning rate.
import math
import torch
import torch.nn as nn
def configure_optimizers(model: nn.Module, weight_decay: float, lr: float, betas=(0.9, 0.95)):
param_dict = {pn: p for pn, p in model.named_parameters() if p.requires_grad}
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
optim_groups = [
{"params": decay_params, "weight_decay": weight_decay},
{"params": nodecay_params, "weight_decay": 0.0},
]
return torch.optim.AdamW(optim_groups, lr=lr, betas=betas)
def get_cosine_lr(step: int, warmup_steps: int, max_steps: int, max_lr: float, min_lr: float) -> float:
if step < warmup_steps:
return max_lr * (step + 1) / max(1, warmup_steps)
if step >= max_steps:
return min_lr
decay_ratio = (step - warmup_steps) / max(1, (max_steps - warmup_steps))
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (max_lr - min_lr)How to Test Your Implementation
Verify leak-free dataset sharding, gradient accumulation equivalence, single-batch memorization, and standalone `MyLLM v1` checkpoint reloading.
Zero Train/Validation Document Leakage Check
Ensure validation perplexity reflects true generalization rather than duplicate documents.
Steps: Compute SHA-256 hashes of all raw documents in train and val splits and assert `len(train_hashes & val_hashes) == 0`.
Pass Criteria: Zero shared documents between splits.
Gradient Accumulation Equivalence Test
Verify that accumulating gradients over 4 micro-batches of size 2 equals 1 batch of size 8.
Steps: Fix RNG seed, run 1 step with `(B=8, accum=1)` in FP32 and 1 step with `(B=2, accum=4)`, and compare parameter gradients.
Pass Criteria: `torch.allclose(g1, g2, atol=1e-6)` is `True`.
Standalone Release Reload & Generation Test
Confirm the exported `myllm_v1/` bundle can be loaded in a fresh Python process.
Steps: Load `myllm_v1/` via `generate.py` with `temperature=0.0` (greedy) twice on the same prompt.
Pass Criteria: Both runs produce identical, coherent token sequences.
Overall Project Acceptance Criteria
- Dataset pipeline deduplicates at the document level and verifies zero document overlap between `train.bin`, `val.bin`, and `test.bin`.
- Pretraining pipeline verifies initial step-0 loss `≈ ln(vocab_size)` and passes the single-batch overfit sanity check (`loss < 0.05`).
- Training loop implements AdamW 2D/1D parameter groups, linear warmup + cosine decay, gradient accumulation, and gradient norm clipping.
- `MyLLM v1` is packaged with custom tokenizer, `config.json`, inference weights, evaluation + ablation report, `MODEL_CARD.md`, and `generate.py` CLI.
Expandable Hints & Debugging Tips
Try implementing each stage on your own first. Expand a hint below if you want to verify tensor shapes, mathematical identities, or common edge cases.
Optional Stretch Challenges
Scaling Law Mini-Sweep (3 Model Sizes)
Train 3 model sizes (e.g., 500K, 2M, and 8M parameters) on the same dataset and plot final validation loss against total training FLOPs ($C \approx 6 N D$).
What You Learn: Empirical power-law scaling curves ($L(C)$) and compute-optimal allocation.
MinHash Near-Duplicate Filtering
Upgrade your exact SHA-256 deduplicator to a 5-gram MinHash LSH near-deduplicator that catches documents with >85% Jaccard similarity.
What You Learn: Web-scale fuzzy deduplication algorithms used in FineWeb and Dolma.
What Your Completed Repository Should Contain
`MyLLM v1` — A complete self-trained language model release containing your custom BPE tokenizer, GPT architecture, data & pretraining pipeline, trained checkpoints, evaluation & ablation report, Model Card, and CLI inference script.
Deliverable Artifacts
- Custom tokenizer (`tokenizer.json`) & memory-mapped dataset preparation script (`prepare_data.py`)
- GPT-style architecture (`model.py`) & configuration files (`config.json`)
- Pretraining pipeline (`train.py`) with AdamW, cosine LR, gradient accumulation, and checkpointing
- Trained model checkpoints (`myllm_v1/model.pt`)
- Quantitative evaluation & ablation report (`EVALUATION_REPORT.md`)
- Standalone streaming CLI inference script (`generate.py`)
- Reproducible documentation & transparent Model Card (`MODEL_CARD.md`)
After Completing This Project, You Should Be Able To...
- Engineer a leak-free, deduplicated, memory-mapped pretraining dataset pipeline.
- Train a GPT-style language model from random weights to coherent text generation.
- Diagnose training curves, gradient norms, learning rate warmup, and validation perplexity.
- Design and interpret controlled ablation experiments and package reproducible LLM artifacts.
GitHub Repository Submission
Link your public or private GitHub repository URL to keep a structured portfolio of your from-scratch AI builds. Stored locally in your browser (self-reported; no GitHub tokens or remote code execution).
Suggested Repository: myllm-v1-pretraining-from-scratch
Recommended Directory Layout:
data/prepare_data.py myllm/model.py myllm/train.py myllm/evaluate.py generate.py myllm_v1/config.json myllm_v1/MODEL_CARD.md EVALUATION_REPORT.md tests/test_pretraining_pipeline.py README.md
Recommended README Documentation Checklist
- Include the training and validation loss/perplexity curve plot.
- Show chronological sample generations from step 0, 25%, 50%, and 100%.
- Include the ablation study comparison table.
- Provide a single-command quickstart for `python generate.py`.