Restore original n.py (was overwritten by delta patch upload)

n.py

@@ -4,8 +4,7 @@
 # Enhanced inference: checkpoint name, tok/s, UK time
 
 from __future__ import annotations
-import argparse, json, math, pathlib, random, time, os, sys, threading, hashlib
-from pathlib import Path
+import argparse, json, math, pathlib, random, time, os, sys
 from contextlib import nullcontext
 from typing import Dict, Any, List, Optional, Tuple
 from datetime import datetime, timezone
@@ -136,8 +135,6 @@ SAT_BLOCK = 2
 LR_CORE, LR_HEAD = 5e-5, 2e-4
 EMIT_LAMBDA = 0.1
 DEFAULT_SAVE_SEC = 24 * 3600
-DEFAULT_DELTA_STEPS = 500      # lightweight weight-only save every N steps
-DEFAULT_MAX_DELTAS = 5         # keep last N deltas (older pruned after full save)
 CKDIR = pathlib.Path("ckpts_expansion")
 
 DEFAULT_PRETRAIN_SOURCES = "OpenTransformer/goddess-crawl,OpenTransformer/agillm-crawl-data,OpenTransformer/web-crawl-2026,OpenTransformer/web-crawl-clean-v2,OpenTransformer/scraped-web-data,OpenTransformer/turbo-crawl,OpenTransformer/sft-data-clean,OpenTransformer/web-crawl-v1"
@@ -486,99 +483,6 @@ def sat_mask_cached(new_len: int, cached_len: int, block=SAT_BLOCK):
 
 
 # ───────────────────────── Checkpoint helpers ─────────────────────────
-
-# ───────────────────────── Delta Checkpoints (weight-only, async) ─────────────────────────
-_delta_lock = threading.Lock()
-_delta_thread: Optional[threading.Thread] = None
-
-def _sha256_file(path: pathlib.Path) -> str:
-    """Compute SHA256 of a file for integrity verification."""
-    h = hashlib.sha256()
-    with open(path, "rb") as f:
-        for chunk in iter(lambda: f.read(1 << 20), b""):
-            h.update(chunk)
-    return h.hexdigest()
-
-def _do_delta_save(tensors: dict, path: pathlib.Path, meta: dict):
-    """Background worker: write weight-only checkpoint + checksum."""
-    try:
-        path.parent.mkdir(exist_ok=True, parents=True)
-        tmp = path.with_suffix(path.suffix + ".dtmp")
-        torch.save({"weights": tensors, **meta}, tmp, _use_new_zipfile_serialization=False)
-        digest = _sha256_file(tmp)
-        tmp.replace(path)
-        # Write sidecar checksum
-        path.with_suffix(".sha256").write_text(f"{digest}  {path.name}\n")
-        print(f"  [delta] saved {path.name} ({digest[:12]}...)")
-    except Exception as e:
-        print(f"  [delta] FAILED {path.name}: {e}")
-
-def save_delta(core, ar_h, sat_h, step: int, seen_tok: int, save_dir: pathlib.Path, phase_name: str):
-    """Save weight-only delta in background thread. Non-blocking."""
-    global _delta_thread
-    # Wait for any previous delta write to finish
-    if _delta_thread is not None and _delta_thread.is_alive():
-        _delta_thread.join(timeout=60)
-    # Snapshot weights to CPU (detach from GPU graph)
-    with _delta_lock:
-        tensors = {
-            "core": {k: v.detach().cpu() for k, v in core.state_dict().items()},
-            "ar":   {k: v.detach().cpu() for k, v in ar_h.state_dict().items()},
-            "sat":  {k: v.detach().cpu() for k, v in sat_h.state_dict().items()},
-        }
-    meta = {"step": step, "seen_tok": seen_tok, "wall_time": time.time(), "delta": True}
-    path = save_dir / f"{phase_name}_delta_step{step:08d}.pt"
-    _delta_thread = threading.Thread(target=_do_delta_save, args=(tensors, path, meta), daemon=True)
-    _delta_thread.start()
-
-def _prune_deltas(save_dir: pathlib.Path, phase_name: str, max_deltas: int):
-    """Keep only the most recent max_deltas delta files."""
-    if max_deltas is None or max_deltas <= 0:
-        return
-    try:
-        pattern = f"{phase_name}_delta_step*.pt"
-        deltas = sorted(
-            [p for p in save_dir.glob(pattern) if p.stat().st_size > 0],
-            key=lambda p: p.stat().st_mtime
-        )
-        excess = len(deltas) - max_deltas
-        if excess > 0:
-            for p in deltas[:excess]:
-                try:
-                    p.unlink()
-                    sha = p.with_suffix(".sha256")
-                    if sha.exists(): sha.unlink()
-                    print(f"  [delta-prune] deleted {p.name}")
-                except Exception:
-                    pass
-    except Exception as e:
-        print(f"  [delta-prune] error: {e}")
-
-def load_delta(path: pathlib.Path, core, ar_h, sat_h):
-    """Load weight-only delta. Returns (step, seen_tok) or raises."""
-    # Verify checksum if sidecar exists
-    sha_path = path.with_suffix(".sha256")
-    if sha_path.exists():
-        expected = sha_path.read_text().split()[0]
-        actual = _sha256_file(path)
-        if expected != actual:
-            raise ValueError(f"Checksum mismatch for {path.name}: expected {expected[:12]}... got {actual[:12]}...")
-        print(f"  [delta] checksum OK for {path.name}")
-    ck = torch.load(path, map_location="cpu", weights_only=False)
-    if not ck.get("delta"):
-        raise ValueError(f"{path.name} is not a delta checkpoint")
-    core.load_state_dict(ck["weights"]["core"])
-    ar_h.load_state_dict(ck["weights"]["ar"])
-    sat_h.load_state_dict(ck["weights"]["sat"])
-    return ck.get("step", 0), ck.get("seen_tok", 0)
-
-def _flush_delta():
-    """Wait for any in-flight delta save to complete."""
-    global _delta_thread
-    if _delta_thread is not None and _delta_thread.is_alive():
-        print("  [delta] flushing in-flight write...")
-        _delta_thread.join(timeout=120)
-
 def save_ckpt(path: pathlib.Path, core, ar_h, sat_h, opt, scaler, meta):
     path.parent.mkdir(exist_ok=True, parents=True)
     tmp = path.with_suffix(path.suffix + ".tmp")
@@ -700,7 +604,6 @@ def _train_phase(
     MAX_OOM_RETRIES = 2
     now_wall = time.time()
     last_save_mono = time.monotonic() - (now_wall - (resume_wall_time or now_wall))
-    last_delta_step = start_step
     print(f"[{phase_name}] Starting. Goal: {total_tokens_needed:,} tokens. Batch={BATCH}, Block={BLOCK}")
     print(f"[{phase_name}] AR_ONLY={args.ar_only}, TIE_WEIGHTS={tie_weights}, STREAMING={streaming}")
     while seen_tok < total_tokens_needed:
@@ -774,19 +677,10 @@ def _train_phase(
             now_mono = time.monotonic()
             if now_mono - last_save_mono >= args.save_every_sec:
                 ck_name = f"{phase_name}_step{step:08d}.pt"
-                _flush_delta()  # wait for any in-flight delta before full save
                 _prune_checkpoints(pathlib.Path(args.save_dir), phase_name, max_ckpts)
                 save_ckpt(pathlib.Path(args.save_dir) / ck_name, core, ar_h, sat_h, opt, scaler,
                           meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights})
                 last_save_mono = now_mono
-                # Prune old deltas after a full save (they're superseded)
-                _prune_deltas(pathlib.Path(args.save_dir), phase_name, args.delta_max_keep)
-                last_delta_step = step  # reset delta counter after full save
-        # ── Delta checkpoint (step-based, weight-only, async) ──
-        if args.delta_every_steps > 0 and (step - last_delta_step) >= args.delta_every_steps:
-            _prune_deltas(pathlib.Path(args.save_dir), phase_name, args.delta_max_keep)
-            save_delta(core, ar_h, sat_h, step, seen_tok, pathlib.Path(args.save_dir), phase_name)
-            last_delta_step = step
         if args.auto_grow:
             steps_since_last_grow += 1
             if steps_since_last_grow >= args.grow_every_steps:
@@ -800,7 +694,6 @@ def _train_phase(
                 except ValueError:
                     grow_plan = sorted(set(grow_plan + [BLOCK]))
     pbar.close()
-    _flush_delta()  # ensure any in-flight delta completes before final save
     save_ckpt(pathlib.Path(args.save_dir) / f"{phase_name}_final.pt", core, ar_h, sat_h, opt, scaler,
               meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights})
     return step, seen_tok, time.time()
@@ -844,11 +737,7 @@ def train(args):
     ])
     scaler = GradScaler(enabled=(args.amp and DEV.type == "cuda"))
     start_step, seen_tok, last_wall = 0, 0, None
-    if args.resume_delta and not args.fresh:
-        delta_step, delta_tok = load_delta(pathlib.Path(args.resume_delta), core, ar_h, sat_h)
-        start_step, seen_tok, last_wall = delta_step, delta_tok, None
-        print(f"Resumed from DELTA at step {start_step} (optimizer state reset — momentum rebuilds in ~100 steps)")
-    elif args.resume and not args.fresh:
+    if args.resume and not args.fresh:
         start_step, seen_tok, last_wall = load_ckpt(pathlib.Path(args.resume), core, ar_h, sat_h, opt, scaler)
         print(f"Resumed from step {start_step}")
     # torch.compile AFTER loading checkpoint (key names differ)
@@ -1046,9 +935,6 @@ def main():
     tr.add_argument("--amp", action="store_true")
     tr.add_argument("--compile", action="store_true", help="Use torch.compile for speedup")
     tr.add_argument("--save_every_sec", type=int, default=DEFAULT_SAVE_SEC)
-    tr.add_argument("--delta_every_steps", type=int, default=DEFAULT_DELTA_STEPS, help="Weight-only delta save every N steps (0=off)")
-    tr.add_argument("--delta_max_keep", type=int, default=DEFAULT_MAX_DELTAS, help="Max delta checkpoints to keep")
-    tr.add_argument("--resume_delta", type=str, help="Resume from a delta (weight-only, no optimizer state)")
     tr.add_argument("--save_dir", default=str(CKDIR))
     tr.add_argument("--resume", type=str)
     tr.add_argument("--x2", action="store_true")