{"slug": "i-trained-a-113m-parameter-earthquake-llm-from-absolute-scratch", "title": "I trained a 113M-parameter earthquake LLM from absolute scratch", "summary": "A developer trained a 113M-parameter language model for earthquake science from scratch using open-access papers, Wikipedia, and FineWeb-Edu on two NVIDIA A30 GPUs. The model achieved a 35% reduction in bits-per-byte on general prose by mixing domain and general text, demonstrating the full pretraining lifecycle from data crawling to inference.", "body_md": "**Train a small GPT for earthquake science — the entire LLM lifecycle, from a blank folder to a talking model, explained block by block.**\n\nCrawl → Clean → Tokenize → Model → Train → Infer, on 2× NVIDIA A30 (48 GB).\n\nSix free data sources → crawl → clean/dedup → 16k BPE → 113M GQA+RoPE decoder → 2-GPU DDP training → streaming inference. Each stage is a section below.\n\nnanoGPT-Seis is a teaching repository. It is not trying to be a great earthquake model — it is trying to make every stage of pretraining a language model legible: where the data comes from, how it is cleaned and deduplicated, how a tokenizer is built, why the Transformer looks the way it does, how it is trained across two GPUs, and how it is served. Every design decision is explained, and every number (perplexity, VRAM, tokens) is one we actually measured on this hardware.\n\nThe corpus mixes earthquake / seismology text (open-access papers via\nCrossref+Unpaywall, arXiv/EarthArXiv preprints, the \"Earthquake Insights\" Substack)\nwith general text (Wikipedia + FineWeb-Edu) for plain-language fluency — about\n24% domain / 76% general. A focused corpus lets a ~100M-param model become\ngenuinely fluent on a single node, so you can see the whole loop close in a day, not\na month. (Why the general mix? See [§1](#1-results-at-a-glance).)\n\nStatus:the pretraining lifecycle is complete — crawl through inference.\n\n[Results at a glance](#1-results-at-a-glance)[Quick start](#2-quick-start)[Stage 1 — Data crawling](#3-stage-1--data-crawling)[Stage 2 — Processing & dedup](#4-stage-2--processing--dedup)[Stage 3 — The BPE tokenizer](#5-stage-3--the-bpe-tokenizer)[Stage 4 — The model (RoPE, GQA, …)](#6-stage-4--the-model)[Stage 5 — Training (DDP, VRAM, LR)](#7-stage-5--training)[Stage 6 — Inference](#8-stage-6--inference)[Repository layout](#9-repository-layout)[Scaling-law experiments](#10-scaling-law-experiments)\n\n| value | |\n|---|---|\n| Corpus | 533,248 docs · 485.7M words · 822.7M training tokens (≈2.4:1 general:domain) |\n| Model | 113M params — decoder-only, GQA + RoPE + RMSNorm + SwiGLU |\n| Hardware | 2× NVIDIA A30 (24 GB each), bf16, DDP — also runs on a single RTX 3090/4090, or 12–16 GB with a smaller batch (\n|\n\n**4096** tokens**0.997**— vs 1.527 for a domain-only base (**−35%**)Three findings worth pausing on:\n\n**Longer context helped.** A controlled A/B (data held fixed) — retraining at 4096 vs 1024 dropped perplexity ~11% (9.74 → 10.93 domain-only) for only ~26% more compute per step — papers have long-range structure a 1024-token window can't see across.**The model uses that context.** In a 4096-token window, loss on tokens at positions 2048–4096 is 25% lower than at 0–64 — it conditions on thousands of preceding tokens (see[§8](#8-stage-6--inference)).**A general-text mix restores fluency.** Adding Wikipedia + FineWeb-Edu (~2.4:1 general:domain) cut bits/byte on general prose by 35% vs a paper-only base — see the data-mix comparison below.\n\nA paper-only base is fluent in paper-register but repetitive/incoherent in plain\nprose. Adding ~540M tokens of Wikipedia + FineWeb-Edu (→ ~823M train tokens,\n~2.4:1 general:domain, ~3.8 epochs — within the ~4-epoch repeat budget shown to be\nnear-lossless by [Muennighoff et al., 2023](https://arxiv.org/abs/2305.16264)) gives v2.\n\nMeasured with bits-per-byte (tokenizer-independent, so v1↔v2 is fair;\n`src/compare_models.py`\n\n):\n\nv2 is far more fluent on general prose (−35% bits/byte) and generates coherent\nnon-earthquake text where v1 emits gibberish, at the cost of some domain sharpness\n(+22%) — the classic fluency↔specialization trade-off. This fluent base is the\nright starting point for SFT; base pretraining alone never yields chat. Domain-only\nweights are kept as `checkpoints/ckpt_v1_domain.pt`\n\n.\n\nThe pretrained 113M checkpoint is hosted on the Hugging Face Hub:\n[ jiazhe868/nanogpt_seis](https://huggingface.co/jiazhe868/nanogpt_seis).\n\n```\n# environment (a working CUDA-12.4 PyTorch; see the note below)\nconda activate nanogpt_seis\npip install -r requirements.txt\n\n# download the checkpoint, tokenizer, and model config into the paths expected by\n# src/inference.py\nhuggingface-cli download jiazhe868/nanogpt_seis \\\n    checkpoints/ckpt.pt \\\n    data/tokenized/tokenizer.json \\\n    data/tokenized/meta.json \\\n    configs/gpt120m_ctx4k.yaml \\\n    --local-dir .\n\npython -m src.inference --prompt \"The 2011 Tohoku earthquake\"\n# environment (a working CUDA-12.4 PyTorch; see the note below)\nconda activate nanogpt_seis\npip install -r requirements.txt\n\n# --- run the whole pipeline ---\n# earthquake-domain sources\npython -m src.crawl.wikipedia        --max-pages 500                              # earthquake-titled pages\npython -m src.crawl.fulltext         --per-journal 3000 --broad 30000 --workers 64\npython -m src.crawl.preprints        --arxiv 3000 --eartharxiv 2000\npython -m src.crawl.substack         --max 500\n# general-text mix for fluency (~540M tokens: Wikipedia + FineWeb-Edu)\npython -m src.crawl.general          --wiki-tokens 300000000 --fineweb-tokens 240000000\n\npython -m src.process.build_corpus   --val-frac 0.005      # clean · dedup · split\npython -m src.tokenizer.train_bpe    --vocab-size 16384    # train the tokenizer\npython -m src.tokenizer.encode                             # → uint16 shards\n\ntorchrun --standalone --nproc_per_node=2 \\\n    -m src.train --config configs/gpt120m_ctx4k.yaml       # train on 2 GPUs\n\npython -m src.inference --prompt \"The 2011 Tohoku earthquake\"   # streams live\n```\n\nA PyTorch built for a newer CUDA than your driver will silently report⚠️ Environment gotcha.`cuda.is_available() == False`\n\nand fall back to CPU. Verify with`python -c \"import torch; print(torch.cuda.is_available())\"`\n\n— it must print`True`\n\n. This project uses`torch 2.6.0+cu124`\n\nto match a CUDA-12.5 driver.\n\n**Goal:** assemble a large, legal, full-text earthquake corpus from free sources.\nCode: `src/crawl/`\n\n.\n\n| source | module | what we pull | mechanism |\n|---|---|---|---|\n| Research papers | `fulltext.py` |\nOA full-text PDFs of earthquake papers | Crossref (DOIs) → Unpaywall (OA PDF) → download → extract |\n| Preprints | `preprints.py` |\narXiv + EarthArXiv full text | arXiv API + OSF/DOI → PDF |\n| Wikipedia | `wikipedia.py` |\npages titled \"earthquake\" | MediaWiki API plaintext extracts |\n| Substack | `substack.py` |\n\"Earthquake Insights\" articles | archive API + HTML body parse |\n| General text | `general.py` |\nWikipedia + FineWeb-Edu (fluency mix) | HF `datasets` streaming to a token budget |\n\nWhy a general-text mix?A ~113M model trained only on research papers becomes fluent in paper-register but repetitive and incoherent in plain prose, and 240M tokens is far below compute-optimal ([Hoffmann et al., 2022]). So we add ~240M tokens of[Wikipedia](encyclopedic) +[FineWeb-Edu]([Penedo et al., 2024]; quality-filtered educational web) for a ~1:1 general:domain mix — a fluent base that the planned SFT stage can then make conversational. (Base pretraining alone never yields chat; that's SFT.)\n\nEvery document is normalized to one schema (`src/crawl/common.py`\n\n):\n\n```\n@dataclass\nclass Doc:\n    source: str          # \"fulltext\" | \"arxiv\" | \"wikipedia\" | ...\n    id: str              # stable per-source id (used for dedup)\n    title: str\n    text: str            # the cleaned body we will tokenize\n    url: str = \"\"\n    date: str = \"\"\n    extra: dict = field(default_factory=dict)   # venue, cited_by, full_text, ...\n```\n\n**Finding the papers (Crossref).** Crossref indexes ~150M scholarly works with a free,\ngenerous API. We page through earthquake journal-articles with a deep cursor\n(no offset limit), filtered by journal ISSN:\n\n```\n# src/crawl/fulltext.py — iter_crossref()\nparams = {\"rows\": 1000, \"cursor\": cursor,\n          \"filter\": \"type:journal-article,issn:0094-8276\",   # e.g. GRL\n          \"query.bibliographic\": \"earthquake\", \"mailto\": EMAIL}\n```\n\n**Finding the open PDF (Unpaywall).** A DOI is not a PDF. Unpaywall (free, ~100k/day)\nmaps a DOI to its legal open-access copies. We try repository (green OA) locations\nfirst — publisher links are frequently bot-blocked stubs:\n\n```\n# repository copies download far more reliably than publisher links\nprio = 0 if loc.get(\"host_type\") == \"repository\" else 1\n```\n\n**Downloading in parallel, politely.** PDFs come from many hosts, so we use a thread\npool but throttle per host — different servers download concurrently while any\nsingle server stays rate-limited:\n\n``` python\nclass HostThrottle:                 # src/crawl/fulltext.py\n    def wait(self, host):\n        with self._guard:                                  # get/create this host's lock\n            host_lock = self._locks.setdefault(host, threading.Lock())\n            self._last.setdefault(host, 0.0)\n        with host_lock:                                    # serialize only this host\n            delta = self.min - (time.monotonic() - self._last[host])\n            if delta > 0: time.sleep(delta)\n            self._last[host] = time.monotonic()\n```\n\n**Validating every download.** Not every \"OA PDF\" is real — many are 5 KB anti-bot\nlanding pages. We accept a download only if it is a real PDF with enough text:\n\n```\nif not pdf_bytes.startswith(b\"%PDF\"): return None   # HTML / stub\nif doc.page_count < 2:                return None    # cover page\nif len(text) < min_chars:             return None    # too little extracted\n```\n\n**Not wasting the budget.** Some journals (Science, Nature) are almost entirely\npaywalled — scanning thousands of their DOIs would burn the Unpaywall budget for zero\nfull text. A low-yield abort gate skips a journal once its hit-rate stays under a\nthreshold:\n\n```\nif scanned >= abort_after and got_ft / max(1, scanned) < min_hit:\n    break        # this venue isn't worth more API calls\n```\n\n**Resumable.** Output is appended line-by-line and already-fetched ids are skipped on\nrestart, so a multi-hour crawl survives interruption:\n\n```\ndone = _load_done_ids(out)          # ids already in the JSONL\n...\nif iid in done: continue            # skip work we already have\n```\n\nWar story (why Crossref + Unpaywall).This project originally enumerated papers via OpenAlex, which changed to a paid credit model mid-build — the free daily budget ran out after ~100 requests. Crossref + Unpaywall is the free, robust replacement, and it is what the code ships with. The lesson — pin your data source assumptions and make the crawler resumable — is baked into the design.\n\nFull-text yield varies a lot by source: arXiv ~99% (open by design), the broad OA pool ~15%, paywalled journals ~0% (abstract fallback). Net corpus: ~20k full-text papers + ~26k abstracts.\n\nThe crawler above is API-driven, but the shape underneath is the classic one: a breadth-first traversal of a link graph, run across many threads. It is worth writing that general pattern out in full, because getting the concurrency right is the hard part. The crux is shared state that every worker touches at once — a frontier of URLs still to visit, and a set of URLs already seen.\n\nThree rules make it correct:\n\n- The frontier is a\n`queue.Queue`\n\n, which is internally synchronised:`get()`\n\nand`put()`\n\nare atomic, and`get()`\n\nblocks until work exists, so workers never busy-wait. - The seen-set gets its own lock. \"Is this URL new? if so, add it\" is a read-modify-write; without the lock two threads can both judge the same URL new and enqueue it twice.\n- Termination uses\n`Queue.join()`\n\n/`task_done()`\n\n. The hard question in a concurrent BFS is knowing when you are done: a worker finding the frontier momentarily empty does not mean the crawl is finished, because another worker may be one instruction away from enqueuing more links.`join()`\n\ncounts outstanding tasks and only releases once every enqueued URL has been fully processed.\n\n``` python\nimport re, threading, queue, requests\nfrom urllib.parse import urljoin, urldefrag, urlparse\n\nclass BFSCrawler:\n    \"\"\"Breadth-first web crawler: one shared frontier, N worker threads.\"\"\"\n\n    def __init__(self, seeds, max_pages=5000, n_workers=16, min_interval=1.0):\n        self.frontier   = queue.Queue()          # thread-safe BFS queue of (url, depth)\n        self.seen       = set(seeds)             # every URL ever enqueued\n        self.seen_lock  = threading.Lock()       # guards `seen` (check-then-add races)\n        self.pages      = []                     # kept (url, html) pairs\n        self.pages_lock = threading.Lock()       # guards `pages` + the page counter\n        self.max_pages  = max_pages\n        self.n_workers  = n_workers\n        self.throttle   = HostThrottle(min_interval)   # §3.2 — politeness, per host\n        for s in seeds:\n            self.frontier.put((s, 0))            # seed the frontier at depth 0\n\n    def _links(self, base, html):                # extract absolute, de-fragmented links\n        for m in re.finditer(r'href=[\"\\'](.*?)[\"\\']', html):\n            url, _ = urldefrag(urljoin(base, m.group(1)))   # relative → absolute, drop #frag\n            if url.startswith((\"http://\", \"https://\")):\n                yield url\n\n    def _enqueue(self, url, depth):\n        with self.seen_lock:              # the check AND the add are ONE critical section,\n            if url in self.seen:          # else two threads both see `url` as new and\n                return                    # enqueue it twice\n            self.seen.add(url)\n        self.frontier.put((url, depth))\n\n    def _worker(self):\n        while True:\n            url, depth = self.frontier.get()      # blocks; (None, _) is the stop sentinel\n            if url is None:\n                self.frontier.task_done()\n                return\n            try:\n                with self.pages_lock:\n                    if len(self.pages) >= self.max_pages:\n                        continue                  # budget spent — drain the rest quietly\n                self.throttle.wait(urlparse(url).netloc)      # rate-limit this host\n                html = requests.get(url, timeout=10).text\n                with self.pages_lock:\n                    if len(self.pages) >= self.max_pages:\n                        continue\n                    self.pages.append((url, html))\n                for link in self._links(url, html):           # BFS: expand the frontier\n                    self._enqueue(link, depth + 1)\n            except Exception:\n                pass                              # one dead link must not kill a worker\n            finally:\n                self.frontier.task_done()         # pairs with the get() above\n\n    def run(self):\n        workers = [threading.Thread(target=self._worker, daemon=True)\n                   for _ in range(self.n_workers)]\n        for t in workers:\n            t.start()\n        self.frontier.join()                      # wait until the frontier is truly drained\n        for _ in workers:                         # then release the idle, blocked workers\n            self.frontier.put((None, 0))\n        for t in workers:\n            t.join()\n        return self.pages\n```\n\n`ThreadPoolExecutor`\n\nreaches the same result with less bookkeeping — submit `_worker`\n\n`n_workers`\n\ntimes and let the pool own the threads — but spelling out the\n`Queue.join()`\n\nhandshake is what makes the termination logic visible.\n\nThis repo's `fulltext.py`\n\nis a one-hop specialisation of exactly this skeleton: the\nfrontier is the list of DOIs Crossref returns, each is \"expanded\" not into page\nout-links but into its Unpaywall PDF locations, and the shared guarded state is the\nresume-set of already-fetched ids rather than a `seen`\n\nset of URLs. The thread pool and\nthe per-host throttle are identical. A full multi-hop BFS would instead follow\n`<a href>`\n\nlinks, which is the version above.\n\nCode: `src/process/`\n\n. Turns messy `data/raw/*.jsonl`\n\ninto clean train/val splits.\n\n``` python\ndef normalize(text):\n    text = html.unescape(text)                 # &#13; &amp; → real chars\n    text = unicodedata.normalize(\"NFKC\", text) # fi ligatures, full-width, …\n    text = _CTRL_RE.sub(\"\", text)              # strip control bytes\n    text = _MULTISPACE_RE.sub(\" \", text)       # collapse whitespace\n    ...\n```\n\nPDFs get extra care: line-break de-hyphenation (`earth-\\nquake → earthquake`\n\n) and\nreference-section stripping (truncate at the last \"References\"/\"Bibliography\"\nheading — the bibliography is noise for a language model).\n\nThree cheap, effective filters (`passes_filters`\n\n):\n\n**length**— drop < 200 chars;** English ratio**— fraction of tokens that are common English stopwords; drops the Spanish-translated Substack posts without a language-ID dependency;**alpha ratio**— fraction of letters/spaces; drops garbled PDF tables and math dumps.\n\nDuplicate training data measurably hurts language models\n([Lee et al., 2021](https://arxiv.org/abs/2107.06499)) — repeated documents get\nover-memorised and waste the token budget. But comparing every document against every\nother is O(N²) and reads every byte, so good dedup is an exercise in doing the cheap\ntest first: eliminate non-candidates with progressively more expensive checks, and pay\nthe expensive one only on what survives.\n\n**The cascade, in its classic form.** The same idea drives the textbook \"find duplicate\nfiles on disk\" problem. Two files can only be identical if they have the same size; of\nthe same-size files, only those whose first kilobyte also matches can be identical; and\nonly those earn a full-content hash. Each tier is strictly more expensive than the last\nand runs on strictly fewer inputs:\n\n``` python\nimport os, hashlib\nfrom collections import defaultdict\n\ndef find_duplicate_files(root: str) -> list[list[str]]:\n    def file_hash(path: str, first_chunk_only=False, chunk=8192) -> str:\n        h = hashlib.md5()\n        try:\n            with open(path, \"rb\") as f:\n                if first_chunk_only:\n                    h.update(f.read(1024))                    # cheap: one small read\n                else:\n                    for block in iter(lambda: f.read(chunk), b\"\"):\n                        h.update(block)                       # full read of the file\n        except OSError:\n            return \"\"\n        return h.hexdigest()\n\n    # Tier 1 — group by size. Just a stat() per file; no bytes are read.\n    by_size: dict[int, list[str]] = defaultdict(list)\n    for dirpath, _, names in os.walk(root):\n        for name in names:\n            path = os.path.join(dirpath, name)\n            if os.path.islink(path):\n                continue\n            try:\n                size = os.path.getsize(path)\n            except OSError:\n                continue\n            if size > 0:\n                by_size[size].append(path)\n\n    # Tier 2 — within each size group, group by a partial (first-1KB) hash.\n    by_partial: dict[str, list[str]] = defaultdict(list)\n    for paths in by_size.values():\n        if len(paths) > 1:                                    # a lone file can't be a dup\n            for p in paths:\n                if (ph := file_hash(p, first_chunk_only=True)):\n                    by_partial[ph].append(p)\n\n    # Tier 3 — only now pay for a full-content hash, on the few real candidates.\n    by_full: dict[str, list[str]] = defaultdict(list)\n    for paths in by_partial.values():\n        if len(paths) > 1:\n            for p in paths:\n                if (fh := file_hash(p, first_chunk_only=False)):\n                    by_full[fh].append(p)\n\n    return [paths for paths in by_full.values() if len(paths) > 1]\n```\n\nWhy it is fast: N files cost N cheap `stat()`\n\ncalls; a partial read happens only for\nfiles that share a size; a full read happens only for files that also share a first\nkilobyte. Almost everything is filtered out before Tier 3, so the expensive\nfull-content pass touches a tiny fraction of the data — the O(N²)-looking problem\ncollapses to roughly O(N). The design rule generalises: cheapest discriminator first,\nnarrow to candidates, confirm last.\n\n**Exact pass, applied to documents.** Text documents are small enough that a single\nhash is the whole cascade — the trick is what you hash. We normalise first (lowercase,\ncollapse whitespace) so copies that differ only in formatting collide, then key on one\nSHA-1. Normalisation is the moral equivalent of the size/partial tiers: it forces \"same\ncontent, different wrapping\" into the same bucket before the hash even runs.\n\n``` php\ndef exact_key(text: str) -> str:                 # src/process/dedup.py\n    normalized = \" \".join(text.lower().split())  # canonicalise whitespace + case\n    return hashlib.sha1(normalized.encode(\"utf-8\")).hexdigest()\n\nseen, unique = set(), []\nfor d in docs:\n    k = exact_key(d[\"text\"])\n    if k in seen:            # O(1) set membership — the whole exact pass is O(N)\n        continue\n    seen.add(k); unique.append(d)\n```\n\n**Near-duplicate pass (MinHash + LSH).** Exact hashing misses documents that are almost\nidentical — e.g. an abstract that also appears verbatim inside its full-text PDF, plus\none extra sentence: one byte differs, so the SHA-1s diverge completely. For those we\nneed a similarity measure, and LSH is the same cascade idea applied to fuzzy matching.\nEach document becomes a set of word 5-grams (shingles); a MinHash\n([Broder, 1997](#references)) signature (128 permutations) estimates the Jaccard\nsimilarity of two shingle sets in constant space; and an LSH index buckets signatures\nso only likely-similar documents are ever compared — the cheap candidate filter that\nreplaces the O(N²) all-pairs comparison, exactly like Tier 1 above. Documents are\nprocessed longest-first, so when a near-dup is found we keep the fullest version:\n\n```\nm = MinHash(num_perm=128)\nfor shingle in word_5grams(text): m.update(shingle.encode())\nif lsh.query(m):          # LSH returns only bucket-mates — the candidate set\n    near_dups += 1        # a ≥70%-similar doc is already kept → drop this one\nelse:\n    lsh.insert(id, m); kept.append(doc)\n```\n\nThe MinHash step is the expensive tier (it reads every surviving document), so\n`dedup.py`\n\ncomputes signatures across a process pool — worker processes inherit the\ndocument list by copy-on-write and return just the signatures — while the stateful LSH\ninsert/query stays serial but cheap (dict ops on precomputed signatures). On the real\ncorpus this removed 1,126 exact + 304 near duplicates. Finally the docs are shuffled\nwith a fixed seed and split into train / val, preserving each doc's metadata.\n\nCode: `src/tokenizer/`\n\n. We train our own tokenizer rather than reuse [GPT-2](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf)'s —\nowning this stage is the point, and a domain vocabulary encodes seismology text more\nefficiently.\n\nBPE (Byte-Pair Encoding — [Sennrich et al., 2015](https://arxiv.org/abs/1508.07909),\nadapting a 1994 compression scheme) starts from single characters and greedily merges the most\nfrequent adjacent pair, over and over, until it reaches the target vocab size. Frequent\nsequences (`earthquake`\n\n, `sub`\n\n, `duction`\n\n) become single tokens; rare ones stay in\npieces. Byte-level ([Wang et al., 2019](https://arxiv.org/abs/1909.03341)) means the\nbase alphabet is the 256 raw bytes, so any input\nis representable — there are no `<unk>`\n\ntokens, ever.\n\n```\n# src/tokenizer/train_bpe.py\ntokenizer = Tokenizer(BPE(unk_token=None))\ntokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)\ntrainer = BpeTrainer(vocab_size=16384, special_tokens=[\"<|endoftext|>\"],\n                     initial_alphabet=ByteLevel.alphabet())   # all 256 bytes\ntokenizer.train_from_iterator(_iter_texts(train_jsonl), trainer)\n```\n\nWe use 16,384 tokens (vs GPT-2's 50,257): smaller vocab → smaller embedding matrix\n(a big fraction of a small model's params) and the domain is narrow. Measured\ncompression on our corpus is ~3.9 chars/token — healthy. The single special token\n`<|endoftext|>`\n\nmarks document boundaries.\n\n`tokenizers`\n\ndoes this for us in fast Rust, but the algorithm is small enough to\nwrite by hand — and understanding it is the point. A BPE tokenizer is just a base\nalphabet plus an ordered list of merge rules, learned in two phases.\n\nThe whole thing fits in one small class — training learns the merges, encoding replays them, decoding concatenates the bytes back:\n\n``` python\nfrom collections import defaultdict\n\nclass BPETokenizer:\n    def __init__(self, vocab_size: int):\n        assert vocab_size >= 256\n        self.vocab_size = vocab_size\n        self.vocab = {}      # token id -> merged bytes\n        self.merges = {}     # (id1, id2) -> new id\n\n    def _get_stats(self, ids: list) -> dict:            # count adjacent id pairs\n        count = defaultdict(int)\n        for i in range(len(ids) - 1):\n            count[(ids[i], ids[i + 1])] += 1\n        return count\n\n    def _merge(self, ids: list, pair: tuple, newidx: int) -> list:\n        output, i = [], 0\n        while i < len(ids):\n            if i < len(ids) - 1 and pair == (ids[i], ids[i + 1]):\n                output.append(newidx); i += 2\n            else:\n                output.append(ids[i]); i += 1\n        return output\n\n    def train(self, text: str):\n        for i in range(256):                            # base vocab = the 256 bytes\n            self.vocab[i] = bytes([i])\n        ids = list(text.encode(\"utf-8\"))\n        for i in range(self.vocab_size - 256):          # each merge adds one token\n            stats = self._get_stats(ids)\n            if not stats:\n                break\n            top_pair = max(stats, key=stats.get)        # most frequent adjacent pair...\n            newidx = 256 + i\n            ids = self._merge(ids, top_pair, newidx)    # ...becomes one new token\n            self.merges[top_pair] = newidx\n            self.vocab[newidx] = self.vocab[top_pair[0]] + self.vocab[top_pair[1]]\n\n    def encode(self, text: str) -> list:                # text -> BPE token ids\n        ids = list(text.encode(\"utf-8\"))\n        while len(ids) >= 2:\n            stats = self._get_stats(ids)\n            # apply the pair whose merge was learned earliest (smallest new id)\n            pair = min(stats, key=lambda p: self.merges.get(p, float(\"inf\")))\n            if pair not in self.merges:\n                break\n            ids = self._merge(ids, pair, self.merges[pair])\n        return ids\n\n    def decode(self, ids: list) -> str:                 # token ids -> text\n        raw = b\"\".join(self.vocab[i] for i in ids)\n        return raw.decode(\"utf-8\", errors=\"replace\")\n```\n\n`train`\n\nstarts from the 256 byte values, then repeatedly counts adjacent id pairs,\nmerges the most frequent pair into a fresh id (`256 + i`\n\n), and records both the merge\nrule and its byte expansion — one new token per iteration, so it stops after\n`vocab_size − 256`\n\nmerges. `encode`\n\nreplays those merges greedily, always applying\nthe one learned earliest (`min`\n\nby new id = highest priority) until no learned pair\nis left. `decode`\n\nconcatenates each id's bytes and UTF-8-decodes. Because the base\nalphabet is all 256 bytes, nothing is ever out-of-vocabulary. (For a longer\nfrom-scratch treatment of this exact algorithm, see Karpathy's\n[minbpe](https://github.com/karpathy/minbpe) and the\n[\"Let's build the GPT Tokenizer\"](https://www.youtube.com/@AndrejKarpathy) video.)\n\nOn our real 16k tokenizer (verified) the same mechanism gives `earthquake`\n\nand\n`subduction`\n\none token each (they're common in the corpus), `newest → [new, est]`\n\n,\n`seismogram → [seism, ogram]`\n\n, and a genuinely rare word splits far more —\n`antidisestablishmentarianism → [ant, id, is, establish, ment, arian, ism]`\n\n.\n\nThis minimal version merges over the raw byte stream; production tokenizers — including\nour `tokenizers`\n\n-based `train_bpe.py`\n\n— add a byte-level pre-tokenizer that first\nsplits text into word-ish chunks so merges never cross word boundaries. But the core\nloop is exactly the above, just in fast Rust so it handles the 3 GB corpus.\n\n`encode.py`\n\nstreams the corpus through the tokenizer and writes one flat array of token\nids per split, documents separated by the eot id:\n\n```\nfor enc in tok.encode_batch(batch):\n    buf.extend(enc.ids); buf.append(eot_id)       # eot between docs\n    if len(buf) >= FLUSH_EVERY:\n        np.asarray(buf, dtype=np.uint16).tofile(fout)   # 2 bytes/token\n```\n\n`uint16`\n\n(0–65535) fits our 16k vocab in 2 bytes/token → 822.7M tokens ≈ 1.6 GB, and the\ntraining loader `np.memmap`\n\ns it (no RAM blow-up). Crucially, this file is\ncontext-length-agnostic — the window width is chosen at training time, which is why\nswitching 1024→4096 needed no re-tokenization.\n\nCode: `src/model/gqa_gpt.py`\n\n. A modern, [Llama](https://arxiv.org/abs/2302.13971)-style [decoder-only Transformer](https://arxiv.org/abs/1706.03762).\n\nBottom-to-top: token ids → embedding → ×16 pre-norm blocks (RMSNorm → GQA+RoPE → residual, then RMSNorm → SwiGLU → residual) → final RMSNorm → weight-tied LM head → logits.\n\n| choice | over the \"classic GPT\" alternative | why |\n|---|---|---|\nRMSNorm (\n|\nLayerNorm | fewer ops, no mean-subtraction, same quality — the Llama default |\nRoPE (\n|\nlearned position embeddings | relative, extrapolates, no position table to learn |\nGQA (\n|\nfull multi-head attention | 3× smaller KV cache for long-context inference, ~no quality loss |\nSwiGLU (\n|\nGELU MLP | gated activation, better quality per parameter |\nweight tying (\n|\nseparate input/output embeddings | saves 12.6M params on a small model, regularizes |\nno biases |\nlinear layers with bias | biases add little; simpler and slightly faster |\n\nAttention is the operation that lets tokens exchange information. Each token's vector is linearly projected into three roles:\n\n**query**`q`\n\n— what this token is looking for;**key**`k`\n\n— what this token offers to others;**value**`v`\n\n— the information it passes on when attended to.\n\nFor one head of width `d_k`\n\n, stack these over the sequence into `Q, K, V`\n\nof shape\n`T × d_k`\n\n(one row per token). Attention scores every query against every key, turns the\nscores into weights, and returns each query's weighted average of the values:\n\n```\nAttention(Q, K, V) = softmax( Q Kᵀ / √d_k + M ) · V\n```\n\n`Q Kᵀ`\n\nis the`T × T`\n\nmatrix of query·key dot products — how strongly token i's query matches token j's key.`softmax`\n\nover each row makes weights that sum to 1, and multiplying by`V`\n\naverages the values under those weights.`M`\n\nis the causal mask:`0`\n\non and below the diagonal,`−∞`\n\nabove, so a token attends only to itself and earlier tokens — it can't peek at the future it is trained to predict.\n\n**Why divide by √d_k.** `q·k`\n\nis a sum of `d_k`\n\nproducts; for roughly unit-variance,\nindependent entries that sum has variance `d_k`\n\n, so raw scores grow like `√d_k`\n\n. Feed\nlarge scores into softmax and it saturates — one weight ≈ 1, the rest ≈ 0 — and its\ngradient vanishes, so learning stalls. Dividing by `√d_k`\n\nrescales scores back to ~unit\nvariance and keeps softmax in a responsive range. (RoPE in §6.3 is applied to `q`\n\nand `k`\n\njust before this product; GQA in §6.4 shares `K`\n\n/`V`\n\nacross heads.)\n\n**Why several heads, not one.** Instead of one attention over the full `d_model`\n\nwidth, we\nsplit into `h`\n\nheads of `d_k = d_model / h`\n\n(here 12 × 64), attend independently, and\nconcatenate. A single head must collapse every kind of relationship — local adjacency,\nsyntax, long-range coreference — into one weight distribution; separate heads let each\nspecialise in a different pattern and a different subspace at the same time. Total compute\nis unchanged (`h · d_k = d_model`\n\n), and the specialisation is visible in the trained model\n(§8.4: some heads are previous-token, some broad, some attention sinks).\n\n**Cost.** Both `Q Kᵀ`\n\nand the weighting of `V`\n\nare `T × T × d_k`\n\nper head, so attention is\n**O(T² · d_model)** in time and **O(T²)** in memory for the score matrix — quadratic in\nsequence length `T`\n\n. That quadratic is why long context is expensive, and why the next two\nsubsections matter: GQA shrinks the per-token state that has to be cached, and\nFlashAttention removes the O(T²) memory term entirely.\n\nA Transformer is permutation-invariant; it needs to be told token order. RoPE encodes\nposition by rotating each query/key vector by an angle proportional to its position.\nThe dot product of a rotated query at position `m`\n\nand key at position `n`\n\nthen depends\nonly on their relative distance `m − n`\n\n— which is exactly what attention should\ncare about, and it lets the model handle positions it can partly extrapolate to.\n\nThe cleanest way to see it: treat each pair of channels as a complex number —\nrotating a 2-D vector by angle θ is just multiplying by the unit complex number\n`e^{iθ}`\n\n. Precompute `e^{i·m·θ_j}`\n\nfor every position `m`\n\nand frequency `θ_j`\n\n, then\nencoding position is a single complex multiply:\n\n``` php\ndef precompute_mtheta(dim: int, seqlen: int, base: float = 1e4) -> torch.Tensor:\n    theta = 1 / (base ** (torch.arange(0, dim, 2) / dim))    # D/2 frequencies\n    m = torch.arange(0, seqlen)                              # S positions\n    mtheta = torch.outer(m, theta)                           # S × D/2 angles\n    return torch.polar(torch.ones_like(mtheta), mtheta)      # unit complex e^{i·mθ}\n\ndef apply_rope(x: torch.Tensor, mtheta: torch.Tensor) -> torch.Tensor:\n    # x: (B, S, H, D);  mtheta: (S, D/2)\n    x_complex = torch.view_as_complex(x.float().unflatten(-1, (-1, 2)))   # B S H D/2\n    mtheta = mtheta[:x.shape[1]].reshape(1, x.shape[1], 1, -1)            # 1 S 1 D/2\n    x_rotated = torch.view_as_real(x_complex * mtheta).flatten(3)         # B S H D\n    return x_rotated.type_as(x)\n```\n\nEach channel pair `(x₂ⱼ, x₂ⱼ₊₁)`\n\nbecomes the complex number `x₂ⱼ + i·x₂ⱼ₊₁`\n\nand is\nrotated by `m·θ_j`\n\n; low-index pairs use large `θ_j`\n\n(rotate fast), high-index pairs\nsmall `θ_j`\n\n(slow). There are no learned parameters, and the rotation preserves norm.\nThe dot product of a rotated query at `m`\n\nand key at `n`\n\nthen depends only on the\nrelative offset `m − n`\n\n— the whole point.\n\nOur repo implements the algebraically-equivalent real-valued `rotate_half`\n\nform\n(same rotation, no complex tensors — friendlier to the compiled / FlashAttention path);\n`cos`\n\n/`sin`\n\nare precomputed once and applied to Q and K just before attention:\n\n``` python\ndef _apply_rope(x, cos, sin):          # x: (B, n_head, T, head_dim)\n    return x * cos + _rotate_half(x) * sin\n```\n\n(One convention pairs adjacent channels, the other pairs channel `j`\n\nwith `j + D/2`\n\n;\nboth give the same relative-position behavior — just a different weight layout.)\n\nTwo consequences (both computed from this model's `head_dim=64, θ=10000`\n\n): the q·k\nscore depends only on the relative distance — curves for queries at positions\n200/800/2000 lie exactly on top of each other (left) — and the channel pairs rotate\nat a geometric spread of frequencies, fast to slow (right), so the model can\nresolve both nearby and far-apart positions.\n\nStandard multi-head attention gives every one of the 12 query heads its own key/value\nhead — so at inference the KV cache (the stored keys/values for every past token)\nholds 12 heads' worth per token. GQA ([Ainslie et al., 2023](https://arxiv.org/abs/2305.13245))\nlets several query heads share a KV head — interpolating between full MHA and\nmulti-query attention ([Shazeer, 2019](https://arxiv.org/abs/1911.02150)). We\nuse 12 query : 4 KV heads, so the KV cache is 3× smaller — cheaper long-context\ninference — at nearly no quality cost, because attention quality is dominated by the\nnumber of query heads (still 12).\n\n```\n# src/model/gqa_gpt.py — GroupedQueryAttention\n# Per-head projection written as an einsum so the contraction is explicit:\n#   out[b, h, t, i] = Σ_d  x[b, t, d] · W[h, i, d]\nq = torch.einsum(\"btd,hid->bhti\", x, self.wq.weight.view(n_head, hd, d))  # 12 heads\nk = torch.einsum(\"btd,hid->bhti\", x, self.wk.weight.view(n_kv,   hd, d))  #  4 heads  ← fewer\nv = torch.einsum(\"btd,hid->bhti\", x, self.wv.weight.view(n_kv,   hd, d))  #  4 heads\nif n_kv != n_head:                       # broadcast KV heads to match Q for SDPA\n    k = k.repeat_interleave(rep, dim=1)  # (the cache itself still stores only n_kv heads)\n    v = v.repeat_interleave(rep, dim=1)\ny = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # FlashAttention kernel\n```\n\nThat one `scaled_dot_product_attention`\n\ncall is just the standard attention math —\nworth spelling out once. Here is the explicit, non-Flash equivalent (also how\n`src/figures/attention_map.py`\n\nrecovers the weights, since the fused kernel doesn't\nexpose them):\n\n``` python\nimport math\n# q, k, v: (B, n_head, T, head_dim) — KV already broadcast to n_head above\nscores = torch.einsum(\"bhqd,bhkd->bhqk\", q, k) / math.sqrt(head_dim)  # (B,n_head,T,T)\nmask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)\nscores = scores.masked_fill(mask, float(\"-inf\"))                      # causal\nattn = torch.softmax(scores, dim=-1)                                  # each row sums to 1\ny = torch.einsum(\"bhqk,bhkd->bhqd\", attn, v)                          # (B,n_head,T,head_dim)\n```\n\nThe two einsums are the query·key similarity (`bhqd,bhkd->bhqk`\n\n) and the value-weighted\nsum (`bhqk,bhkd->bhqd`\n\n) — identical to `q @ kᵀ`\n\nand `attn @ v`\n\n, just with the contracted\naxis named.\n\n**Why FlashAttention** ([Dao et al., 2022](https://arxiv.org/abs/2205.14135); [FlashAttention-2, 2023](https://arxiv.org/abs/2307.08691)). `F.scaled_dot_product_attention`\n\ncomputes exactly the above\n(verified identical to ~1e-7), but attention written this way is memory-bandwidth\nbound, not compute bound: the naive version writes the `T×T`\n\n`scores`\n\nand `attn`\n\nmatrices out to the GPU's large-but-slow main memory (HBM) and reads them back, and\nshuttling those O(T²) intermediates dominates the wall-clock. FlashAttention never\nbuilds them. It splits Q, K, V into blocks small enough to fit in on-chip SRAM — the\nfast per-SM scratchpad (tens of KB, ~10–20× the bandwidth of HBM) — and for each query\nblock streams over the key/value blocks, accumulating the result tile by tile so the\nO(T²) intermediates never leave SRAM. Only the O(T) output is written back to HBM.\n\n**The online softmax, in code.** The one hard part is the softmax. An ordinary softmax\nneeds the whole row at once — subtract the row max for numerical stability, exponentiate,\ndivide by the sum — but a query block only ever sees one key block at a time. The fix is\nan online softmax that keeps three running quantities per query row and updates them\nblock by block: the running max `m`\n\n, the running normaliser `l = Σ exp(·)`\n\n, and the\nrunning un-normalised output `o`\n\n. When a new block raises the max, the earlier\naccumulators are rescaled by `exp(m_old − m_new)`\n\nto rebase them onto the new max:\n\n``` python\nimport torch, math\n\nclass FlashAttention:\n    \"\"\"Tiled attention with an online softmax — the FlashAttention forward pass in\n    plain PyTorch. Never materialises the full S×S score matrix, only Br×Bc tiles.\"\"\"\n\n    def __init__(self, block_q: int, block_kv: int):\n        self.Br = block_q        # query-block size (rows of the score matrix)\n        self.Bc = block_kv       # key/value-block size (columns)\n\n    def __call__(self, q, k, v):\n        # q, k, v: (B, H, S, D)  ->  out: (B, H, S, D)\n        B, H, S, D = q.shape\n        scale = 1.0 / math.sqrt(D)\n        out = torch.zeros_like(q)\n\n        for i in range(0, S, self.Br):                    # outer loop: query blocks\n            q_i = q[..., i:i + self.Br, :]                # (B,H,Br,D) — stays in SRAM\n            o_i = torch.zeros_like(q_i)                   # running un-normalised output\n            l_i = torch.zeros(*q_i.shape[:-1], 1, device=q.device, dtype=q.dtype)  # running Σexp\n            m_i = torch.full_like(l_i, float(\"-inf\"))     # running row-max\n\n            for j in range(0, S, self.Bc):                # inner loop: stream K/V blocks\n                k_j = k[..., j:j + self.Bc, :]\n                v_j = v[..., j:j + self.Bc, :]\n\n                s_ij  = torch.matmul(q_i, k_j.transpose(-2, -1)) * scale   # (B,H,Br,Bc) scores\n                m_new = torch.maximum(m_i, s_ij.max(dim=-1, keepdim=True).values)   # new row-max\n                p_ij  = torch.exp(s_ij - m_new)           # tile probabilities vs the new max\n                corr  = torch.exp(m_i - m_new)            # rescale earlier accumulators (≤ 1)\n\n                l_i = l_i * corr + p_ij.sum(dim=-1, keepdim=True)   # rebase Σexp, add this tile\n                o_i = o_i * corr + torch.matmul(p_ij, v_j)          # rebase output, add p·V\n                m_i = m_new                               # advance the running max\n\n            out[..., i:i + self.Br, :] = o_i / l_i        # normalise once, write O(T) to HBM\n\n        return out\n```\n\nOne inner step reads a `Br×Bc`\n\nscore tile `s_ij`\n\n, folds its row-max into the running\nmax (`m_new`\n\n), forms that tile's probabilities `p_ij = exp(s_ij − m_new)`\n\nagainst the\nnew max, and — before adding them in — multiplies the earlier `l`\n\nand `o`\n\nby the\ncorrection `corr = exp(m_i − m_new) ≤ 1`\n\n. That correction is the whole trick: it\nretroactively puts the running sum and output on the latest max's scale, so after the\nfinal block `o / l`\n\nis bit-for-bit a full softmax (max abs diff ~1e-15 vs\n`scaled_dot_product_attention`\n\n) while nothing larger than one `Br×Bc`\n\ntile is ever held.\nThe production CUDA kernel is this exact recurrence with the tiles living in\nregisters/SRAM. For a causal decoder like ours the inner loop simply skips future\nkey blocks (`j > i`\n\n) and triangular-masks only the diagonal block, so causality roughly\nhalves the work as well.\n\nFar fewer HBM round-trips → both O(T) memory and a big speedup — which is exactly what makes 4096-token context affordable on a 24 GB card.\n\nEvery sublayer is wrapped in normalisation. We use **RMSNorm**: scale each token vector to\nunit root-mean-square, then apply a learned per-channel gain `g`\n\n— no mean subtraction, no\nbias:\n\n```\nRMSNorm(x) = x / sqrt( mean(x²) + ε ) · g\n```\n\nLayerNorm additionally centres the vector (`(x − μ)/σ · γ + β`\n\n); RMSNorm drops the mean\nand the bias, which is cheaper and works as well in practice — the Llama default.\n\nWhere the norm sits matters as much as which one. The original Transformer is\n**post-norm** — normalise after the residual add; modern decoders (GPT-2, Llama, us) are\n**pre-norm** — normalise inside the residual branch, before the sublayer:\n\n```\npost-norm:  x  ←  Norm( x + Sublayer(x) )\npre-norm:   x  ←  x + Sublayer( Norm(x) )\n```\n\n**Why pre-norm.** In pre-norm the residual path `x + …`\n\nis an unbroken identity highway:\ngradients flow from the loss straight back to every layer without passing through a\nnormaliser, so deep stacks train stably from scratch. Post-norm puts a norm on that main\npath, so the gradient is rescaled at every layer on the way down — deep post-norm\nTransformers need careful learning-rate warmup and initialisation just to converge. The\nprice of pre-norm is that the residual stream's variance grows with depth (each block adds\nto it without renormalising the sum); we offset that in the weight init (§6.7).\n\nBetween attention layers, each token is transformed on its own by a feed-forward network.\nThe classic version is `W₂ · GELU(W₁ x)`\n\n: project up, apply a pointwise activation, project\ndown. We use **SwiGLU**, a gated variant with three matrices:\n\n```\nSwiGLU(x) = W₂ ( SiLU(W₁ x) ⊙ (W₃ x) ),     SiLU(z) = z · sigmoid(z)\n```\n\nThe up-projection `W₃ x`\n\nis multiplied element-wise (`⊙`\n\n) by a data-dependent gate\n`SiLU(W₁ x)`\n\n: the network learns, per token and per hidden unit, how much signal to let\nthrough. That multiplicative interaction is more expressive than a single fixed activation\nand gives better quality per parameter. Because SwiGLU uses three matrices instead of two,\nits hidden width is set to `≈ 8/3 · d_model`\n\n(rounded to a multiple of 256) rather than\n`4 · d_model`\n\n, so the parameter count matches a plain MLP — the `⌈8/3·d⌉`\n\nrule. Code:\n`src/model/gqa_gpt.py`\n\n's `SwiGLU`\n\n(`w1`\n\ngate, `w3`\n\nup, `w2`\n\ndown).\n\nInitialisation decides whether the first steps are stable. Three rules:\n\n**Linear and embedding weights:** A small standard deviation keeps activations and the initial logits in a sane range. At init the model should be maximally unsure, so the loss should be`Normal(0, 0.02)`\n\n.`ln(vocab_size) = ln(16384) ≈ 9.70`\n\n; ours starts at**9.85**✓ — the quickest check that the model and loss are wired correctly.** Residual output projections scaled by**(the GPT-2 trick): the attention`1/√(2·n_layer)`\n\n`Wo`\n\nand the SwiGLU`W₂`\n\n— the two matrices that write back into the residual stream — are initialised with`std = 0.02/√(2·n_layer)`\n\n. In a pre-norm net (§6.5) each of the`2·n_layer`\n\nsublayers adds to the residual stream without renormalising the sum, so without this its variance would grow with depth and the logits would blow up; shrinking those writes keeps the residual variance roughly constant from the first block to the last.**No biases anywhere**, and RMSNorm gains start at`1`\n\n.\n\n(For the scaling experiments, muP rescales both this init and the learning rate by width; §10.2.)\n\n`configs/gpt120m_ctx4k.yaml`\n\n:\n\n```\nn_layer: 16      d_model: 768      n_head: 12   n_kv_head: 4     # GQA 12:4\nblock_size: 4096   vocab_size: 16384   ffn hidden: 2048 (SwiGLU)\n```\n\n**d_model 768 / 12 heads / 16 layers**≈ GPT-2-small shape → ~113M params, a size that trains comfortably on one node and is big enough to be fluent.** head_dim 64**(768/12) — the sweet spot FlashAttention is tuned for.** SwiGLU hidden 2048**— the Llama`⌈8/3·d⌉`\n\nrule rounded to a multiple of 256.\n\nCode: `src/train.py`\n\n. Data-parallel across both A30s.\n\nWe launch with `torchrun --nproc_per_node=2`\n\n, which starts one process per GPU.\nDistributedDataParallel keeps a replica of the model on each GPU; each processes a\ndifferent micro-batch, and gradients are all-reduced (averaged) across GPUs before\nthe optimizer step — so both replicas stay identical. The effective batch is the sum\nacross GPUs.\n\n```\ninit_process_group(backend=\"nccl\")\nmodel = torch.compile(model)                       # fuse kernels\nmodel = DDP(model, device_ids=[local_rank])        # replicate + all-reduce grads\n```\n\n**Gradient accumulation** multiplies the batch further: we run several micro-batches\nbefore stepping, and — importantly — only sync gradients on the last one, so the\nin-between backward passes don't pay the all-reduce cost:\n\n```\nfor micro in range(grad_accum):\n    model.require_backward_grad_sync = (micro == grad_accum - 1)\n    with autocast(bfloat16):\n        _, loss = model(x, y)\n    (loss / grad_accum).backward()\ntorch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # clip exploding grads\noptimizer.step()\n```\n\nGlobal batch = `batch(4) × grad_accum(12) × gpus(2) × ctx(4096)`\n\n= 393,216\ntokens/iter. bf16 autocast ([mixed-precision training](https://arxiv.org/abs/1710.03740))\nhalves memory and doubles throughput vs fp32 with no\nloss-scaling needed (bf16 has fp32's exponent range).\n\nDDP has a hard ceiling: it replicates the whole model and its training state on every GPU,\nso the biggest model it can train is the one that fits on one card. Mixed-precision Adam\ncosts about **16 bytes per parameter** — 2 (bf16 weights) + 2 (bf16 grads) + 12 (fp32\noptimizer: master weights 4, momentum 4, variance 4) — so a 7B model needs ~112 GB of\nstate before a single activation, far past one GPU. Beyond that you have to shard the model\nitself. The strategies below are orthogonal axes — large runs combine several — differing\nmainly in what they split and how much they must communicate (N = number of GPUs sharing a\nshard).\n\n| strategy | splits across GPUs | what it cuts | communication | wants |\n|---|---|---|---|---|\nDDP |\nnothing (full replica) | throughput only | all-reduce grads | any link |\nZeRO-1 |\noptimizer states | opt memory → 12/N B/param | ≈ DDP | any |\nZeRO-2 |\n+ gradients | + grad memory → 2/N | ≈ DDP | any |\nZeRO-3 / FSDP |\n+ parameters | all state → ~16/N | ~1.5× DDP | fast interconnect |\nTensor (Megatron) |\nweight matrices, within a layer | params + activations + layer compute | 2 all-reduce per layer, each pass | NVLink, intra-node |\nPipeline |\nlayers (by depth) | params + activations by depth | point-to-point activations | tolerates inter-node |\nSequence / context |\nthe token dimension | activation memory; enables long context | K/V exchange in attention | medium |\nExpert (MoE) |\nexperts (MoE only) | grows params at fixed FLOPs | 2× all-to-all | NVLink, intra-node |\n\n**ZeRO / FSDP — shard the data-parallel redundancy.** Still data parallelism (every rank\nruns the whole model on its own micro-batch), but instead of replicating the training state\nN times it shards it across the N ranks ([Rajbhandari et al. 2019](https://arxiv.org/abs/1910.02054)).\nZeRO-1 shards the optimizer states (the 12 bytes/param) at essentially DDP communication —\nthe cheapest, near-free win. ZeRO-2 also shards gradients (reduce-scatter instead of\nall-reduce, same volume). ZeRO-3 also shards the parameters: each rank keeps 1/N of the\nweights and all-gathers a layer's shard just-in-time in the forward and the backward, then\nfrees it — per-GPU state falls to ~16/N bytes/param (near-linear memory scaling) at about\n1.5× DDP's communication. PyTorch FSDP is ZeRO-3 ([Zhao et al. 2023](https://arxiv.org/abs/2304.11277)).\nWhat ZeRO does not do: shard activations, or make any layer faster — per-rank compute is\nunchanged. Caveat: ZeRO-3's extra parameter all-gathers need a fast interconnect or\nthroughput sags.\n\n**Tensor parallelism (Megatron) — split inside a layer.** Shard individual weight matrices\nso each GPU holds a slice and computes part of every layer, recombining with a collective\n([Shoeybi et al. 2019](https://arxiv.org/abs/1909.08053)). The MLP splits its up-projection\ncolumn-wise and its down-projection row-wise (one all-reduce); attention splits across\nheads, which are independent (one all-reduce on the output projection). It cuts parameters,\nactivations, and per-layer compute together, and speeds up the layer itself. Caveat: two\nall-reduces per layer in the forward and two in the backward sit on the critical path —\nvery heavy communication — so tensor parallelism is only efficient within a node over\nNVLink, and its degree is usually capped at the GPUs in one node (≤ 8). It also needs\nMegatron-style layer code, not a transparent wrapper.\n\n**Pipeline parallelism (GPipe, 1F1B) — split by depth.** Consecutive layers live on\ndifferent GPUs; only the activations at stage boundaries are passed, as small point-to-point\nmessages that tolerate slower inter-node links ([Huang et al. 2018](https://arxiv.org/abs/1811.06965)).\nTo stop GPUs idling, the batch is chopped into micro-batches and pipelined. Caveat: the\npipeline bubble — GPUs sit idle while the pipeline fills and drains, a fraction\n≈ (P−1)/(m + P−1) for P stages and m micro-batches, so you need many micro-batches to\namortize it — and balancing the per-stage load is fiddly.\n\n**Sequence / context parallelism — split the token dimension.** In Megatron form it\ncomplements tensor parallelism by also sharding the norm/dropout regions TP leaves\nreplicated, cutting activation memory further ([Korthikanti et al. 2022](https://arxiv.org/abs/2205.05198)).\nIn the ring-attention / Ulysses form it shards the sequence itself, so each GPU holds part\nof the tokens and exchanges K/V across ranks during attention\n([Liu et al. 2023](https://arxiv.org/abs/2310.01889)) — the route to very long contexts\n(hundreds of thousands of tokens) whose O(T) activations or O(T²) attention would never fit\notherwise. Caveat: attention now carries cross-rank communication.\n\n**Expert parallelism — for Mixture-of-Experts only.** In an MoE the feed-forward is replaced\nby many expert FFNs and each token is routed to a few (top-k). Expert parallelism puts\ndifferent experts on different GPUs: an all-to-all dispatches each token to its expert's GPU,\nand another all-to-all brings the result back ([Fedus et al. 2021](https://arxiv.org/abs/2101.03961)).\nThis decouples parameter count from per-token compute — scale to trillions of parameters at\nfixed FLOPs per token. Caveats: all-to-all is bandwidth-hungry; token loads across experts\nare uneven (stragglers), needing a load-balancing loss and per-expert capacity caps; and it\napplies only if you chose an MoE architecture.\n\n**Combining them, and choosing by size.** The axes multiply — total GPUs =\nDP × TP × PP (× expert × context) — and the rule that matters is to put the most\ncommunication-hungry axis on the fastest link ([Narayanan et al. 2021](https://arxiv.org/abs/2104.04473)):\ntensor parallelism (and MoE all-to-all) stay inside a node on NVLink; pipeline spans nodes,\nsince its point-to-point transfers tolerate slower links; data parallelism with ZeRO wraps\nthe outside to scale throughput and shard whatever state is left. Activation recomputation\n(gradient checkpointing) trades recompute for activation memory and composes with all of\nthem — usually the first lever to pull. Rational defaults for a dense model:\n\n**≤ ~1B (fits one GPU with state):** plain DDP. Our 113M lives here.**Fits, but state is tight or you want a bigger batch:** ZeRO-1/2 — near-free memory at DDP-level communication.**Doesn't fit one GPU, single node (~a few B–13B):** ZeRO-3 / FSDP to shard everything across the node (simplest), or tensor parallelism ≤ 8 if you are compute-bound.**~10B–100B, multi-node:** 3D — tensor parallel within each node, pipeline across nodes, data parallel / ZeRO on top; FSDP alone starts to choke on inter-node all-gathers here.**100B+ or very long context:** full 3D plus sequence/context parallelism, and expert parallelism if the model is MoE.\n\nLinear warmup (400 iters) then cosine decay ([Loshchilov & Hutter, 2016](https://arxiv.org/abs/1608.03983)) to a floor. Warmup avoids early\ndivergence while Adam's statistics are cold; cosine decay squeezes out the last\nimprovements at low LR. Plotting both against the training step makes the schedule\nand its effect legible on one figure:\n\nThe LR (orange) ramps up over the first 400 steps then follows a cosine to its floor, while the loss (blue) crashes during warmup and keeps refining through the long low-LR tail — which is where the final best val loss came from.\n\nTraining memory has a fixed part and a batch-dependent part.\n\n**Fixed (independent of batch):** parameters + gradients + Adam's two moments.\nWith fp32 master weights and fused AdamW ([Loshchilov & Hutter, 2017](https://arxiv.org/abs/1711.05101)):\n\n```\nparams (fp32)   113M × 4 B = 0.45 GB\ngrads  (fp32)   113M × 4 B = 0.45 GB\nAdam m,v (fp32) 113M × 8 B = 0.90 GB\n                            ─────────\nfixed ≈ 1.8 GB\n```\n\n**Batch-dependent (activations):** grows with `batch × seq_len × n_layer × d_model`\n\n.\nThis is the hard-to-predict part (FlashAttention, autocast and `torch.compile`\n\nall\nchange it), so we measure it empirically with a one-shot probe — run a single\nforward+backward and read the peak:\n\n```\ntorch.cuda.reset_peak_memory_stats()\nloss = model(x, y)[1]; loss.backward(); opt.step()\nprint(torch.cuda.max_memory_allocated() / 1e9, \"GB\")\n```\n\nThat's how we picked the batch size: batch 4 @ 4096 → 15.2 GB (fits 24 GB with headroom); batch 8 @ 4096 OOMs. This probe-first habit beats guessing every time.\n\nMeasured at run time: ~12 GB/GPU, both at 100% utilization, ~2.9 s/iter.\n\nNothing here is A30-specific. Because DDP replicates the whole model on each GPU, a\nsingle 24 GB card already holds a full training copy — so this runs unchanged on a single\nRTX 3090 or RTX 4090 (drop `torchrun`\n\n; `train.py`\n\ntakes the single-process path when no\n`RANK`\n\nis set). Memory is dominated by activations, which scale with batch × context, so a\nsmaller card just uses a smaller micro-batch; gradient accumulation then restores the same\n393,216-token global batch (it runs micro-batches sequentially, so it costs time, not\nmemory). Measured peaks:\n\n| card | setting | peak VRAM | command |\n|---|---|---|---|\n| 24 GB (RTX 3090 / 4090 / A5000) | batch 4 @ 4096 (shipped) | ~16 GB | `python -m src.train --config configs/gpt120m_ctx4k.yaml` |\n| 16 GB | batch 2 @ 4096 | 8.9 GB | `… --batch-size 2 --grad-accum 48` |\n| 12 GB | batch 1 @ 4096 | 5.3 GB | `… --batch-size 1 --grad-accum 96` |\n\nAmpere (3090, A5000) and Ada (4090) run on the pinned `torch 2.6.0+cu124`\n\nas-is; the\nnewest Blackwell cards need a CUDA-12.8 / PyTorch ≥2.7 build. Rough wall-clock for the full\n8,000-iter run (measured on A30, spec-scaled for the rest): ~6.5 h on 2× A30, ~13 h on one\nA30, ~10–13 h on a 4090, and ~20–30 h on a 3090 or A5000 — throughput scales with the\ncard, but the recipe is identical.\n\nCode: `src/inference.py`\n\n(engine + tests), `src/sample.py`\n\n(simple sampler).\n\nNaive generation recomputes the whole sequence every step — O(T²). With a KV-cache we keep each layer's past keys/values and feed only the new token each step — O(T), so throughput stays flat (~83 tok/s) as the sequence grows toward 4096, where the naive path degrades to ~63 tok/s. Correctness is verified against the naive path (greedy output is byte-identical).\n\nGeneration streams token-by-token (`generate_stream`\n\nyields one token per step; the\nCLI prints with `flush=True`\n\n), so you see text ~176 ms after hitting enter instead of\nwaiting for the whole block. Decoding uses a \"decode-the-running-list, emit-the-suffix\"\ntrick so multi-byte characters that span BPE tokens render correctly.\n\nThe sampler also supports anti-repetition controls (on by default in the CLIs):\n`--repetition-penalty`\n\n(downweights already-seen tokens) and `--no-repeat-ngram`\n\n(hard-bans repeating n-grams). At 0.1B params the base model occasionally falls\ninto degenerate loops (e.g. inside a paper table); these decisively fix it —\n`Mj Mj Mj…`\n\n→ a coherent citation — without changing correctness (both default\nto off inside the model, on in the CLI).\n\n`generate_annotated()`\n\nrecords, for every emitted token, the raw probability the\nmodel assigned it and the top-8 alternatives. Shading a real sample by that\nconfidence makes the model's behavior visible — and shows what temperature\nsampling does:\n\nDark = confident (`Mw 9.0`\n\n, `the`\n\n, `earthquake`\n\n), light = uncertain. The bottom\npanel explains one light token: the model's top choice was `area`\n\n(p 0.59), but\nwith temperature 0.8 it sampled the rarer `plane`\n\n(p 0.01) — \"rupture plane\",\nstill perfectly valid seismology. This is exactly the exploration/quality trade-off\nthat `--temperature`\n\ncontrols.\n\nThe decisive test (`--test`\n\n, `context_utilization`\n\n): measure loss by position within a\nfull window. If long context helps, later positions — which have more preceding tokens\nto condition on — should be predicted better:\n\n| position in window | mean loss |\n|---|---|\n| 0–64 (little context) | 2.89 |\n| 1024–2048 | 2.22 |\n| 2048–4096 (full context) | 2.16 |\n\n−25% from early to late — the model demonstrably exploits the long context.\n\n```\npython -m src.inference --test           # generation + perplexity + context-utilization\npython -m src.inference --interactive    # streaming REPL\npython -m src.inference --perplexity-text \"...\"   # score any text\n```\n\nThe token embeddings encode real domain structure — a 2-D PCA cleanly separates\nseismology terms from everyday words, and nearest neighbours are sensible\n(`tsunami → liquefaction, landslide, evacuation`\n\n):\n\nAnd different attention heads specialize — some are local (previous-token), some are \"attention-sink\" heads, some spread broadly — all causal (lower-triangular):\n\nA nice sanity check of specialization: perplexity on a seismology sentence is ~13–16, but on generic English it is ~600–900 — the model is sharply tuned to its domain.\n\n```\nnanogpt_seis/\n├── README.md                  ← you are here\n├── configs/                   ← YAML model+training configs\n│   ├── gpt120m.yaml           (1024-ctx baseline)\n│   ├── gpt120m_ctx4k.yaml     (4096-ctx, the main model)\n│   └── gpt_small.yaml         (~32M, for fast pipeline iteration)\n├── src/\n│   ├── crawl/                 ← Stage 1: data crawlers\n│   │   ├── common.py          (Doc schema, cached polite HTTP)\n│   │   ├── fulltext.py        (Crossref + Unpaywall, concurrent, resumable)\n│   │   ├── preprints.py       (arXiv + EarthArXiv)\n│   │   ├── general.py         (Wikipedia + FineWeb-Edu general-text mix)\n│   │   ├── wikipedia.py  substack.py  arxiv_openalex.py\n│   ├── process/               ← Stage 2: clean / filter / dedup / split\n│   │   ├── clean.py  dedup.py  build_corpus.py\n│   ├── tokenizer/             ← Stage 3: BPE\n│   │   ├── train_bpe.py  encode.py\n│   ├── model/gqa_gpt.py       ← Stage 4: the Transformer\n│   ├── train.py               ← Stage 5: DDP training loop\n│   ├── inference.py sample.py ← Stage 6: serving (streaming, KV-cache, anti-repeat)\n│   ├── compare_models.py      ← v1-vs-v2 A/B (bits-per-byte)\n│   ├── plot_training.py       ← training-curve plots\n│   ├── scaling/               ← Stage 10: IsoFLOP scaling-law sweep (muP)\n│   └── figures/               ← the diagrams in this README\n├── assets/                    ← generated figures\n├── data/                      ← raw / processed / tokenized (git-ignored)\n└── checkpoints/               ← weights, logs, plots (weights git-ignored)\n```\n\nRegenerate every figure in this README (each writes `.png`\n\n+ `.pdf`\n\n):\n\n```\n# diagrams (graphviz) + data plots (matplotlib) — no GPU\npython -m src.figures.workflow\npython -m src.figures.architecture\npython -m src.figures.gqa_vs_mha\npython -m src.figures.training_curves\npython -m src.figures.corpus_composition\npython -m src.figures.rope\n# analysis figures that run the model — need a trained ckpt + GPU\nCUDA_VISIBLE_DEVICES=0 python -m src.figures.generation_example\nCUDA_VISIBLE_DEVICES=0 python -m src.figures.embedding_space\nCUDA_VISIBLE_DEVICES=0 python -m src.figures.bpb_comparison\nCUDA_VISIBLE_DEVICES=0 python -m src.figures.context_utilization\nCUDA_VISIBLE_DEVICES=0 python -m src.figures.attention_map\n```\n\nCode: `src/scaling/`\n\n. The 113M model is a single point; the interesting question is\nthe shape of the curve. This is a small, self-contained harness that measures the\ncompute-optimal frontier at 3M–85M scale, the [IsoFLOP](https://arxiv.org/abs/2203.15556)\nway (Hoffmann et al. 2022): train a family of sizes under fixed compute budgets and\nsee which size wins at each budget.\n\n**Vary only N and D.** Context length (1024) and the global batch (131,072 tokens/iter) are fixed across every run, so the only moving parts are model size N and data D. Each run then costs exactly`C ≈ 6·N·D`\n\nFLOPs — that is what makes a budget an \"IsoFLOP\" set of points.**N is non-embedding parameters.** With a fixed 16k tied embedding, the vocab matrix would dominate tiny models and distort the power law, so N excludes it (Kaplan et al. convention). The tokenizer and embedding size are identical in every run.**One profile per budget.** For each C, every model whose implied`D = C/(6N)`\n\nis trainable and under ~4 epochs of the corpus is run; the size with the lowest loss on that profile is the compute-optimal N_opt(C). Connecting the minima across budgets gives the frontier, and its slope is the scaling exponent.\n\nThe family spans a ~165× range in size (head_dim fixed at 64, so d_model = 64·n_head):\n\n| model | n_layer | d_model | ~non-emb N |\n|---|---|---|---|\n| xxs | 3 | 128 | ~0.7M |\n| xs | 4 | 256 | ~3.1M |\n| s | 6 | 384 | ~9.4M |\n| m | 8 | 512 | ~24M |\n| l | 10 | 640 | ~44M |\n| xl | 12 | 768 | ~76M |\n| 2xl | 14 | 896 | ~122M |\n\nDefault budgets span `6e15`\n\n–`5.4e17`\n\n(five, a 90× range in compute) over these seven\nsizes → 25 cells after the data-window gate. The sweep is resumable (a `DONE`\n\nmarker per\nfinished run), so widening the grid only trains the new cells.\n\nThe confound in any size sweep is the learning rate: the best LR shifts with width, so\nreusing one LR across a 24× range would penalise some sizes and pollute the fit.\nMaximal-Update Parametrization ([muP, Yang et al. 2022](https://arxiv.org/abs/2203.03466))\nfixes this — tune the LR once at a base width and it transfers to the rest. For Adam the\ncore rule is that hidden-matrix weights (whose fan-in grows with width) get a learning\nrate scaled by `base_width / d_model`\n\n, while embeddings and norms keep the base LR; the\nsame weights are initialised smaller by `sqrt(base_width / d_model)`\n\n, and the output\nlogits are scaled by `base_width / d_model`\n\nto keep their magnitude width-stable.\n\n```\n# src/model/gqa_gpt.py — muP is a per-group LR multiplier + matched init\nmult    = d_model / mup_base_width\nlr_mult = 1 / mult if name.endswith(_HIDDEN_SUFFIXES) else 1.0  # embeds/norms: 1.0\n#   init:   hidden weights *= sqrt(1 / mult)\n#   forward: logits *= 1 / mult\n```\n\nBecause head_dim is fixed across the sweep, muP's `1/head_dim`\n\nattention-scale change is\nunnecessary — the usual `1/sqrt(head_dim)`\n\nis already width-independent here. muP is\nopt-in (`mup: true`\n\nin the config) and an exact no-op at the base width, so the 113M\nmodel and every earlier result are unchanged. One caveat: with tied input/output\nembeddings the shared matrix takes the embedding parametrization (constant LR), so the\noutput multiplier — not a separate readout LR — carries muP's width stabilisation.\n\n```\n# 1. write configs/scaling/*.yaml + manifest, and print the run matrix\npython -m src.scaling.run_sweep --generate\n\n# 2. launch pending runs on 2 GPUs (resumable: DONE markers + --resume)\npython -m src.scaling.run_sweep --run --nproc 2      # --budgets 2e16 6e16 to subset\n\n# 3. fit the frontier, then plot it\npython -m src.scaling.fit                            # → checkpoints/scaling/scaling_fit.json\npython -m src.figures.scaling_laws                   # → assets/scaling_laws.png\n```\n\n`fit`\n\nfits a parabola to each budget's loss-vs-N profile, takes the vertex as the\ncompute-optimal (N_opt, D_opt, L_opt), and then fits `N_opt ∝ C^a`\n\nand `D_opt ∝ C^b`\n\nacross budgets — Chinchilla predicts `a ≈ b ≈ 0.5`\n\n. The figure shows the IsoFLOP\nprofiles (loss vs N, one valley per budget) and the fitted frontier. With the exponents\nin hand you can check whether the actual 113M / 823M-token run sits on the frontier you\nmeasured, or off it.\n\nAcross five budgets over a 90× compute range (left), each profile traces a loss valley: at fixed compute a too-small model underfits and a too-large one is undertrained, and the minimum marches to larger N as compute grows. From the parabola vertices:\n\n| budget C (FLOPs) | N_opt (non-emb) | D_opt (tokens) | L_opt | tokens/param |\n|---|---|---|---|---|\n| 6e15 | 2.0M | 511M | 3.55 | 261 |\n| 2e16 | 5.7M | 587M | 3.29 | 104 |\n| 6e16 | 10.8M | 928M | 3.00 | 86 |\n| 1.8e17 | 25.0M | 1.20B | 2.94 | 48 |\n| 5.4e17 | 44.3M | 2.03B | 2.75 | 46 |\n\nFitting the frontier gives **N_opt ∝ C^0.69** and **D_opt ∝ C^0.31** (they sum to 1, as\nthey must). Two readings:\n\n- The frontier is real and tight — five points on a clean power law (right panel) — but its slope is steeper than Chinchilla's ≈ 0.5: over this 0.7M–122M, 6e15–5.4e17 range the marginal FLOP buys more by growing the model than by training it longer. A fixed 1024-token context (which does not grow with the model), the small scale, short runs, and a single muP-transferred base LR all plausibly steepen it, so the exact exponent isn't universal.\n- tokens/param at the optimum falls from ~260 at the smallest budget toward ~46 at the largest — the more compute, the closer to Chinchilla-like ratios, as expected if the frontier only approaches its asymptotic slope at larger scale.\n\nThe practical implication at this scale: given a compute budget, under-sizing the model costs more than under-training it — spend the marginal FLOP on parameters before extra epochs. That is the opposite emphasis from the deliberately over-trained \"inference-optimal\" small models often shipped today, which leave the compute-optimal frontier on purpose to cut serving cost — a reminder that \"optimal\" depends on whether you are optimising training compute or lifetime inference cost.\n\n**Where the real 113M model lands.** Overlaying the actual run (red star) — N = 100.7M\nnon-embedding, D = 3.15B tokens processed, C ≈ 1.9e18 FLOPs — the extrapolated frontier\npredicts a compute-optimal size N_opt ≈ 117M there, so the real model is **0.86×\ncompute-optimal**: slightly smaller than ideal, having seen a bit more data per parameter\nthan the frontier prefers (31 vs ~46 tokens/param). It sat close to the measured\nfrontier, erring marginally toward too-small / over-trained — reassuring for a model\nsized by hand.\n\nCaveats, plainly: the top budget's minimum is a left-boundary argmin (the smaller models that would pin it fall outside the under-4-epoch data window at 5.4e17), so it may over-state N_opt there; and the real-model point is a 3.5× extrapolation beyond the largest budget, at a different context length (4096) on the mixed corpus, so only its (C, N) position is comparable, not its loss. Read the 0.86× as \"near-optimal, slightly small,\" not a precise verdict. Widening the range further (more budgets, an even bigger model) would firm it up.\n\n**Architecture & attention**\n\n- Vaswani et al., \"Attention Is All You Need\" (2017) — the Transformer.\n[arXiv:1706.03762](https://arxiv.org/abs/1706.03762) - Su et al., \"RoFormer: Enhanced Transformer with Rotary Position Embedding\" (2021) — RoPE.\n[arXiv:2104.09864](https://arxiv.org/abs/2104.09864) - Shazeer, \"Fast Transformer Decoding: One Write-Head Is All You Need\" (2019) — MQA.\n[arXiv:1911.02150](https://arxiv.org/abs/1911.02150) - Ainslie et al., \"GQA: Training Generalized Multi-Query Transformer Models…\" (2023).\n[arXiv:2305.13245](https://arxiv.org/abs/2305.13245) - Dao et al., \"FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness\" (2022).\n[arXiv:2205.14135](https://arxiv.org/abs/2205.14135)· Dao, \"FlashAttention-2\" (2023).[arXiv:2307.08691](https://arxiv.org/abs/2307.08691) - Zhang & Sennrich, \"Root Mean Square Layer Normalization\" (2019) — RMSNorm.\n[arXiv:1910.07467](https://arxiv.org/abs/1910.07467) - Shazeer, \"GLU Variants Improve Transformer\" (2020) — SwiGLU.\n[arXiv:2002.05202](https://arxiv.org/abs/2002.05202) - Press & Wolf, \"Using the Output Embedding to Improve Language Models\" (2016) — weight tying.\n[arXiv:1608.05859](https://arxiv.org/abs/1608.05859)\n\n**Tokenization**\n\n- Sennrich et al., \"Neural Machine Translation of Rare Words with Subword Units\" (2015) — BPE for NMT.\n[arXiv:1508.07909](https://arxiv.org/abs/1508.07909) - Wang et al., \"Neural Machine Translation with Byte-Level Subwords\" (2019) — byte-level BPE (popularized by GPT-2).\n[arXiv:1909.03341](https://arxiv.org/abs/1909.03341) - P. Gage, \"A New Algorithm for Data Compression\" (1994) — the original BPE.\n\n**Language models**\n\n- Radford et al., \"Improving Language Understanding by Generative Pre-Training\" (2018) — GPT.\n[OpenAI](https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf) - Radford et al., \"Language Models are Unsupervised Multitask Learners\" (2019) — GPT-2.\n[OpenAI](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) - Brown et al., \"Language Models are Few-Shot Learners\" (2020) — GPT-3.\n[arXiv:2005.14165](https://arxiv.org/abs/2005.14165) - Touvron et al., \"LLaMA: Open and Efficient Foundation Language Models\" (2023).\n[arXiv:2302.13971](https://arxiv.org/abs/2302.13971) - Touvron et al., \"Llama 2: Open Foundation and Fine-Tuned Chat Models\" (2023).\n[arXiv:2307.09288](https://arxiv.org/abs/2307.09288) - Meta AI, \"The Llama 3 Herd of Models\" (2024).\n[arXiv:2407.21783](https://arxiv.org/abs/2407.21783)\n\n**Training & scaling**\n\n- Loshchilov & Hutter, \"Decoupled Weight Decay Regularization\" (2017) — AdamW.\n[arXiv:1711.05101](https://arxiv.org/abs/1711.05101) - Loshchilov & Hutter, \"SGDR: Stochastic Gradient Descent with Warm Restarts\" (2016) — cosine schedule.\n[arXiv:1608.03983](https://arxiv.org/abs/1608.03983) - Micikevicius et al., \"Mixed Precision Training\" (2017).\n[arXiv:1710.03740](https://arxiv.org/abs/1710.03740) - Yang et al., \"Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer\" (2022) — muP / muTransfer.\n[arXiv:2203.03466](https://arxiv.org/abs/2203.03466) - Hoffmann et al., \"Training Compute-Optimal Large Language Models\" (2022) — Chinchilla.\n[arXiv:2203.15556](https://arxiv.org/abs/2203.15556) - Muennighoff et al., \"Scaling Data-Constrained Language Models\" (2023) — how far to repeat data.\n[arXiv:2305.16264](https://arxiv.org/abs/2305.16264) - Lee et al., \"Deduplicating Training Data Makes Language Models Better\" (2021).\n[arXiv:2107.06499](https://arxiv.org/abs/2107.06499) - Broder, \"On the Resemblance and Containment of Documents\" (1997) — MinHash.\n\n**Parallelism & systems**\n\n- Rajbhandari et al., \"ZeRO: Memory Optimizations Toward Training Trillion Parameter Models\" (2019).\n[arXiv:1910.02054](https://arxiv.org/abs/1910.02054) - Shoeybi et al., \"Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism\" (2019).\n[arXiv:1909.08053](https://arxiv.org/abs/1909.08053) - Huang et al., \"GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism\" (2018).\n[arXiv:1811.06965](https://arxiv.org/abs/1811.06965) - Narayanan et al., \"Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM\" (2021) — 3D parallelism.\n[arXiv:2104.04473](https://arxiv.org/abs/2104.04473) - Korthikanti et al., \"Reducing Activation Recomputation in Large Transformer Models\" (2022) — sequence parallelism + selective checkpointing.\n[arXiv:2205.05198](https://arxiv.org/abs/2205.05198) - Zhao et al., \"PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel\" (2023).\n[arXiv:2304.11277](https://arxiv.org/abs/2304.11277) - Liu et al., \"Ring Attention with Blockwise Transformers for Near-Infinite Context\" (2023).\n[arXiv:2310.01889](https://arxiv.org/abs/2310.01889) - Jacobs et al., \"DeepSpeed-Ulysses: System Optimizations for Extreme Long Sequence Transformer Models\" (2023).\n[arXiv:2309.14509](https://arxiv.org/abs/2309.14509) - Fedus et al., \"Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity\" (2021) — MoE.\n[arXiv:2101.03961](https://arxiv.org/abs/2101.03961) - Lepikhin et al., \"GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding\" (2020) — MoE.\n[arXiv:2006.16668](https://arxiv.org/abs/2006.16668)\n\n**Data & datasets**\n\n- Penedo et al., \"The FineWeb Datasets\" (2024).\n[arXiv:2406.17557](https://arxiv.org/abs/2406.17557)·[🤗 HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) - Wikipedia dump —\n[🤗 wikimedia/wikipedia](https://huggingface.co/datasets/wikimedia/wikipedia) [Crossref REST API](https://www.crossref.org/documentation/retrieve-metadata/rest-api/)·[Unpaywall API](https://unpaywall.org/products/api)·[arXiv API](https://info.arxiv.org/help/api/)·[EarthArXiv](https://eartharxiv.org/)(via[OSF](https://api.osf.io/)) ·[Earthquake Insights](https://earthquakeinsights.substack.com/)\n\n**Tools & learning resources**\n\n- Andrej Karpathy — \"Neural Networks: Zero to Hero\"\n[YouTube](https://www.youtube.com/@AndrejKarpathy)(esp. \"Let's build GPT\" and \"Let's build the GPT Tokenizer\") [karpathy/nanoGPT](https://github.com/karpathy/nanoGPT)·[karpathy/minbpe](https://github.com/karpathy/minbpe)·[jingyaogong/minimind](https://github.com/jingyaogong/minimind)[HuggingFace](https://github.com/huggingface/tokenizers)&`tokenizers`\n\n·`datasets`\n\n(MinHash/LSH) ·`datasketch`\n\n[PyTorch](https://pytorch.org/)(SDPA, DDP)\n\nInspired by Andrej Karpathy's nanoGPT and the minimind project's teach-by-building philosophy. Data courtesy of Crossref, Unpaywall, arXiv, EarthArXiv/OSF, Wikipedia, and the Earthquake Insights Substack — thanks to the open-science infrastructure that makes a project like this possible.\n\n**Code** is MIT — see[LICENSE](/jiazhe868/nanogpt-seis/blob/main/LICENSE).**The corpus is not redistributed.**`data/`\n\nis git-ignored; you regenerate it by running the crawlers. Each source keeps its own terms — arXiv/EarthArXiv (per-paper licenses), Crossref/Unpaywall (open-access PDFs only), Wikipedia (CC BY-SA), FineWeb-Edu (ODC-By), and the Earthquake Insights Substack (author's content). Respect each source's license and robots/rate-limit etiquette; the crawlers send a polite identifying`User-Agent`\n\n/`mailto`\n\n.**Model weights** are a derived artifact, distributed (if at all) via a GitHub Release or the Hugging Face Hub, not committed to git.\n\nMIT — see [LICENSE](/jiazhe868/nanogpt-seis/blob/main/LICENSE).", "url": "https://wpnews.pro/news/i-trained-a-113m-parameter-earthquake-llm-from-absolute-scratch", "canonical_source": "https://github.com/jiazhe868/nanogpt-seis", "published_at": "2026-07-12 21:58:38+00:00", "updated_at": "2026-07-12 22:05:02.236804+00:00", "lang": "en", "topics": ["artificial-intelligence", "large-language-models", "ai-research", "ai-tools", "ai-infrastructure"], "entities": ["nanoGPT-Seis", "NVIDIA A30", "Hugging Face", "jiazhe868/nanogpt_seis", "FineWeb-Edu", "Earthquake Insights", "arXiv", "EarthArXiv"], "alternates": {"html": "https://wpnews.pro/news/i-trained-a-113m-parameter-earthquake-llm-from-absolute-scratch", "markdown": "https://wpnews.pro/news/i-trained-a-113m-parameter-earthquake-llm-from-absolute-scratch.md", "text": "https://wpnews.pro/news/i-trained-a-113m-parameter-earthquake-llm-from-absolute-scratch.txt", "jsonld": "https://wpnews.pro/news/i-trained-a-113m-parameter-earthquake-llm-from-absolute-scratch.jsonld"}}