Run the sweep yourself

Which GPU

Colab assigns the GPU, and may give a different one from the one you asked for, so check: the notebook's test cell prints device, and every run log records it in its first line ("hardware": {… "device": …}). The experiment's 33 runs used:

• NVIDIA A100-SXM4-40GB: 27 runs
• NVIDIA L4: 3 runs
• NVIDIA A100-SXM4-80GB: 2 runs
• Tesla T4: 1 run

Drive copies and resuming

The work happens on Colab's local disk; results, checkpoints and the tokenised data are copied to MyDrive/DeMo_Fixed_Band_Experiment_state with rsync every 10 minutes and when the sweep ends, and restored from there at the start of each session. If the session disconnects, reconnect and run all again: finished runs are skipped and an interrupted run resumes from its last checkpoint. To stop early, interrupt the sweep cell and run save_to_drive() before disconnecting.

How long each configuration took

Wall time per run of 3,000 steps, mean over seeds, from the run logs (results.json's seconds), with the GPUs the runs used. Times vary with the GPU.

1. DeMo fixed-band experiment on a Colab GPU

   Runs the experiment in DeMo_Fixed_Band_Experiment on a GPU, at the tinystories-gpu size (about 16M parameters, 3,000 steps per run).

   1. Runtime > Change runtime type > GPU (T4 is free; an A100 is faster).
   2. Runtime > Run all. The first cells check the GPU, run the tests and a smoke run, and time the model; the sweep cell then runs the whole grid.
   3. If the session disconnects, reconnect and *Run all* again. Each session pulls the latest code from GitHub and restores the finished runs from Drive; the timing cell is skipped once the sweep has started.
   4. To stop early, interrupt the sweep cell and run save_to_drive() in a new cell before disconnecting, so the last few minutes are kept.

   The work happens on Colab's local disk (training directly on the Drive mount made the mount drop). Results, checkpoints and the tokenised data are copied to MyDrive/DeMo_Fixed_Band_Experiment_state every 10 minutes and when the sweep ends, and restored from there at the start of each session, so finished runs are skipped and an interrupted run resumes from its last checkpoint.

   Optional: to push results to GitHub as each run finishes, add a Colab secret (key icon, left) named GH_TOKEN holding a fine-grained token with write access to this repository's contents.

2. In [1]

   !nvidia-smi --query-gpu=name,memory.total --format=csv

3. In [2]

   from google.colab import drive
   import os, subprocess, threading, time
   drive.mount('/content/drive')
   REPO = '/content/DeMo_Fixed_Band_Experiment'
   STORE = '/content/drive/MyDrive/DeMo_Fixed_Band_Experiment_state'
   if not os.path.exists(REPO):
       !git clone -q https://github.com/BrendanJamesLynskey/DeMo_Fixed_Band_Experiment.git {REPO}
   %cd {REPO}
   !git config user.name "Colab sweep" && git config user.email "colab-sweep@users.noreply.github.com"
   # fetch and reset rather than pull: results committed to the repository would otherwise clash with
   # the same files restored from Drive, and the pull would fail, leaving the old code in place
   !git fetch -q origin && git reset -q --hard origin/main
   !git log --oneline -1
   !which rsync > /dev/null || apt-get -qq install -y rsync > /dev/null
   os.makedirs(STORE, exist_ok=True)
   # restore what earlier sessions saved: results, checkpoints and the tokenised data
   !rsync -a {STORE}/results/ results/ 2>/dev/null; rsync -a {STORE}/data/ data/ 2>/dev/null; true
   import glob
   print('restored from Drive:', len(glob.glob('results/tinystories-gpu/*.jsonl')), 'run logs,',
         len(glob.glob('results/ckpt/*.pt')), 'checkpoints')
   
   def save_to_drive():
       """Copy results, checkpoints and tokenised data to Drive (not the raw text downloads)."""
       if not os.path.isdir(STORE):
           print('Drive is not mounted; skipping the copy')
           return
       # a finished run deletes its checkpoint locally, but rsync does not propagate deletions:
       # remove checkpoints of finished runs from both copies, or they fill the Drive quota
       import glob, json as _json
       for ck in glob.glob('results/ckpt/*.pt') + glob.glob(f'{STORE}/results/ckpt/*.pt'):
           stem = os.path.basename(ck)[:-3]
           for log in glob.glob(f'results/*/{stem}.jsonl'):
               lines = open(log).read().strip().splitlines()
               if lines and _json.loads(lines[-1]).get('type') == 'done':
                   os.remove(ck)
                   break
       for src in ('results', 'data'):
           if os.path.isdir(src):
               subprocess.run(['rsync', '-a', '--exclude', '*.txt', '--exclude', '*.part', src + '/', f'{STORE}/{src}/'])
   
   def _saver(period=600):
       while True:
           time.sleep(period)
           try:
               save_to_drive()
           except Exception as e:
               print('Drive copy failed:', e)
   
   if not globals().get('_saver_started'):
       threading.Thread(target=_saver, daemon=True).start()
       _saver_started = True

4. In [3]

   # PyTorch with CUDA is preinstalled on Colab; add the rest
   !pip -q install tokenizers ruff pytest einops

5. In [4]

   # Tests (they run on the CPU) and a short run of every variant on synthetic data, on the GPU
   !python -m pytest -q tests
   !python experiment/sweep.py --grid ci --dataset synthetic
   !python -c "import json,glob; f=sorted(glob.glob('results/synthetic/*.jsonl'))[0]; print('device:', json.loads(open(f).readline())['hardware']['device'])"

6. In [5]

   # Timing: 30 steps of AdamW and of DeMo with a fixed band, then an estimate for the whole grid.
   # Skipped once the sweep has started (set FORCE_TIMING = True to run it anyway).
   FORCE_TIMING = False
   if FORCE_TIMING or not os.path.isdir('results/tinystories-gpu'):
       import json, sys, time, dataclasses
       sys.path.insert(0, 'experiment')
       import sweep, train
       from pathlib import Path
       prof = sweep.SIZES['tinystories-gpu']
       rate = {}
       for v in ('adamw', 'band', 'optics'):
           cfg = train.RunConfig(**{**prof, 'variant': v, 'steps': 30, 'eval_every': 30, 'diag_every': 30, 'ckpt_every': 1000,
                                    'chunk_mode': '1d' if v == 'optics' else '2d', 'chunk': 256 if v == 'optics' else 64})
           out = Path('/content/timing') / f'{v}.jsonl'
           if out.exists(): out.unlink()
           train.run(cfg, out)
           ev = [json.loads(l) for l in out.read_text().splitlines() if '"eval"' in l][-1]
           rate[v] = ev['seconds'] / 30
           print(f'{v}: {rate[v]:.2f} s/step')
       total = 0
       for stage in ('lr', 'core', 'optics', 'shapes'):
           for c in sweep.grid(stage, 'tinystories-gpu'):
               total += c.steps * rate['adamw' if c.variant == 'adamw' else 'optics' if c.variant == 'optics' else 'band']
       print(f'estimated time for the whole grid: {total / 3600:.1f} h')
   else:
       print('sweep already started: timing skipped')
   

7. In [6]

   # The whole experiment: learning-rate sweep, then the core, optics and band-shape stages.
   # Long-running. Finished runs are skipped on a rerun; an interrupted one resumes.
   FLAGS = ''
   try:
       from google.colab import userdata
       token = userdata.get('GH_TOKEN')
       !git remote set-url origin "https://x-access-token:{token}@github.com/BrendanJamesLynskey/DeMo_Fixed_Band_Experiment.git"
       FLAGS = '--commit --push'
       print('results will be committed and pushed after each run')
   except Exception:
       print('no GH_TOKEN secret: results stay in Drive')
   !python experiment/sweep.py --grid all --dataset tinystories-gpu $FLAGS 2>&1 | tee -a results/colab_sweep.log
   save_to_drive()

8. In [7]

   !python experiment/analyse.py --dataset tinystories-gpu
   from IPython.display import Image, display
   for f in ('loss_vs_steps_and_bytes.png', 'enob_tolerance.png'):
       p = f'results/tinystories-gpu/{f}'
       if os.path.exists(p): display(Image(p))