GitHub - jiazhe868/nanogpt-seis

English | 中文 Train a small GPT for earthquake science — the entire LLM lifecycle, from a blank folder to a talking model, explained block by block. Crawl → Clean → Tokenize → Model → Train → Infer, on 2× NVIDIA A30 (48 GB). Six free data sources → crawl → clean/dedup → 16k BPE → 113M GQA+RoPE decoder → 2-GPU DDP training → streaming inference. Each stage is a section below. nanoGPT-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. The corpus mixes earthquake / seismology text (open-access papers via Crossref+Unpaywall, arXiv/EarthArXiv preprints, the "Earthquake Insights" Substack) with general text (Wikipedia + FineWeb-Edu) for plain-language fluency — about 24% domain / 76% general. A focused corpus lets a ~100M-param model become genuinely fluent on a single node, so you can see the whole loop close in a day, not a month. (Why the general mix? See §1.) Status: the pretraining lifecycle is complete — crawl through inference. Table of contents Results at a glance Quick start Stage 1 — Data crawling Stage 2 — Processing & dedup Stage 3 — The BPE tokenizer Stage 4 — The model (RoPE, GQA, …) Stage 5 — Training (DDP, VRAM, LR) Stage 6 — Inference Repository layout Scaling-law experiments 1.

Results at a glance value Corpus 533,248 docs · 485.7M words · 822.7M training tokens (≈2.4:1 general:domain) Model 113M params — decoder-only, GQA + RoPE + RMSNorm + SwiGLU 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 (§7.5) Context length 4096 tokens Training 8,000 iters (~3.8 epochs), ~6.5 h, ~2.9 s/iter Fluency (bits/byte, general text) 0.997 — vs 1.527 for a domain-only base (−35%) Inference KV-cached streaming, ~176 ms to first token, anti-repeat sampler Three findings worth pausing on: 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). 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. Data mix: domain-only (v1) → general + domain (v2) A paper-only base is fluent in paper-register but repetitive/incoherent in plain prose.

Adding ~540M tokens of Wikipedia + FineWeb-Edu (→ ~823M train tokens, ~2.4:1 general:domain, ~3.8 epochs — within the ~4-epoch repeat budget shown to be near-lossless by Muennighoff et al., 2023) gives v2. Measured with bits-per-byte (tokenizer-independent, so v1↔v2 is fair; src/compare_models.py): v2 is far more fluent on general prose (−35% bits/byte) and generates coherent non-earthquake text where v1 emits gibberish, at the cost of some domain sharpness (+22%) — the classic fluency↔specialization trade-off. This fluent base is the right starting point for SFT; base pretraining alone never yields chat. Domain-only weights are kept as checkpoints/ckpt_v1_domain.pt. 2. Quick start 2.1 Try the pretrained model The pretrained 113M checkpoint is hosted on the Hugging Face Hub: jiazhe868/nanogpt_seis. # environment (a working CUDA-12.4 PyTorch; see the note below) conda activate nanogpt_seis pip install -r requirements.txt # download the checkpoint, tokenizer, and model config into the paths expected by # src/inference.py huggingface-cli download jiazhe868/nanogpt_seis \ checkpoints/ckpt.pt \ data/tokenized/tokenizer.json \ data/tokenized/meta.json \ configs/gpt120m_ctx4k.yaml \ --local-dir . python -m src.inference --prompt "The 2011 Tohoku earthquake" 2.2 Reproduce pretraining from scratch # environment (a working CUDA-12.4 PyTorch; see the note below) conda activate nanogpt_seis pip install -r requirements.txt # --- run the whole pipeline --- # earthquake-domain

sources python -m src.crawl.wikipedia --max-pages 500 # earthquake-titled pages python -m src.crawl.fulltext --per-journal 3000 --broad 30000 --workers 64 python -m src.crawl.preprints --arxiv 3000 --eartharxiv 2000 python -m src.crawl.substack --max 500 # general-text mix for fluency (~540M tokens: Wikipedia + FineWeb-Edu) python -m src.crawl.general --wiki-tokens 300000000 --fineweb-tokens 240000000 python -m src.process.build_corpus --val-frac 0.005 # clean · dedup · split python -m src.tokenizer.train_bpe --vocab-size 16384 # train the tokenizer python -m src.tokenizer.encode # → uint16 shards torchrun --standalone --nproc_per_node=2 \ -m src.train --config configs/gpt120m_ctx4k.yaml # train on 2 GPUs python -m src.inference --prompt "The 2011 Tohoku earthquake" # streams live ⚠ Environment gotcha. A PyTorch built for a newer CUDA than your driver will silently report cuda.is_available() == False and fall back to CPU. Verify with python -c "import torch; print(torch.cuda.is_available())" — it must print True. This project uses torch 2.6.0+cu124 to match a CUDA-12.5 driver. 3. Stage 1 — Data crawling Goal: assemble a large, legal, full-text earthquake corpus from free sources.

Code: src/crawl/. 3.1 The sources and how each is fetched source module what we pull mechanism Research papers fulltext.py OA full-text PDFs of earthquake papers Crossref (DOIs) → Unpaywall (OA PDF) → download → extract Preprints preprints.py arXiv + EarthArXiv full text arXiv API + OSF/DOI → PDF Wikipedia wikipedia.py pages titled "earthquake" MediaWiki API plaintext extracts Substack substack.py "Earthquake Insights" articles archive API + HTML body parse General text general.py Wikipedia + FineWeb-Edu (fluency mix) HF datasets streaming to a token budget Why 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.) Every document is normalized to one schema (src/crawl/common.py): @dataclass class Doc: source: str # "fulltext" | "arxiv" | "wikipedia" | ... id: str # stable per-source id (used for dedup) title: str text: str # the cleaned body we will tokenize url: str = "" date: str = "" extra: dict = field(default_factory=dict) # venue, cited_by, full_text, ... 3.2 How the web crawling actually works Finding the papers (Crossref). Crossref indexes ~150M scholarly works with a free, generous API.

We page through earthquake journal-articles with a deep cursor (no offset limit), filtered by journal ISSN: # src/crawl/fulltext.py — iter_crossref() params = {"rows": 1000, "cursor": cursor, "filter": "type:journal-article,issn:0094-8276", # e.g. GRL "query.bibliographic": "earthquake", "mailto": EMAIL} Finding the open PDF (Unpaywall). A DOI is not a PDF. Unpaywall (free, ~100k/day) maps a DOI to its legal open-access copies. We try repository (green OA) locations first — publisher links are frequently bot-blocked stubs: # repository copies download far more reliably than publisher links prio = 0 if loc.get("host_type") == "repository" else 1 Downloading in parallel, politely. PDFs come from many hosts, so we use a thread pool but throttle per host — different servers download concurrently while any single server stays rate-limited: class HostThrottle: # src/crawl/fulltext.py def wait(self, host): with self._guard: # get/create this host's lock host_lock = self._locks.setdefault(host, threading.Lock()) self._last.setdefault(host, 0.0) with host_lock: # serialize only this host delta = self.min - (time.monotonic() - self._last[host]) if delta > 0: time.sleep(delta) self._last[host] = time.monotonic() Validating every download. Not every "OA PDF" is real — many are 5 KB anti-bot landing pages. We accept a download only if it is a real PDF with enough text: if not pdf_bytes.startswith(b"%PDF"): return None # HTML / stub if doc.page_count < 2: return None # cover page if len(text) < min_chars: return None # too little extracted Not wasting the budget. Some journals (Science, Nature) are almost entirely paywalled — scanning thousands of their DOIs would burn the Unpaywall budget for zero full text.

A low-yield abort gate skips a journal once its hit-rate stays under a threshold: if scanned >= abort_after and got_ft / max(1, scanned) < min_hit: break # this venue isn't worth more API calls Resumable. Output is appended line-by-line and already-fetched ids are skipped on restart, so a multi-hour crawl survives interruption: done = _load_done_ids(out) # ids already in the JSONL ... if iid in done: continue # skip work we already have War 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. Full-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. 3.3 The general pattern — a thread-safe BFS crawler The 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. Three rules make it correct: The frontier is a queue.Queue, which is internally synchronised: get() and put() are atomic, and get() blocks 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. Termination uses Queue.join() / task_done(). 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.