fast local checkpoints, fewer bucket writes

scugnizz-llama.py

@@ -45,6 +45,9 @@ KEEP (optimizer resume):
 SAFE TO DELETE (weights-only snapshots):
   - checkpoint_weights_last.pt
   - old model_final.pt after a newer run finishes
+
+Checkpoints write to CKPT_LOCAL_DIR (default /tmp) first, then resume is copied to out-dir.
+Set WEIGHTS_SAVE_INTERVAL=0 for long runs (default).
 """
 
 
@@ -276,9 +279,13 @@ def parse_args():
     p.add_argument("--dropout", type=float, default=0.0)
     p.add_argument("--eval-interval", type=int, default=50)
     p.add_argument("--eval-batches", type=int, default=10)
-    p.add_argument("--save-interval", type=int, default=50)
-    p.add_argument("--weights-save-interval", type=int, default=10,
-                   help="Save weights-only checkpoint every N steps (fast, ~3GB)")
+    p.add_argument("--save-interval", type=int, default=100)
+    p.add_argument("--weights-save-interval", type=int, default=0,
+                   help="Weights-only snapshot every N steps (0=off; slows bucket I/O)")
+    p.add_argument("--ckpt-local-dir", default="/tmp/scugnizz-ckpts",
+                   help="Fast local dir for checkpoint writes; synced to out-dir periodically")
+    p.add_argument("--bucket-sync-interval", type=int, default=1,
+                   help="Sync resume checkpoint to out-dir every N save-intervals (1=every save)")
     p.add_argument("--log-interval", type=int, default=10)
     p.add_argument("--pcs-a", type=float, default=0.8309193524478643)
     p.add_argument("--pcs-b", type=float, default=0.0)
@@ -439,11 +446,11 @@ def save_ckpt(path, model, opt, step, best, rank, world_size, mirror_resume=Fals
     barrier(rank, world_size)
 
 
-def copy_best_ckpt(out, rank, world_size):
+def copy_best_ckpt(local, rank, world_size):
     if not is_main(rank):
         barrier(rank, world_size)
         return
-    last, best = out / "checkpoint_last.pt", out / "checkpoint_best.pt"
+    last, best = local / "checkpoint_last.pt", local / "checkpoint_best.pt"
     if last.exists():
         shutil.copy2(last, best)
         print("SAVED", best, "(copy of checkpoint_last)", flush=True)
@@ -488,12 +495,12 @@ def load_init_payload(path, dev):
     return ck, None, False
 
 
-def resume_training(model, out, a, rank, world_size, dev):
+def resume_training(model, local, out, a, rank, world_size, dev):
     start, best, opt_state = 0, float("inf"), None
     if a.no_resume:
         return start, best, opt_state
 
-    resume_path = find_resume_ckpt(out)
+    resume_path = find_resume_ckpt(local, out)
     init_path = a.init_from or fetch_init_hub(a, rank, world_size)
     weights_path = out / "model_final.pt"
 
@@ -565,11 +572,40 @@ def resume_training(model, out, a, rank, world_size, dev):
     return start, best, opt_state
 
 
-def find_resume_ckpt(out: Path):
-    for name in ("checkpoint_last.pt", RESUME_CKPT, "checkpoint_best.pt"):
-        p = out / name
-        if p.exists():
-            return p
+def sync_file(src, dst, rank, world_size):
+    if not is_main(rank):
+        barrier(rank, world_size)
+        return
+    if src.exists():
+        dst.parent.mkdir(parents=True, exist_ok=True)
+        tmp = dst.with_suffix(dst.suffix + ".tmp")
+        shutil.copy2(src, tmp)
+        tmp.replace(dst)
+        print("SYNCED", dst, flush=True)
+    barrier(rank, world_size)
+
+
+def sync_bucket_ckpts(local, out, rank, world_size, names):
+    if not is_main(rank):
+        barrier(rank, world_size)
+        return
+    for name in names:
+        src, dst = local / name, out / name
+        if src.exists():
+            dst.parent.mkdir(parents=True, exist_ok=True)
+            tmp = dst.with_suffix(dst.suffix + ".tmp")
+            shutil.copy2(src, tmp)
+            tmp.replace(dst)
+            print("SYNCED", dst, flush=True)
+    barrier(rank, world_size)
+
+
+def find_resume_ckpt(local, out):
+    for base in (local, out):
+        for name in ("checkpoint_last.pt", RESUME_CKPT, "checkpoint_best.pt"):
+            p = base / name
+            if p.exists():
+                return p
     return None
 
 
@@ -651,8 +687,10 @@ def main():
     set_seed(a.seed, rank)
     dt = dtype_of(a.dtype, dev)
     out = Path(a.out_dir)
+    local = Path(a.ckpt_local_dir)
     if is_main(rank):
         out.mkdir(parents=True, exist_ok=True)
+        local.mkdir(parents=True, exist_ok=True)
         (out / "CHECKPOINT_POLICY.txt").write_text(KEEP_NOTE, encoding="utf-8")
     barrier(rank, world_size)
 
@@ -674,7 +712,7 @@ def main():
     params = sum(p.numel() for p in model.parameters())
     log0(rank, f"params {params/1e9:.3f}B ({params/1e6:.1f}M)", flush=True)
 
-    start, best, opt_state = resume_training(model, out, a, rank, world_size, dev)
+    start, best, opt_state = resume_training(model, local, out, a, rank, world_size, dev)
     if is_main(rank):
         (out / "args.json").write_text(json.dumps(vars(a), indent=2), encoding="utf-8")
     barrier(rank, world_size)
@@ -694,17 +732,20 @@ def main():
         opt.load_state_dict(opt_state)
         log0(rank, "loaded optimizer state", flush=True)
 
-    last = out / "checkpoint_last.pt"
-    bestp = out / "checkpoint_best.pt"
-    weights_last = out / "checkpoint_weights_last.pt"
+    last = local / "checkpoint_last.pt"
+    bestp = local / "checkpoint_best.pt"
+    weights_last = local / "checkpoint_weights_last.pt"
+    saves_done = 0
 
     log0(rank, "initializing data streams...", flush=True)
+    log0(rank, f"ckpt_local={local} out_dir={out}", flush=True)
     train = Batcher(a, tok, dev, rank, log=is_main(rank))
     val = Batcher(a, tok, dev, 10_000, log=False) if is_main(rank) else None
     log0(rank, f"starting training loop from step {start}...", flush=True)
     t0 = time.time()
     roll = 0.0
     n = 0
+    log_t0 = t0
     model.train()
 
     for step in range(start, a.max_steps + 1):
@@ -727,13 +768,17 @@ def main():
 
         if step % a.log_interval == 0:
             avg = roll / max(1, n)
-            elapsed = max(1e-6, time.time() - t0)
+            now = time.time()
+            elapsed = max(1e-6, now - t0)
+            interval = max(1e-6, now - log_t0)
             tok_s = max(1, step - start + 1) * tps / elapsed
+            instant = a.log_interval * tps / interval
             log0(
                 rank,
-                f"step {step:08d} | loss {avg:.4f} | ppl {math.exp(min(20, avg)):.1f} | lr {lr:.2e} | tok/s {tok_s:.0f} | trained_tokens {step*tps} | rows {train.rows}",
+                f"step {step:08d} | loss {avg:.4f} | ppl {math.exp(min(20, avg)):.1f} | lr {lr:.2e} | tok/s {tok_s:.0f} ({instant:.0f} inst) | trained_tokens {step*tps} | rows {train.rows}",
                 flush=True,
             )
+            log_t0 = now
             roll = 0.0
             n = 0
 
@@ -761,13 +806,22 @@ def main():
         if step > 0 and step % a.save_interval == 0:
             save_ckpt(last, model, opt, step, best, rank, world_size, mirror_resume=True)
             if val_improved:
-                copy_best_ckpt(out, rank, world_size)
+                copy_best_ckpt(local, rank, world_size)
+            saves_done += 1
+            if saves_done % max(1, a.bucket_sync_interval) == 0:
+                names = [RESUME_CKPT, "checkpoint_last.pt"]
+                if (local / "checkpoint_best.pt").exists():
+                    names.append("checkpoint_best.pt")
+                sync_bucket_ckpts(local, out, rank, world_size, names)
             if is_main(rank) and a.push_every_save:
                 upload_final(str(out), a.hub_repo_id, a.hub_path, f"pretrain checkpoint step {step}")
 
     save_ckpt(last, model, opt, a.max_steps, best, rank, world_size, mirror_resume=True)
+    sync_bucket_ckpts(local, out, rank, world_size, (RESUME_CKPT, "checkpoint_last.pt", "checkpoint_best.pt"))
     if is_main(rank):
-        torch.save(raw_model(model).state_dict(), out / "model_final.pt")
+        final_local = local / "model_final.pt"
+        torch.save(raw_model(model).state_dict(), final_local)
+        sync_file(final_local, out / "model_final.pt", rank, world_size)
         tok.save_pretrained(out / "tokenizer")
         if a.push:
             upload_final(str(out), a.hub_repo_id, a.hub_path, "final ScugnizzLLM Llama-PCS pretrain")