Aion v1 — final build

Model, tokenizer, and training code for the 1B parameter build, targeting ~20B tokens (~32-48GB) on Kaggle TPU v5e-8 (primary) with T4 x2 GPU (fallback). Matches configs/aion_v1.yaml exactly — that file is the single source of truth; every script here reads from it rather than hardcoding values.

Structure

aion-final/
├── model/model.py              # 1B param model, TPU+GPU aware, experimental flags off by default
├── tokenizer/setup_tokenizer.py # Llama2 tokenizer + fertility check
├── training/train.py           # main training script, auto-detects TPU/GPU/CPU
├── data_pipeline/               # downloader (HF datasets) + crawler (docs sites)
├── data/                        # EMPTY — this is where all downloaded/crawled data lands
│   ├── code/stack_v2/
│   ├── general/fineweb_edu/
│   ├── general/wikipedia/
│   ├── math/finemath/
│   ├── math/openwebmath/
│   └── docs/technical_docs/
├── configs/aion_v1.yaml        # locked config, single source of truth
└── checkpoints/                 # training checkpoints land here

---

Step 1 — Install dependencies

pip install torch datasets huggingface_hub transformers requests beautifulsoup4 pyyaml

On Kaggle TPU notebooks specifically, also install:

pip install torch_xla

(TPU support via torch_xla — the training script auto-detects it; if not installed, it falls back to GPU/CPU automatically.)

---

Step 2 — Download data from Hugging Face

These are the exact datasets locked in configs/aion_v1.yaml's data mix. Run via the pipeline's downloader — it handles resume, dedup, and sharding automatically (see data_pipeline/README section below for how resume works).

cd data_pipeline

# Download everything (will take a while for the full ~20B token target —
# start with a small slice for Phase 2 validation first, see note below)
python downloader/hf_downloader.py

# Or download one dataset at a time:
python downloader/hf_downloader.py stack_v2_smol      # code, ~45-50% of mix
python downloader/hf_downloader.py finemath_4plus      # math, part of the 20%
python downloader/hf_downloader.py openwebmath          # math, part of the 20%
python downloader/hf_downloader.py fineweb_edu          # general text, ~30-35% of mix
python downloader/hf_downloader.py wikipedia             # general text, ~30-35% of mix

Before running the full pull: confirm exact dataset/config names on huggingface.co directly — names in hf_downloader.py are marked # VERIFY where they may have changed since the research was compiled. Specifically check:

• bigcode/the-stack-v2-dedup (or current equivalent — check bigcode's HF org page)
• HuggingFaceTB/stack-edu
• HuggingFaceTB/finemath (config: finemath-4plus)

Start small first. For Phase 2 validation (per the build plan), pull only a 1-2GB slice before committing to the full run:

python downloader/hf_downloader.py stack_v2_smol   # let it run briefly, then Ctrl+C

The resume system means stopping early and restarting later just continues — no data is lost or corrupted by stopping partway through.

---

Step 3 — What to crawl (technical docs)

The HF datasets above cover code, general text, and math. They do not cover library/API documentation — that's the one piece that needs the crawler in data_pipeline/crawler/doc_crawler.py.

Recommended crawl targets (add these to SEED_SITES in doc_crawler.py, alongside the Python docs and MDN examples already there):

Site	Why	Allowed domain
Python official docs	Already configured	docs.python.org
MDN Web Docs	Already configured	developer.mozilla.org
Node.js API docs	JS/Node coverage, matches your language priority list	nodejs.org
Rust std docs	Rust is in your language priority list	doc.rust-lang.org
Go docs	Go is in your language priority list	go.dev (check current docs path)
React docs	Extremely common in real-world JS code, teaches library usage patterns	react.dev
PostgreSQL docs	SQL is in your priority list, teaches real query patterns	www.postgresql.org
Django docs	Common Python web framework, teaches idiomatic usage	docs.djangoproject.com

Before crawling any of these: check robots.txt and current terms of use — Python docs and MDN are well-documented as crawl-friendly for reasonable automated access; verify the others individually before running at scale, since terms vary and can change.

cd data_pipeline
python crawler/doc_crawler.py python_docs mdn_web_docs

---

Step 4 — Tokenizer fertility check

Run this after you have at least a small sample of data in each category — don't skip this before scaling to the full pull.

cd tokenizer
python setup_tokenizer.py --sample_dir ../data --n_samples 200

This reports tokens/word per category (code, general, math, docs). If code fertility exceeds 1.5, it flags for human review rather than silently proceeding — per project rules, that decision isn't automated.

---

Step 5 — Training

Do not run this on the full dataset until Gate 2 (small-slice validation) has passed. Run on your small slice first:

cd training
python train.py --config ../configs/aion_v1.yaml --data_dir ../data --max_steps 200

--max_steps 200 caps it to a short validation run. Remove that flag only after Gate 2 passes, for the real full-scale run:

python train.py --config ../configs/aion_v1.yaml --data_dir ../data --resume

The script auto-detects TPU vs GPU vs CPU and applies the correct precision automatically (bf16 on TPU, fp16+GradScaler on GPU, fp32 on CPU-emergency-only). Checkpoints save every 500 steps by default (--checkpoint_every), including full optimizer state — a killed session resumes with --resume and continues the loss curve rather than restarting it.

If you ever see a NaN loss message, the script stops immediately and saves a debug checkpoint — per project rules, don't manually restart past a NaN without reviewing what caused it first.

---

Known gaps — still needs decisions before the full run

• Exact HF dataset/config names need final verification (marked # VERIFY in hf_downloader.py)
• Crawl target list above is a starting point — confirm each site's current terms before adding
• Experimental flags (use_mod, gqa_qk_norm, z_loss in the config) are all false by default — only enable after A/B testing on the validation slice per Phase 2 of the build plan
• Fuzzy/MinHash deduplication (beyond the exact-hash dedup already built into the pipeline) is a separate pass — see docs/research/deduplication_pipeline_report.md from the launch kit for the DataTrove-based recommendation