Skip to content

Commit 3182b47

Browse files
charlesmartin14charlesmartin14
committed
add nanochat d12 reference baseline notebook
1 parent 7cbb948 commit 3182b47

1 file changed

Lines changed: 166 additions & 0 deletions

File tree

‎baseline/notebooks/NanoChat_D12_Reference_Baseline.ipynb‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# nanochat d12 — strong reference baseline\n",
8+
"\n",
9+
"This notebook runs a **pinned upstream nanochat d12** baseline suitable for RG optimizer comparisons.\n",
10+
"\n",
11+
"nanochat uses depth 12 as its reference/tuning scale. At d12, upstream chooses a 768-wide transformer and transfers the tuned optimization recipe using its scaling-law / μP-style rules. We preserve the upstream architecture, initialization, separate embedding/unembedding/scalar/matrix learning rates, hybrid AdamW+Muon optimizer, automatic token horizon and batch size, LR/momentum/weight-decay schedules, BOS-aligned packing, and tokenizer pipeline.\n",
12+
"\n",
13+
"Three independent seeds are run. The RG wrapper pins nanochat commit `92d63d4e8bb4df75c3b71618f31ddde2378b2bcd` and only changes the upstream hard-coded seed 42 to read `NANOCHAT_SEED`; it does not replace nanochat training logic.\n"
14+
]
15+
},
16+
{
17+
"cell_type": "code",
18+
"execution_count": null,
19+
"metadata": {},
20+
"outputs": [],
21+
"source": [
22+
"from pathlib import Path\n",
23+
"import os, sys\n",
24+
"import pandas as pd\n",
25+
"\n",
26+
"ROOT = None\n",
27+
"for path in [Path.cwd(), *Path.cwd().parents]:\n",
28+
" candidate = path / 'baseline'\n",
29+
" if (candidate / 'rg_baselines').is_dir():\n",
30+
" ROOT = candidate\n",
31+
" break\n",
32+
" if (path / 'rg_baselines').is_dir():\n",
33+
" ROOT = path\n",
34+
" break\n",
35+
"if ROOT is None:\n",
36+
" raise RuntimeError('Run from a clone of CalculatedContent/rg_optimizers.')\n",
37+
"ROOT = ROOT.resolve()\n",
38+
"if str(ROOT) not in sys.path:\n",
39+
" sys.path.insert(0, str(ROOT))\n",
40+
"\n",
41+
"from rg_baselines.nanochat_reference import (\n",
42+
" DEFAULT_NANOCHAT_SEEDS, NANOCHAT_COMMIT, NanoChatD12Config,\n",
43+
" analyze_weightwatcher_checkpoints, collect_metrics, ensure_checkout,\n",
44+
" ensure_environment, prepare_data, run_seed,\n",
45+
")\n",
46+
"\n",
47+
"WORK_ROOT = Path(os.environ.get('RG_NANOCHAT_WORK_ROOT', ROOT / 'nanochat_work')).expanduser().resolve()\n",
48+
"CHECKOUT = WORK_ROOT / 'upstream_nanochat'\n",
49+
"CACHE = Path(os.environ.get('NANOCHAT_BASE_DIR', WORK_ROOT / 'cache')).expanduser().resolve()\n",
50+
"RUN_DIR = Path(os.environ.get('RG_BASELINE_RUN_ROOT', ROOT / 'runs')).expanduser().resolve() / 'nanochat_d12_reference'\n",
51+
"WORK_ROOT.mkdir(parents=True, exist_ok=True)\n",
52+
"RUN_DIR.mkdir(parents=True, exist_ok=True)\n",
53+
"CONFIG = NanoChatD12Config()\n",
54+
"SEEDS = DEFAULT_NANOCHAT_SEEDS\n",
55+
"NPROC_PER_NODE = int(os.environ.get('RG_NANOCHAT_NPROC', '8'))\n",
56+
"print('Pinned nanochat commit:', NANOCHAT_COMMIT)\n",
57+
"print('d12 width:', CONFIG.model_dim, 'seeds:', SEEDS, 'GPU processes:', NPROC_PER_NODE)\n",
58+
"display(pd.DataFrame([CONFIG.__dict__]))\n"
59+
]
60+
},
61+
{
62+
"cell_type": "markdown",
63+
"metadata": {},
64+
"source": [
65+
"## 1. Pin nanochat and create its environment\n",
66+
"\n",
67+
"The checkout is detached at the audited commit. nanochat's own `uv` environment and dependency configuration are used.\n"
68+
]
69+
},
70+
{
71+
"cell_type": "code",
72+
"execution_count": null,
73+
"metadata": {},
74+
"outputs": [],
75+
"source": [
76+
"CHECKOUT = ensure_checkout(CHECKOUT)\n",
77+
"ensure_environment(CHECKOUT, gpu=True)\n",
78+
"print('nanochat checkout:', CHECKOUT)\n"
79+
]
80+
},
81+
{
82+
"cell_type": "markdown",
83+
"metadata": {},
84+
"source": [
85+
"## 2. Prepare upstream data and tokenizer\n",
86+
"\n",
87+
"This follows nanochat's miniseries setup: 1000 dataset shards and a 32,768-token tokenizer trained from up to 2B characters. Run once; later replicates reuse the cache.\n"
88+
]
89+
},
90+
{
91+
"cell_type": "code",
92+
"execution_count": null,
93+
"metadata": {},
94+
"outputs": [],
95+
"source": [
96+
"prepare_data(CHECKOUT, CACHE, CONFIG)\n",
97+
"print('nanochat cache:', CACHE)\n"
98+
]
99+
},
100+
{
101+
"cell_type": "markdown",
102+
"metadata": {},
103+
"source": [
104+
"## 3. Run three d12 reference replicates\n",
105+
"\n",
106+
"These are full reference runs, not smoke tests. Upstream nanochat computes the training horizon from 12 tokens per scaling parameter, auto-computes total token batch size, and applies its internal depth/batch LR and weight-decay scaling. Checkpoints and validation are emitted every 250 steps; CORE is evaluated at the final step.\n"
107+
]
108+
},
109+
{
110+
"cell_type": "code",
111+
"execution_count": null,
112+
"metadata": {},
113+
"outputs": [],
114+
"source": [
115+
"logs = []\n",
116+
"for seed in SEEDS:\n",
117+
" log_path = run_seed(CHECKOUT, CACHE, RUN_DIR, CONFIG, seed=seed, nproc_per_node=NPROC_PER_NODE)\n",
118+
" logs.append((seed, log_path))\n",
119+
"\n",
120+
"metrics = collect_metrics(logs, RUN_DIR / 'training_metrics_all_seeds.csv')\n",
121+
"display(metrics.tail(30))\n"
122+
]
123+
},
124+
{
125+
"cell_type": "markdown",
126+
"metadata": {},
127+
"source": [
128+
"## 4. Offline WeightWatcher analysis\n",
129+
"\n",
130+
"WeightWatcher runs after training so spectral diagnostics do not contaminate timed baseline performance. Every saved checkpoint is analyzed with `ERG=True, randomize=True`; all returned columns are retained, including alpha, randomized correlation-trap fields, and ERG metrics when supplied by WeightWatcher.\n"
131+
]
132+
},
133+
{
134+
"cell_type": "code",
135+
"execution_count": null,
136+
"metadata": {},
137+
"outputs": [],
138+
"source": [
139+
"spectral_frames = []\n",
140+
"for seed in SEEDS:\n",
141+
" frame = analyze_weightwatcher_checkpoints(\n",
142+
" CHECKOUT, CACHE, seed=seed, output_csv=RUN_DIR / f'weightwatcher_seed{seed}.csv'\n",
143+
" )\n",
144+
" spectral_frames.append(frame)\n",
145+
"spectral = pd.concat(spectral_frames, ignore_index=True)\n",
146+
"spectral.to_csv(RUN_DIR / 'weightwatcher_all_seeds.csv', index=False)\n",
147+
"display(spectral.tail(50))\n"
148+
]
149+
},
150+
{
151+
"cell_type": "markdown",
152+
"metadata": {},
153+
"source": [
154+
"## Baseline contract\n",
155+
"\n",
156+
"Do not silently update the nanochat commit in an optimizer comparison. A new upstream commit constitutes a new baseline version. The persisted logs, checkpoints, configuration snapshots, training/validation metrics, CORE score, and WeightWatcher diagnostics define the reference control.\n"
157+
]
158+
}
159+
],
160+
"metadata": {
161+
"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
162+
"language_info": {"name": "python", "version": "3"}
163+
},
164+
"nbformat": 4,
165+
"nbformat_minor": 5
166+
}

0 commit comments

Comments
 (0)