1 files changed, 9 insertions, 1 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py
index 7647fae..069d858 100644
--- a/ep_run/casc_bp_train.py
+++ b/ep_run/casc_bp_train.py
@@ -21,6 +21,7 @@ ap.add_argument('--cosine', action='store_true')         # warmup then cosine de
 ap.add_argument('--lr_min_ratio', type=float, default=0.1)
 ap.add_argument('--qk_norm', action='store_true')        # RMS-norm q,k per head before scores (OLMo2-style; bounds logits, analog-friendly)
 ap.add_argument('--final_ln', action='store_true')       # final LayerNorm before readout (standard GPT; bounds sig_tok growth -> keeps beta/estimator healthy on long runs)
+ap.add_argument('--resume', default='')                  # path to a ckpt (tok/pos/blocks) to continue from; step taken from ckpt
 args = ap.parse_args()
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -76,6 +77,12 @@ blocks = nn.ModuleList([Block(args.C, args.H, args.qk_norm) for _ in range(args.
 mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1)
 ln_f = nn.LayerNorm(args.C).to(dev) if args.final_ln else nn.Identity()
 params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + list(ln_f.parameters())
+start_step = 0
+if args.resume:
+    _ck = torch.load(args.resume, map_location=dev, weights_only=False)
+    tok.load_state_dict(_ck['tok']); pos.load_state_dict(_ck['pos']); blocks.load_state_dict(_ck['blocks'])
+    start_step = int(_ck.get('step', 0))
+    print(f'[resume] loaded {args.resume} at step {start_step}', flush=True)
 if args.opt == 'muon':
     from muon import build_hybrid
     opt, sched = build_hybrid(blocks, params, args.lr, args.muon_lr, args.warmup)
@@ -116,7 +123,8 @@ n = sum(p.numel() for p in params)
 print(f'[{args.tag}] cascade-BP L{args.L} C{args.C} H{args.H} T{args.T} | {n/1e6:.2f}M params | {dev}', flush=True)
 best, t0 = 1e9, time.time()
 outdir = Path('runs'); outdir.mkdir(exist_ok=True)
-for step in range(args.steps + 1):
+for _ in range(start_step): sched.step()   # advance LR schedule to the resumed step
+for step in range(start_step, args.steps + 1):
     x, y = get_batch('train')
     loss = F.cross_entropy(fwd(x).reshape(-1, vocab), y.reshape(-1))
     opt.zero_grad(set_to_none=True); loss.backward()