# 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.
```