Add Level 0 three-by-three epoch diagnostics notebook

level_0_wwpgd/notebooks/04_repeat_3x3_epoch_diagnostics.ipynb

@@ -0,0 +1 @@
+{"cells":[{"cell_type":"markdown","id":"728017ed","metadata":{},"source":["# Level 0 repeat: 3 AdamW vs 3 AdamW + WWPGD\n","\n","Runs the original **4-block, 4-head, width-128, context-256** Level 0 experiment for paired seeds `1337, 2027, 4099`. It then plots train/validation/test cross-entropy, perplexity and token accuracy, plus layerwise WeightWatcher alpha, KS `D`, alpha uncertainty, fitted-tail support, ERG gap, and WWPGD intervention size.\n","\n","The x-axis is **effective epoch = tokens seen / train-split tokens**. Because training samples random windows, this is a token-budget equivalent rather than a literal without-replacement pass. Checkpoint test metrics are computed after training and never used for optimization or model selection."]},{"cell_type":"code","execution_count":null,"id":"76c2a266","metadata":{},"outputs":[],"source":["from __future__ import annotations\n","import json, math, os, subprocess, sys\n","from pathlib import Path\n","import matplotlib.pyplot as plt\n","import numpy as np\n","import pandas as pd\n","import torch, yaml\n","from IPython.display import display\n","\n","SEEDS=(1337,2027,4099)\n","def repo_root(p=Path.cwd()):\n","    for q in (p.resolve(),*p.resolve().parents):\n","        if (q/\"level_0_baseline\").is_dir() and (q/\"level_0_wwpgd\").is_dir(): return q\n","    raise FileNotFoundError(\"Run inside nanogpt-experiments\")\n","ROOT=repo_root()\n","BASE_CFG=ROOT/\"level_0_baseline/configs/level0.yaml\"\n","WW_CFG=ROOT/\"level_0_wwpgd/configs/level0.yaml\"\n","RUNNER=ROOT/\"scripts/run_isolated_level0_pair.sh\"\n","PAIR=Path(os.getenv(\"NANOGPT_LEVEL0_REPEAT_ROOT\",\"/tmp/nanogpt-level0-repeat-4b4h-3x3\"))\n","DATA=Path(os.getenv(\"NANOGPT_LEVEL0_DATA_ROOT\",\"/tmp/nanogpt-level0-bpe/data\"))\n","ANALYSIS=PAIR/\"analysis\"; ANALYSIS.mkdir(parents=True,exist_ok=True)\n","DEVICE=os.getenv(\"NANOGPT_LEVEL0_REPEAT_DEVICE\", \"cuda\" if torch.cuda.is_available() else (\"mps\" if torch.backends.mps.is_available() else \"cpu\"))\n","b=yaml.safe_load(BASE_CFG.read_text()); w=yaml.safe_load(WW_CFG.read_text())\n","for section in (\"model\",\"training\",\"analysis\"): assert b[section]==w[section]\n","assert {k:b[\"model\"][k] for k in (\"n_layer\",\"n_head\",\"n_embd\",\"block_size\",\"vocab_size\")}=={\"n_layer\":4,\"n_head\":4,\"n_embd\":128,\"block_size\":256,\"vocab_size\":50257}\n","assert w[\"wwpgd\"][\"apply_mode\"]==\"event_projection\" and w[\"wwpgd\"][\"interval\"]==1 and w[\"wwpgd\"][\"target_alpha\"]==2.0\n","display(pd.DataFrame({\"setting\":[\"blocks\",\"heads\",\"width\",\"context\",\"steps\",\"seeds\"],\"value\":[4,4,128,256,b[\"training\"][\"max_steps\"],str(SEEDS)]}))\n","def run(phase):\n","    env=os.environ.copy(); env.update({\"NANOGPT_LEVEL0_PAIRED_SEEDS\":\",\".join(map(str,SEEDS)),\"NANOGPT_LEVEL0_DATA_ROOT\":str(DATA),\"NANOGPT_LEVEL0_DEVICE\":DEVICE,\"NANOGPT_LEVEL0_PAIR_STATE_FILE\":str(PAIR/\"pair_state.txt\")})\n","    subprocess.run([\"bash\",str(RUNNER),phase,str(PAIR)],cwd=ROOT,env=env,check=True)\n","if os.getenv(\"NANOGPT_LEVEL0_REPEAT_RUN\",\"1\")!=\"0\":\n","    run(\"baseline\"); run(\"wwpgd\")\n","run(\"verify\")"]},{"cell_type":"code","execution_count":null,"id":"08dc9f15","metadata":{},"outputs":[],"source":["META=json.loads((DATA/\"meta.json\").read_text()); TRAIN_TOKENS=int(META[\"splits\"][\"train\"])\n","SPECS={\"AdamW\":(PAIR/\"baseline/results\",\"adamw_seed_*\"),\"AdamW + WWPGD\":(PAIR/\"wwpgd/results\",\"adamw_wwpgd_seed_*\")}\n","runs=[]; frames=[]\n","for arm,(root,pattern) in SPECS.items():\n","    for r in sorted(root.glob(pattern)):\n","        seed=int(r.name.rsplit(\"_seed_\",1)[1])\n","        if json.loads((r/\"run_complete.json\").read_text()).get(\"completed\") is True:\n","            runs.append((arm,seed,r))\n","            d=pd.read_csv(r/\"metrics.csv\"); d[\"arm\"]=arm; d[\"seed\"]=seed; d[\"effective_epoch\"]=d.tokens_seen/TRAIN_TOKENS; frames.append(d)\n","assert {(a,s) for a,s,_ in runs}=={(a,s) for a in SPECS for s in SEEDS}\n","metrics=pd.concat(frames,ignore_index=True)\n","\n","def bands(frame, specs, title, name):\n","    fig,axs=plt.subplots(len(specs),1,figsize=(11,3.7*len(specs)),squeeze=False)\n","    for ax,(col,label) in zip(axs[:,0],specs):\n","        for arm,g in frame.groupby(\"arm\"):\n","            a=g.groupby(\"effective_epoch\")[col].agg([\"mean\",\"std\"]).reset_index(); sd=a[\"std\"].fillna(0)\n","            line,=ax.plot(a.effective_epoch,a[\"mean\"],label=arm); ax.fill_between(a.effective_epoch,a[\"mean\"]-sd,a[\"mean\"]+sd,alpha=.18,color=line.get_color())\n","            for _,q in g.groupby(\"seed\"): ax.plot(q.effective_epoch,q[col],alpha=.22,linewidth=.7,color=line.get_color())\n","        ax.set(xlabel=\"effective epoch\",ylabel=label); ax.grid(alpha=.25); ax.legend()\n","    fig.suptitle(title); fig.tight_layout(); fig.savefig(ANALYSIS/f\"{name}.png\",dpi=150); plt.show()\n","\n","bands(metrics,[(\"train_loss\",\"train CE\"),(\"val_loss\",\"validation CE\")],\"Cross-entropy\",\"cross_entropy\")\n","bands(metrics,[(\"train_perplexity\",\"train perplexity\"),(\"val_perplexity\",\"validation perplexity\")],\"Perplexity\",\"perplexity\")\n","bands(metrics,[(\"train_accuracy\",\"train top-1\"),(\"val_accuracy\",\"validation top-1\"),(\"val_generalization_gap\",\"validation CE - train CE\")],\"Accuracy and generalization\",\"accuracy_gaps\")"]},{"cell_type":"code","execution_count":null,"id":"e0a9182b","metadata":{},"outputs":[],"source":["for p in (ROOT/\"level_0_baseline/src\",ROOT/\"level_0_wwpgd/src\"):\n","    if str(p) not in sys.path: sys.path.insert(0,str(p))\n","from level0_baseline.model import GPT,GPTConfig\n","\n","dtype=np.dtype(str(META.get(\"dtype\",\"uint16\"))); test=np.memmap(DATA/\"test.bin\",dtype=dtype,mode=\"r\")\n","def probe(seed,n_batches):\n","    gen=torch.Generator().manual_seed(seed+3001); out=[]; bs=int(b[\"training\"][\"batch_size\"]); block=int(b[\"model\"][\"block_size\"])\n","    for _ in range(n_batches):\n","        starts=torch.randint(len(test)-block-1,(bs,),generator=gen).tolist()\n","        x=torch.stack([torch.from_numpy(np.asarray(test[i:i+block],dtype=np.int64)) for i in starts])\n","        y=torch.stack([torch.from_numpy(np.asarray(test[i+1:i+1+block],dtype=np.int64)) for i in starts]); out.append((x,y))\n","    return out\n","@torch.inference_mode()\n","def evaluate(model,batches,device):\n","    model.eval(); n=ce=c1=c5=0\n","    for x,y in batches:\n","        x=x.to(device); y=y.to(device); logits,loss=model(x,y); m=y.numel(); n+=m; ce+=float(loss.cpu())*m\n","        c1+=int((logits.argmax(-1)==y).sum().cpu()); c5+=int(logits.topk(5,dim=-1).indices.eq(y[...,None]).any(-1).sum().cpu())\n","    ce/=n\n","    return {\"test_cross_entropy\":ce,\"test_perplexity\":math.exp(min(20,ce)),\"test_bits_per_token\":ce/math.log(2),\"test_top1_accuracy\":c1/n,\"test_top5_accuracy\":c5/n}\n","cache=ANALYSIS/\"checkpoint_test_metrics.csv\"\n","test_df=pd.read_csv(cache) if cache.exists() and os.getenv(\"NANOGPT_LEVEL0_REPEAT_FORCE_TEST\",\"0\")!=\"1\" else pd.DataFrame()\n","done={(str(r.arm),int(r.seed),int(r.step)) for r in test_df.itertuples()} if len(test_df) else set()\n","new=[]; dev=torch.device(DEVICE); cfg=GPTConfig(**b[\"model\"]); batches_n=int(os.getenv(\"NANOGPT_LEVEL0_REPEAT_TEST_EVAL_BATCHES\",b[\"training\"][\"eval_batches\"]))\n","tokens_step=int(b[\"training\"][\"batch_size\"])*int(b[\"model\"][\"block_size\"])*int(b[\"training\"][\"grad_accum_steps\"])\n","for arm,seed,r in runs:\n","    batches=probe(seed,batches_n); paths=sorted(r.glob(\"checkpoint_[0-9]*.pt\"))+[r/\"checkpoint_final.pt\"]; by_step={}\n","    for p in paths:\n","        if p.is_file():\n","            q=torch.load(p,map_location=\"cpu\",weights_only=False); by_step[int(q[\"step\"])]=p\n","    for step,p in sorted(by_step.items()):\n","        if (arm,seed,step) in done: continue\n","        q=torch.load(p,map_location=\"cpu\",weights_only=False); model=GPT(cfg).to(dev); model.load_state_dict(q[\"model\"])\n","        row={\"arm\":arm,\"seed\":seed,\"step\":step,\"tokens_seen\":step*tokens_step,\"effective_epoch\":step*tokens_step/TRAIN_TOKENS,**evaluate(model,batches,dev)}\n","        new.append(row); del model\n","        pd.concat([test_df,pd.DataFrame(new)],ignore_index=True).to_csv(cache,index=False)\n","test_df=pd.concat([test_df,pd.DataFrame(new)],ignore_index=True).drop_duplicates([\"arm\",\"seed\",\"step\"]).sort_values([\"arm\",\"seed\",\"step\"])\n","test_df.to_csv(cache,index=False)\n","bands(test_df,[(\"test_cross_entropy\",\"test CE\"),(\"test_perplexity\",\"test perplexity\"),(\"test_bits_per_token\",\"test bits/token\"),(\"test_top1_accuracy\",\"test top-1\"),(\"test_top5_accuracy\",\"test top-5\")],\"Post-hoc checkpoint test metrics\",\"checkpoint_test_metrics\")\n","gaps=test_df.merge(metrics[[\"arm\",\"seed\",\"step\",\"train_loss\",\"train_accuracy\",\"val_loss\",\"val_accuracy\",\"val_generalization_gap\"]],on=[\"arm\",\"seed\",\"step\"],how=\"left\")\n","gaps[\"test_ce_gap\"]=gaps.test_cross_entropy-gaps.train_loss; gaps[\"test_accuracy_gap\"]=gaps.train_accuracy-gaps.test_top1_accuracy\n","bands(gaps,[(\"val_generalization_gap\",\"validation CE - train CE\"),(\"test_ce_gap\",\"test CE - train CE\"),(\"test_accuracy_gap\",\"train top-1 - test top-1\")],\"Generalization gaps\",\"generalization_gaps\")"]},{"cell_type":"code","execution_count":null,"id":"aa6f254d","metadata":{},"outputs":[],"source":["parts=[]\n","for arm,seed,r in runs:\n","    for p in sorted(r.glob(\"weightwatcher_step_*.csv\")):\n","        d=pd.read_csv(p); d[\"arm\"]=arm; d[\"seed\"]=seed; parts.append(d)\n","ww=pd.concat(parts,ignore_index=True); ww[\"effective_epoch\"]=pd.to_numeric(ww.tokens_seen,errors=\"coerce\")/TRAIN_TOKENS\n","def col(*names):\n","    m={str(c).lower():str(c) for c in ww.columns}\n","    return next((m[n.lower()] for n in names if n.lower() in m),None)\n","aliases={\"alpha\":col(\"alpha\"),\"D\":col(\"D\",\"ks_distance\"),\"sigma\":col(\"sigma\",\"alpha_sigma\"),\"xmin\":col(\"xmin\"),\"tail_count\":col(\"num_pl_spikes\",\"num_evals_in_tail\",\"tail_size\"),\"num_evals\":col(\"num_evals\",\"M\"),\"detX_num\":col(\"detX_num\",\"num_ERG_spikes\"),\"ERG_gap\":col(\"ERG_gap\")}\n","for dst,src in aliases.items():\n","    ww[dst]=pd.to_numeric(ww[src],errors=\"coerce\") if src else np.nan\n","if aliases[\"ERG_gap\"] is None and aliases[\"detX_num\"] and aliases[\"tail_count\"]: ww[\"ERG_gap\"]=ww.detX_num-ww.tail_count\n","ww[\"tail_fraction\"]=ww.tail_count/ww.num_evals.replace(0,np.nan)\n","display(pd.DataFrame(aliases.items(),columns=[\"metric\",\"source\"]))\n","\n","def wwplot(metric,label,title,name,ref=None):\n","    if not ww[metric].notna().any(): print(\"skip\",metric); return\n","    types=sorted(ww.matrix_type.dropna().unique()); fig,axs=plt.subplots(math.ceil(len(types)/2),2,figsize=(14,4*math.ceil(len(types)/2)),squeeze=False)\n","    for ax,t in zip(axs.flat,types):\n","        z=ww[ww.matrix_type==t]\n","        for (arm,block),g in z.groupby([\"arm\",\"block\"]):\n","            u=g.groupby([\"seed\",\"effective_epoch\"],as_index=False)[metric].mean(); a=u.groupby(\"effective_epoch\")[metric].agg([\"mean\",\"std\"]).reset_index(); sd=a[\"std\"].fillna(0)\n","            line,=ax.plot(a.effective_epoch,a[\"mean\"],label=f\"{arm} B{int(block)}\"); ax.fill_between(a.effective_epoch,a[\"mean\"]-sd,a[\"mean\"]+sd,alpha=.12,color=line.get_color())\n","        if ref is not None: ax.axhline(ref,ls=\"--\",lw=1)\n","        ax.set(title=t,xlabel=\"effective epoch\",ylabel=label); ax.grid(alpha=.2)\n","    for ax in axs.flat[len(types):]: ax.axis(\"off\")\n","    h,l=axs.flat[0].get_legend_handles_labels()\n","    if h: fig.legend(h,l,bbox_to_anchor=(1,0.5),loc=\"center left\",fontsize=8)\n","    fig.suptitle(title); fig.tight_layout(); fig.savefig(ANALYSIS/f\"{name}.png\",dpi=150,bbox_inches=\"tight\"); plt.show()\n","wwplot(\"alpha\",\"alpha\",\"Layer alpha\",\"alpha_by_layer\",2)\n","wwplot(\"D\",\"KS D (lower better)\",\"Alpha-fit KS quality\",\"fit_D_by_layer\")\n","wwplot(\"sigma\",\"alpha sigma (lower better)\",\"Alpha-fit uncertainty\",\"fit_sigma_by_layer\")\n","wwplot(\"tail_fraction\",\"tail fraction\",\"Alpha-fit support\",\"tail_fraction_by_layer\")\n","wwplot(\"ERG_gap\",\"ERG gap\",\"ERG gap\",\"ERG_gap_by_layer\",0)\n","wwplot(\"detX_num\",\"detX retained count\",\"Trace-log retained count\",\"detX_by_layer\")\n","wwplot(\"tail_count\",\"power-law tail count\",\"Power-law tail size\",\"tail_count_by_layer\")"]},{"cell_type":"code","execution_count":null,"id":"bd7e17d0","metadata":{},"outputs":[],"source":["ps=[]\n","for arm,seed,r in runs:\n","    p=r/\"wwpgd_projection.csv\"\n","    if p.exists():\n","        d=pd.read_csv(p); d[\"seed\"]=seed; ps.append(d)\n","proj=pd.concat(ps,ignore_index=True)\n","for c in (\"optimizer_step\",\"tokens_seen\",\"relative_frobenius_change_requested\",\"relative_frobenius_change_applied\",\"projection_runtime_seconds\"): proj[c]=pd.to_numeric(proj[c],errors=\"coerce\")\n","proj[\"effective_epoch\"]=proj.tokens_seen/TRAIN_TOKENS\n","dose=proj.groupby([\"seed\",\"optimizer_step\"],as_index=False).agg(effective_epoch=(\"effective_epoch\",\"max\"),requested=(\"relative_frobenius_change_requested\",\"mean\"),applied=(\"relative_frobenius_change_applied\",\"mean\"),applied_max=(\"relative_frobenius_change_applied\",\"max\"),runtime=(\"projection_runtime_seconds\",\"sum\")).sort_values([\"seed\",\"optimizer_step\"])\n","dose[\"cumulative_applied\"]=dose.groupby(\"seed\").applied.cumsum()\n","fig,axs=plt.subplots(4,1,figsize=(11,14))\n","for ax,(m,label) in zip(axs,[(\"requested\",\"requested relative change\"),(\"applied\",\"applied relative change\"),(\"applied_max\",\"maximum applied change\"),(\"cumulative_applied\",\"cumulative applied dose\")]):\n","    a=dose.groupby(\"effective_epoch\")[m].agg([\"mean\",\"std\"]).reset_index(); sd=a[\"std\"].fillna(0); ax.plot(a.effective_epoch,a[\"mean\"]); ax.fill_between(a.effective_epoch,a[\"mean\"]-sd,a[\"mean\"]+sd,alpha=.2); ax.set(xlabel=\"effective epoch\",ylabel=label); ax.grid(alpha=.25)\n","fig.suptitle(\"WWPGD intervention size\"); fig.tight_layout(); fig.savefig(ANALYSIS/\"wwpgd_dose.png\",dpi=150); plt.show()\n","\n","final=test_df.sort_values(\"step\").groupby([\"arm\",\"seed\"],as_index=False).tail(1)\n","display(final.groupby(\"arm\")[[\"test_cross_entropy\",\"test_perplexity\",\"test_top1_accuracy\",\"test_top5_accuracy\"]].agg([\"mean\",\"std\"]))\n","for m in (\"test_cross_entropy\",\"test_perplexity\",\"test_top1_accuracy\",\"test_top5_accuracy\"):\n","    p=final.pivot(index=\"seed\",columns=\"arm\",values=m); p[\"WWPGD_minus_AdamW\"]=p[\"AdamW + WWPGD\"]-p[\"AdamW\"]; print(m); display(p)\n","\n","for name,df in {\"training_metrics.csv\":metrics,\"checkpoint_test_metrics.csv\":test_df,\"generalization_gaps.csv\":gaps,\"weightwatcher_metrics.csv\":ww,\"wwpgd_projection.csv\":proj,\"wwpgd_dose.csv\":dose}.items(): df.to_csv(ANALYSIS/name,index=False)\n","print(\"Analysis written to\",ANALYSIS)"]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10"}},"nbformat":4,"nbformat_minor":5}