From 6f4c0ab54945a7ef06d550c90716f817933e71b2 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Thu, 17 Sep 2026 15:37:22 -0700 Subject: [PATCH 1/4] Add controlled Muon 100k continuation from final checkpoint --- .../continue_muon_100k.py | 99 +++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py diff --git a/baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py b/baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py new file mode 100644 index 0000000..2c4101d --- /dev/null +++ b/baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Continue one completed Muon run from step 10k to 100k without replaying canaries.""" +from __future__ import annotations +import argparse, copy, json, math, time +from collections import Counter +from pathlib import Path +import torch + +from am_data import Dataset +from am_metrics import batches, suffix_losses +from am_runtime import baseline, imports, make_model, environment, model_hash, atomic_json, atomic_save, safe_clip +from am_train import audit, snapshot, restore + + +def main(): + p=argparse.ArgumentParser() + p.add_argument('--root',required=True,help='Completed 10-run study root') + p.add_argument('--seed',type=int,default=1337) + p.add_argument('--target-step',type=int,default=100000) + p.add_argument('--device',choices=['mps','cpu','cuda'],default='mps') + p.add_argument('--checkpoint-every',type=int,default=5000) + p.add_argument('--behavior-every',type=int,default=5000) + p.add_argument('--resume',action='store_true') + a=p.parse_args() + + root=Path(a.root).resolve(); original_cfg=json.loads((root/'protocol.json').read_text()) + start_run=root/'muon'/f'seed_{a.seed}'; source_checkpoint=start_run/'checkpoint_latest.pt' + if not source_checkpoint.exists(): raise FileNotFoundError(source_checkpoint) + source_saved=torch.load(source_checkpoint,map_location='cpu',weights_only=True) + start_step=int(source_saved['step']) + if start_step!=int(original_cfg['steps']): raise ValueError(f'Expected completed checkpoint at {original_cfg["steps"]}, found {start_step}') + if a.target_step<=start_step: raise ValueError('target-step must exceed the source checkpoint step') + + # Freeze the original 0..5000 canary acquisition schedule while extending only ordinary training. + cfg=copy.deepcopy(original_cfg) + cfg['steps']=a.target_step + cfg['withdrawal_step']=int(original_cfg['steps'])//2 + cfg['checkpoint_every']=a.checkpoint_every + cfg['behavior_every']=a.behavior_every + cfg['expensive_every']=max(a.behavior_every,25000) + + out=root/'extensions'/f'muon_seed_{a.seed}_to_{a.target_step}' + out.mkdir(parents=True,exist_ok=True) + atomic_json(out/'continuation_protocol.json',cfg) + atomic_json(out/'source.json',{'source_run':str(start_run),'source_checkpoint':str(source_checkpoint),'source_step':start_step}) + + runtime=environment(a.device); source=baseline(original_cfg) + _,_,make_handles,set_lrs,zero_grad,optimizer_step=imports() + model=make_model(source,original_cfg,a.seed,a.device) + profile=copy.deepcopy(source['optimizer_profiles']['muon']); handles=make_handles(model,profile) + data=Dataset(cfg,a.seed,source['training']['batch_size']*source['training']['grad_accum_steps']) + + original_manifest=json.loads((start_run/'manifest.json').read_text()) + if data.fingerprint!=original_manifest['data_sha256']: + raise ValueError('Extended dataset changed the original canary/data schedule; refusing continuation.') + + latest=out/'checkpoint_latest.pt'; completed=start_step; attempted=int(source_saved['attempted']); counts=Counter(source_saved['counts']) + if a.resume: + if not latest.exists(): raise ValueError('No continuation checkpoint exists to resume.') + saved=torch.load(latest,map_location='cpu',weights_only=True); completed=int(saved['step']); attempted=int(saved['attempted']); counts=Counter(saved['counts']); restore(saved,model,handles) + else: + if latest.exists(): raise ValueError('Continuation already exists; use --resume.') + restore(source_saved,model,handles) + if model_hash(model)!=source_saved['model_sha256']: raise ValueError('Source checkpoint tensor hash mismatch.') + state=snapshot(model,handles,counts,completed,attempted); state.update(model_sha256=model_hash(model),source_step=start_step) + atomic_save(latest,state); atomic_save(out/f'model_{completed:08d}.pt',{'model':state['model'],'step':completed,'source_step':start_step}) + + # Preserve the original Muon LR schedule. At 10k it has reached its configured minimum, + # so the continuation remains at that minimum rather than stretching/restarting the schedule. + schedule_steps=math.ceil(source['dataset']['train_tokens']*profile['lr_schedule_epochs']/(data.batch_size*source['model']['block_size'])) + warmup=math.ceil(schedule_steps*profile['warmup_fraction']) + print(f'Continuing Muon seed={a.seed}: {completed} -> {a.target_step} on {a.device}',flush=True) + print(f'Canary withdrawal remains frozen at step {data.withdrawal}; no new canary presentations occur.',flush=True) + print(f'Original LR schedule length={schedule_steps}; continuation uses the configured minimum LR after schedule end.',flush=True) + print(f'Output: {out}',flush=True) + started=time.monotonic() + + while completed Date: Thu, 17 Sep 2026 15:37:38 -0700 Subject: [PATCH 2/4] Allow frozen withdrawal horizon for controlled continuations --- baseline/experiments/nanogpt_alpha_memorization/am_data.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/baseline/experiments/nanogpt_alpha_memorization/am_data.py b/baseline/experiments/nanogpt_alpha_memorization/am_data.py index c893e95..2a7bbff 100644 --- a/baseline/experiments/nanogpt_alpha_memorization/am_data.py +++ b/baseline/experiments/nanogpt_alpha_memorization/am_data.py @@ -24,9 +24,10 @@ class Record: class Dataset: def __init__(self, cfg, seed, batch_size): self.cfg, self.seed, self.batch_size = cfg, seed, batch_size - self.steps, self.withdrawal = int(cfg['steps']), int(cfg['steps']) // 2 - if self.withdrawal < 1: - raise ValueError('At least two updates are needed.') + self.steps = int(cfg['steps']) + self.withdrawal = int(cfg.get('withdrawal_step', self.steps // 2)) + if self.withdrawal < 1 or self.withdrawal > self.steps: + raise ValueError('withdrawal_step must be within the training horizon.') rng = np.random.default_rng(cfg['data_seed']) self.train, self.audit, self.canaries = [], [], [] p = cfg['modulus'] From 86c96916c1cd771c3f218d8759377de2073c93dd Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Thu, 17 Sep 2026 15:38:04 -0700 Subject: [PATCH 3/4] Fix continuation presentation-count validation --- .../nanogpt_alpha_memorization/continue_muon_100k.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py b/baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py index 2c4101d..9d330b5 100644 --- a/baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py +++ b/baseline/experiments/nanogpt_alpha_memorization/continue_muon_100k.py @@ -87,7 +87,9 @@ def main(): counts.update(r.id for r in records); completed=step+1 if completed%100==0: print(f'muon seed={a.seed} step={completed}/{a.target_step} loss={last_loss:.6f} grad_norm={norm:.3g} norm_fallback={overflow}',flush=True) if completed%a.checkpoint_every==0 or completed==a.target_step: - if counts!=data.planned_counts(completed): raise RuntimeError('Presentation counts changed during continuation.') + planned=data.planned_counts(completed) + for r in data.canaries: + if counts[r.id]!=planned[r.id]: raise RuntimeError(f'Canary presentation count changed for {r.id}.') state=snapshot(model,handles,counts,completed,attempted); state.update(model_sha256=model_hash(model),source_step=start_step) atomic_save(latest,state); atomic_save(out/f'model_{completed:08d}.pt',{'model':state['model'],'step':completed,'source_step':start_step}) if completed%a.behavior_every==0 or completed==a.target_step: From 3b2e883f1751446a0785ef930af13d74cdda9faf Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Thu, 17 Sep 2026 15:38:13 -0700 Subject: [PATCH 4/4] Test frozen canary schedule for long Muon continuation --- .../tests/test_continuation.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 baseline/experiments/nanogpt_alpha_memorization/tests/test_continuation.py diff --git a/baseline/experiments/nanogpt_alpha_memorization/tests/test_continuation.py b/baseline/experiments/nanogpt_alpha_memorization/tests/test_continuation.py new file mode 100644 index 0000000..0c6e7ee --- /dev/null +++ b/baseline/experiments/nanogpt_alpha_memorization/tests/test_continuation.py @@ -0,0 +1,33 @@ +from __future__ import annotations +import copy, json +from pathlib import Path +import sys + +HERE=Path(__file__).resolve().parents[1] +sys.path.insert(0,str(HERE)) +from am_data import Dataset + +CFG=json.loads((HERE/'protocol.json').read_text()) + + +def test_extended_horizon_preserves_original_canary_schedule_and_fingerprint(): + batch_size=32 + original=Dataset(CFG,1337,batch_size) + extended_cfg=copy.deepcopy(CFG) + extended_cfg['steps']=100000 + extended_cfg['withdrawal_step']=CFG['steps']//2 + extended=Dataset(extended_cfg,1337,batch_size) + assert original.withdrawal==extended.withdrawal==5000 + assert original.fingerprint==extended.fingerprint + assert {k:v.id for k,v in original.schedule.items()}=={k:v.id for k,v in extended.schedule.items()} + assert extended.planned_counts(100000)==original.planned_counts(10000) + + +def test_extended_batches_after_10k_contain_no_canaries(): + extended_cfg=copy.deepcopy(CFG) + extended_cfg['steps']=100000 + extended_cfg['withdrawal_step']=5000 + data=Dataset(extended_cfg,1337,32) + canary_ids={r.id for r in data.canaries} + for step in (10000,25000,50000,99999): + assert not any(r.id in canary_ids for r in data.batch(step))