1 files changed, 117 insertions, 0 deletions
diff --git a/docs/COLAB_50M.md b/docs/COLAB_50M.md
new file mode 100644
index 0000000..9236049
--- /dev/null
+++ b/docs/COLAB_50M.md
@@ -0,0 +1,117 @@
+# Colab 50M EP run — self-contained notebook cells
+
+Goal: train the 50M (C=2048) EP energy-transformer on TinyStories-BPE on a Colab A100/H100,
+which fixes the two things timan1's A6000 can't: fp32 throughput and the big-width init instability
+(more headroom to tune the curriculum + bigger batch). Checkpoints to Google Drive with full-state
+resume, so the 12 h Colab timeout never loses progress — just re-run the training cell to continue.
+
+PREP (once, on your laptop): download from timan1 `~/ept/ept_colab.tar.gz` (16 KB, the code) and
+upload it to your Google Drive root as `ept_colab.tar.gz`. Data is regenerated in-notebook (or
+upload `~/ept/lt_ep_code/.../tinystories_bpe` bins to Drive to skip the ~40 min prep — optional).
+
+────────────────────────────────────────────────────────────────────────
+## Cell 1 — setup, Drive, GPU, deps
+```python
+import torch, subprocess, os
+print(torch.__version__, torch.cuda.get_device_name(0))
+assert torch.__version__ >= "2.1", "need torch>=2.1 for torch.func/compile"
+print(subprocess.run(["nvidia-smi","--query-gpu=name,memory.total","--format=csv,noheader"],
+                     capture_output=True,text=True).stdout)
+from google.colab import drive; drive.mount('/content/drive')
+!pip -q install tokenizers
+WORK="/content/work"; DRIVE="/content/drive/MyDrive"; os.makedirs(WORK, exist_ok=True)
+!tar xzf {DRIVE}/ept_colab.tar.gz -C {WORK}
+print("code:", os.listdir(WORK))
+```
+
+## Cell 2 — data (regenerate, cached to Drive; skip if bins already uploaded)
+```python
+import os
+DATA="/content/drive/MyDrive/ept_data/tinystories_bpe"
+if os.path.exists(f"{DATA}/train.bin"):
+    print("BPE bins found on Drive — reusing.")
+else:
+    os.makedirs("/content/drive/MyDrive/ept_data/tinystories", exist_ok=True)
+    %cd /content/drive/MyDrive/ept_data/tinystories
+    !test -f train.txt || wget -q -O train.txt https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt
+    !test -f valid.txt || wget -q -O valid.txt https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt
+    # point the prep script at the Drive paths
+    import re; src=open(f"{WORK}/prepare_tinystories_bpe.py").read()
+    src=src.replace("/tmp/lt_ep/data/tinystories","/content/drive/MyDrive/ept_data/tinystories")
+    src=src.replace("/tmp/lt_ep/data/tinystories_bpe",DATA)
+    open(f"{WORK}/prep_bpe_colab.py","w").write(src)
+    %cd {WORK}
+    !python prep_bpe_colab.py
+print("data:", os.listdir(DATA))
+```
+
+## Cell 3 — STABILITY SMOKE (always run first; ~15 min). Must survive past warmup without abort.
+```python
+%cd {WORK}
+# C=2048 starting curriculum (muP-scaled from C=1024's lr 4e-4 -> ~2e-4; longer warmup; gentler resinit)
+!python lt_ep_train.py --mode ep --attn_mode thick --B 16 --C 2048 --H 16 --T 512 \
+  --c 1.0 --jacreg 1.0 --jr_floor 0.1 --res_target 1.5e-3 --jr_max 64 --res_ema 0.9 \
+  --holo 2 --hr 0.02 --pema 0.999 --t1max 300 --res_est 1e-4 --t2sel 60 --res_gate 5e-3 \
+  --qknorm --resinit 0.05 --warmup 2500 --compile --T1 150 --T2 20 --lr 2e-4 \
+  --steps 1200 --log 100 --data {DATA}
+# READ THE OUTPUT: if it ABORTs or res spikes >0.1 repeatedly through steps 600-1200, the curriculum
+# is still too hot -> lower lr to 1e-4 and/or resinit 0.03 and/or warmup 4000, re-run this cell.
+# If res stays <1e-2 and val descends past step 1000, the curriculum is good -> go to Cell 4.
+```
+
+## Cell 3b — KEEP-ALIVE (run once, then it auto-clicks connect every 60s to beat the ~90min idle kill)
+Open the browser JS console (F12 → Console) on the Colab tab and paste:
+```javascript
+function keepAlive(){
+  document.querySelector("colab-connect-button")?.shadowRoot?.querySelector("#connect")?.click();
+}
+setInterval(keepAlive, 60000);
+```
+This beats ONLY the idle timeout. The HARD cap (free 12h / Pro 24h, and Pro+ background execution
+is unreliable in 2026) is unbeatable — which is why Cell 4 is built to RESUME. When Colab drops you,
+just reconnect and re-run Cell 4; it continues from the last `--save_every` checkpoint on Drive.
+
+## Cell 4 — FULL RUN with Drive full-state resume. Re-run this exact cell after EVERY disconnect.
+```python
+%cd {WORK}
+ST="/content/drive/MyDrive/ept_ckpt/s4_50m.state"; CK="/content/drive/MyDrive/ept_ckpt/s4_50m.best.pt"
+os.makedirs("/content/drive/MyDrive/ept_ckpt", exist_ok=True)
+# --resume loads ST (weights+optimizer+sched+step+jr+best) if present -> idempotent across timeouts.
+# --save_every 100 = atomic full-state save every 100 steps -> a kill loses at most ~100 steps.
+!python lt_ep_train.py --mode ep --attn_mode thick --B 16 --C 2048 --H 16 --T 512 \
+  --c 1.0 --jacreg 1.0 --jr_floor 0.1 --res_target 1.5e-3 --jr_max 64 --res_ema 0.9 \
+  --holo 2 --hr 0.02 --pema 0.999 --t1max 300 --res_est 1e-4 --t2sel 60 --res_gate 5e-3 \
+  --qknorm --resinit 0.05 --warmup 2500 --compile --T1 150 --T2 20 --lr 2e-4 \
+  --steps 24000 --log 200 --save_every 100 --data {DATA} --ckpt {CK} --state {ST} --resume
+# IMPORTANT: match every flag here to the curriculum that PASSED Cell 3 (esp. lr/warmup/resinit).
+# On the FIRST run ST won't exist (fresh start, prints init residual); every re-run prints "[resume] from ...".
+```
+
+### Checkpointing guarantees (tested on timan1)
+- `--state` writes the FULL state (weights + AdamW moments + LR-schedule position + step + λ + best)
+  to `ST.tmp` then `os.replace` → **atomic**: a kill mid-write leaves the previous good `ST` intact.
+- `--resume` continues the LR schedule and optimizer momentum exactly (not a cold warm-start):
+  verified step 150 → resumed 151 with val still descending monotonically.
+- State size at 50M ≈ ~1 GB (weights+pema+opt); `--save_every 100` ≈ a 1 GB Drive write every
+  ~20 min of A100 wall-clock (well under Drive's daily quota). Lower to 50 if you want ≤10-min loss.
+- `--ckpt` (CK) separately keeps the best-val weights for sampling (Cell 5), updated only on improvement.
+
+## Cell 5 — sample stories from the best checkpoint (run anytime; reads CK from Drive)
+```python
+%cd {WORK}
+CK="/content/drive/MyDrive/ept_ckpt/s4_50m.best.pt"
+!python sample_eq.py --ckpt {CK} --data {DATA} --C 2048 --H 16 --T 512 --use_pema --n 4 \
+  --prompt "Once upon a time" --temp 0.8 --topk 40
+```
+Note: sample_eq.py reads vocab from meta.pkl; for BPE it prints token ids unless decoded — if it
+shows numbers not text, ping me and I'll add the BPE decode (tokenizer.json is in {DATA}).
+
+────────────────────────────────────────────────────────────────────────
+NOTES
+- The curriculum in Cells 3/4 is a STARTING GUESS for C=2048 (we never got it stable on timan1).
+  Cell 3 is there precisely to dial it in fast on the better GPU before committing Cell 4's long run.
+- Full-state resume tested on timan1 (step 150 → resumed 151, optimizer/schedule intact).
+- Expected cost: A100 fp32 ~2-3x an A6000 → ~0.06-0.1 it/s → 24k steps ~3-4 days of wall-clock
+  ACROSS resumes (so leave it, re-run Cell 4 whenever Colab drops you). H100 faster.
+- sample_eq.py BPE-decode gap is the one known rough edge; tell me if Cell 5 prints ids.
+```