microme-125m

MicroMe-125M

A 125-million-parameter GPT-style language model trained entirely from scratch on a single RTX 4060 laptop GPU (8 GB) — no pretrained weights, no framework, ~600 lines of hand-written PyTorch. Built as a learning project to understand how a language model is actually made, end to end: tokenizer → transformer → optimizer → pretraining → chat fine-tuning → retrieval.

Honest scope: at 125M parameters this model has fluent language and follows chat format, but it is not a reliable knowledge store — it will state facts confidently and get them wrong. For factual use, pair it with retrieval (RAG). See Limitations.

Highlights

Architecture

Modern-Llama skeleton with speedrun upgrades, all hand-written in model.py: RMSNorm (pre-norm) · RoPE · QK-norm · ReLU² MLP (no gate) · untied input/output embeddings · logit softcap · SDPA/FlashAttention · zero-init residual projections.

Training

1. Tokenizer — 24,576 byte-level BPE on FineWeb-Edu (4.44 chars/token, exact round-trip).
2. Pretraining — 3.0B FineWeb-Edu tokens, Muon+AdamW, warmup + linear cooldown → base val 2.97.
3. Chat SFT — SmolTalk (391M tokens), assistant-only masked loss → follows multi-turn conversations.
4. Persona — a small identity/small-talk set (with a SmolTalk anchor to avoid forgetting), so it answers greetings and knows its name is MicroMe.

The Muon optimizer was reimplemented by hand (~60 LOC) and beat AdamW ~1.3× on sample efficiency in a controlled ablation.

Evaluation

0-shot on the base model via lm-evaluation-harness (self-reported), against standard peers in the 125M class. Peer numbers are approximate published values; MicroMe was run through the same harness.

Competitive with GPT-2-124M and Pythia-160m on 3 of 4 tasks despite training on only ~3B tokens (10–100× less than those models). ARC-Easy is notably strong — a direct payoff of pretraining on educational FineWeb-Edu text. LAMBADA is the weak spot: last-word prediction is penalized by the small 24k vocab and the short training run.

Intended use & limitations

Good at: fluent English, chat format, multi-turn context, short factual answers when grounded with retrieval.

Not good at: reliable world knowledge, arithmetic, multi-step reasoning, or strict formats (e.g. haiku). Three separate fine-tuning experiments confirmed that you cannot fine-tune knowledge into 125M parameters — style lives in the weights, facts must come from retrieval. Use it for learning, experimentation, and RAG demos, not as a knowledge source.

How to use

This is a custom architecture, so load it with the bundled model.py:

import torch, importlib.util
from huggingface_hub import snapshot_download
from tokenizers import Tokenizer
from safetensors.torch import load_file

d = snapshot_download("Ankitgdes/microme-125m")
spec = importlib.util.spec_from_file_location("mm", f"{d}/model.py")
mm = importlib.util.module_from_spec(spec); spec.loader.exec_module(mm)

model = mm.GPT(mm.PRESETS["micro_125m"]).eval()
model.load_state_dict(load_file(f"{d}/model.safetensors"), strict=True)
tok = Tokenizer.from_file(f"{d}/tokenizer.json")
S = {k: tok.token_to_id(f"<|{k}|>") for k in ("bos","user","assistant","end","system")}

# Chat format: <|bos|><|system|>..<|end|><|user|>..<|end|><|assistant|>
# The system prompt carries MicroMe's persona — include it for identity/chitchat.
SYSTEM = "You are MicroMe, a small, friendly AI assistant. Answer briefly and helpfully."

def chat(user, max_new=60):
    ids = ([S["bos"], S["system"]] + tok.encode(SYSTEM).ids + [S["end"]]
           + [S["user"]] + tok.encode(user).ids + [S["end"], S["assistant"]])
    out = model.generate(torch.tensor([ids]), max_new, temperature=0.4, top_k=40, rep_penalty=1.3)[0].tolist()[len(ids):]
    return tok.decode(out[:out.index(S["end"])] if S["end"] in out else out)

print(chat("What is your name?"))          # -> My name is MicroMe.
print(chat("What is the largest planet?")) # weak on facts alone — pair with retrieval (RAG) for accuracy

For accurate factual answers, retrieve relevant context and prepend it to the user turn — a 125M model is not a reliable knowledge store on its own (see Limitations).

Training data

Pretraining on FineWeb-Edu (ODC-By, attribution) and chat/persona tuning on SmolTalk (Apache-2.0). No share-alike data is used, so the model is released under Apache-2.0.

Acknowledgements

Recipe inspired by nanoGPT / nanochat (Karpathy) and the modded-nanogpt speedrun community (Muon, QK-norm, ReLU²). Built as an end-to-end learning project.