"""Config-driven pretraining loop for Stoicheia. Single-process or torchrun/DDP. One JSON config freezes size, schedule, and budget. Two-phase curriculum: 3-tier stable phase -> gold-only anneal in the WSD decay window. Checkpoint/resume by step (every ckpt_every steps) — safe to run as a chain of independent, dependency-linked SLURM jobs instead of one long request. torchrun --nproc_per_node=4 -m train.train --config configs/stoicheia.json """ from __future__ import annotations import argparse, contextlib, json, os, sys, time from pathlib import Path import torch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from model.char_bert import CharBertConfig, CharBertEncoder, num_params from train.collate import pack_batch from train.data import MultiTierLoader, stable_cfg, anneal_cfg from train.loss import compute_loss from train.noising import NoiseConfig from train.schedule import wsd_dyn def ddp_setup(): if "RANK" in os.environ: import datetime import torch.distributed as dist # Generous timeout: rank 0 writes a ~5GB checkpoint to shared disk and then # evaluates, while every other rank waits in the next collective. Under # filesystem contention (many concurrent jobs) that write alone has exceeded # 60 min and tripped the NCCL watchdog, SIGABRT-ing whole runs mid-campaign. # 180 min buys tolerance for a slow filesystem. Do NOT instead bound the eval # with signal.alarm(): SIGALRM interrupts the write and torch.save fails with # EINTR, converting a slow-but-recoverable step into a hard failure. dist.init_process_group("nccl", timeout=datetime.timedelta(minutes=180)) rank = dist.get_rank(); world = dist.get_world_size() torch.cuda.set_device(rank % torch.cuda.device_count()) return rank, world, True return 0, 1, False def eval_stalled(eval_path: Path, window=8, eps=0.002, split="train"): """True if held-out bits_per_char has stopped improving: best of the most recent window/2 evals is not at least eps bits better than the best of the window/2 before. Eval masking is deterministic (fixed seed + records), so the curve is low-noise and a small eps suffices. Only entries of the given eval split are considered, so a run that switched eval source (train-holdout -> dev) restarts its stall window cleanly.""" if not eval_path.exists(): return False bpc = [] for line in eval_path.read_text().splitlines(): try: e = json.loads(line) if e.get("split", "train") == split: bpc.append(e["bits_per_char"]) except Exception: continue if len(bpc) < window: return False half = window // 2 old_best = min(bpc[-window:-half]) new_best = min(bpc[-half:]) return new_best > old_best - eps def save_ckpt(obj, path: Path): """Atomic checkpoint write: a wall-time kill mid-save must never corrupt the resume point (the chain of jobs depends on last.pt always being loadable).""" tmp = path.with_name(path.name + ".tmp") torch.save(obj, tmp) os.replace(tmp, path) def infinite_records(loader, chunk=256): while True: for r in loader.records(chunk): yield r class BatchDataset(torch.utils.data.IterableDataset): """Produces packed batches under a multiprocess DataLoader — N worker processes build batches in PARALLEL. Single-threaded collate (masking is a per-record Python loop) is slow enough to starve the GPU otherwise. Each (rank, worker) gets a disjoint data shard and its own RNG, so ranks/workers never collide and runs stay reproducible.""" def __init__(self, gdata, tier_weights, gold_only, ncfg, T, rows, seed, rank, world, exclude_holdout=True): super().__init__() self.__dict__.update(locals()) def __iter__(self): info = torch.utils.data.get_worker_info() wid = info.id if info else 0 nw = info.num_workers if info else 1 gshard = self.rank * nw + wid gtot = self.world * nw cfg = (anneal_cfg(self.gdata, window=self.T, seed=self.seed, exclude_holdout=self.exclude_holdout) if self.gold_only else stable_cfg(self.gdata, w=tuple(self.tier_weights), window=self.T, seed=self.seed, exclude_holdout=self.exclude_holdout)) loader = MultiTierLoader(cfg, rank=gshard, world_size=gtot) g = torch.Generator().manual_seed(self.seed * 100003 + gshard) gen = infinite_records(loader) while True: yield pack_batch(gen, self.ncfg, self.T, self.rows, g) def make_loader(gdata, tier_weights, gold_only, ncfg, T, rows, seed, rank, world, num_workers, exclude_holdout=True): ds = BatchDataset(gdata, tier_weights, gold_only, ncfg, T, rows, seed, rank, world, exclude_holdout=exclude_holdout) return torch.utils.data.DataLoader( ds, batch_size=None, num_workers=num_workers, prefetch_factor=(2 if num_workers > 0 else None), persistent_workers=(num_workers > 0), pin_memory=True) def main(): ap = argparse.ArgumentParser() ap.add_argument("--config", required=True) a = ap.parse_args() cfg = json.loads(Path(a.config).read_text()) cfg["out_dir"] = os.path.expandvars(cfg["out_dir"]) # allow "$STOICHEIA_DATA/..." in configs rank, world, is_ddp = ddp_setup() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") torch.manual_seed(cfg["seed"] + rank) torch.set_float32_matmul_precision("high") torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True gdata = os.environ["STOICHEIA_DATA"] # dev-driven regime: eval_shards points at the fold's REAL val split (unseen works); # all training decisions (anneal stall, early stop, best.pt) then key off dev, and # train_holdout=false returns the intra-train mod-200 holdout to training. eval_shards = cfg.get("eval_shards") if eval_shards: eval_shards = os.path.expandvars(eval_shards) eval_split = "val" if eval_shards else "train" T = cfg["seq_len"]; rows = cfg.get("rows", cfg["micro_batch"]) total = cfg["total_steps"] out = Path(cfg["out_dir"]); out.mkdir(parents=True, exist_ok=True) metrics_f = out / f"metrics_rank{rank}.jsonl" mcfg = CharBertConfig(attn_impl=cfg.get("attn", "flex"), d_model=cfg["d_model"], n_heads=cfg["d_model"] // 64, depth=cfg["depth"], char_window=cfg["char_window"], qk_norm=cfg.get("qk_norm", True)) model = CharBertEncoder(mcfg).to(device) if rank == 0: print(f"params={num_params(model)/1e6:.1f}M attn={mcfg.attn_impl} " f"T={T} rows={rows} world={world} total_steps={total}") if is_ddp: from torch.nn.parallel import DistributedDataParallel as DDP model = DDP(model, device_ids=[rank % torch.cuda.device_count()]) fwd = torch.compile(model) if cfg.get("compile", True) and device.type == "cuda" else model # decay only matrices; 1-D params (RMSNorm gains) are shape-constrained, not capacity decay_p = [p for p in model.parameters() if p.ndim >= 2] nodecay_p = [p for p in model.parameters() if p.ndim < 2] opt = torch.optim.AdamW( [{"params": decay_p, "weight_decay": cfg["wd"]}, {"params": nodecay_p, "weight_decay": 0.0}], lr=cfg["lr"], betas=(0.9, 0.95), fused=(device.type == "cuda")) ncfg = NoiseConfig() # tuned defaults: span/word/elastic/iid/halfword/substitute w = tuple(cfg.get("tier_weights", [1.0, 1.0, 0.3])) nworkers = cfg.get("num_workers", 16) step0 = 0 ckpt = out / "last.pt" if ckpt.exists(): sd = torch.load(ckpt, map_location=device) (model.module if is_ddp else model).load_state_dict(sd["model"]) opt.load_state_dict(sd["opt"]); step0 = sd["step"] if rank == 0: print(f"resumed from step {step0}") # completion check must mirror the dynamic schedule below: the run ends at # anneal_start + decay_len (marker-decided or capped), not at nominal total_steps. _decay_len = total - int(total * (1 - cfg.get("decay_frac", 0.2))) _hard_max = max(cfg.get("hard_max_steps", total), total) _marker = out / "anneal_start.json" if cfg.get("auto_anneal", True) and _marker.exists(): _end = min(json.loads(_marker.read_text())["step"], _hard_max - _decay_len) + _decay_len elif cfg.get("auto_anneal", True): _end = _hard_max else: _end = total if step0 >= _end: if rank == 0: print(f"run already complete (resume step {step0} >= end step {_end}); " f"nothing to do") if is_ddp: import torch.distributed as dist dist.destroy_process_group() return # fold the resume step into every RNG so each chained job sees a FRESH data order and # fresh noise masks instead of replaying the stream from position 0 (the loader is an # infinite sampler — there is no cheap "skip to batch step0", a reseed is equivalent). data_seed = cfg["seed"] + step0 torch.manual_seed(cfg["seed"] + rank + step0 * 7919) # drives DataLoader worker seeds excl_holdout = cfg.get("train_holdout", True) stable_it = iter(make_loader(gdata, list(w), False, ncfg, T, rows, data_seed, rank, world, nworkers, exclude_holdout=excl_holdout)) # anneal data: default = classic gold-only switch. anneal_phases staggers the mix over # the decay window instead — [[end_frac, [w_gold, w_silver, w_bronze]], ...] — so the # highest-repetition data only dominates once the LR is too small to overfit on it # (the hard gold-only switch at near-peak LR degraded held-out bpc by +0.04 in run 1). phases = cfg.get("anneal_phases") if phases: phase_ends, phase_its, _cache = [], [], {tuple(w): stable_it} for f, pw in phases: key = tuple(pw) if key not in _cache: _cache[key] = iter(make_loader(gdata, list(pw), False, ncfg, T, rows, data_seed, rank, world, nworkers, exclude_holdout=excl_holdout)) phase_ends.append(f); phase_its.append(_cache[key]) else: phase_ends = [1.0] phase_its = [iter(make_loader(gdata, list(w), True, ncfg, T, rows, data_seed, rank, world, nworkers, exclude_holdout=excl_holdout))] # gradient accumulation: global batch per GPU = rows * grad_accum. The swept setup is # 8 rows/GPU; running it as 4 x 2 HALVES activation memory (the 405M model at T=8192, # rows=8, depth=32 needs ~80GB of activations — over the GH200's 95GB once eval and # allocator overhead are counted) with identical training math. accum = cfg.get("grad_accum", 1) # auto-anneal, SYMMETRIC: the anneal starts when the held-out curve stalls — which can # be EARLIER than the planned step (don't burn budget on a flat curve) or LATER (don't # undertrain a model that is still learning; the stable phase extends past the planned # point up to hard_max_steps - decay_len). The decay window always keeps its planned # length. The decision persists in a marker file (survives chain-job restarts) and is # broadcast so all ranks switch at the same step. total_steps is thus a NOMINAL budget; # the actual end is anneal_start + decay_len, capped by hard_max_steps. planned_anneal = int(total * (1 - cfg.get("decay_frac", 0.2))) decay_len = total - planned_anneal hard_max = max(cfg.get("hard_max_steps", total), total) max_start = hard_max - decay_len # latest possible anneal start auto = cfg.get("auto_anneal", True) min_start = int(total * cfg.get("min_anneal_frac", 0.35)) marker = out / "anneal_start.json" anneal_start = None # None = not yet decided (stall or cap decides) if not auto: anneal_start = planned_anneal elif marker.exists(): anneal_start = min(json.loads(marker.read_text())["step"], max_start) if rank == 0: print(f"auto-anneal marker: anneal from step {anneal_start}") model.train() t0 = time.time(); seen_chars = 0 consec_skips = 0; total_skips = 0; diverged = False; stop_early = False end_step = (anneal_start if anneal_start is not None else max_start) + decay_len for step in range(step0, hard_max): if anneal_start is None and step >= max_start: anneal_start = max_start # cap reached: deterministic on all ranks a0 = anneal_start if anneal_start is not None else max_start end_step = a0 + decay_len if step >= end_step: break # anneal finished: run is complete anneal = step >= a0 # durable branch point: the last flat-LR state before the anneal. Lets us re-run a # LONGER or different anneal (data mix, decay shape, soup of several) afterwards # for 20% of the cost, instead of repeating the whole stable phase. if rank == 0 and step == a0 and not (out / "pre_anneal.pt").exists(): core = model.module if is_ddp else model save_ckpt(dict(model=core.state_dict(), opt=opt.state_dict(), step=step, cfg=cfg), out / "pre_anneal.pt") if not marker.exists(): marker.write_text(json.dumps(dict(step=a0, reason="cap"))) print(f" saved pre_anneal.pt at step {step} (anneal branch point)", flush=True) if anneal: fr = (step - a0) / max(decay_len, 1) pi = 0 while pi < len(phase_ends) - 1 and fr >= phase_ends[pi]: pi += 1 if rank == 0 and pi != getattr(main, "_pi", -1): main._pi = pi print(f" anneal phase {pi + 1}/{len(phase_ends)} from step {step} " f"(mix {phases[pi][1] if phases else 'gold-only'})", flush=True) it = phase_its[pi] else: it = stable_it lr = cfg["lr"] * wsd_dyn(step, total, a0, cfg.get("warmup_frac", 0.04), cfg.get("decay_frac", 0.2)) for pg in opt.param_groups: pg["lr"] = lr opt.zero_grad(set_to_none=True) for micro in range(accum): batch = next(it) batch = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()} # skip the allreduce on all but the last micro-batch. DDP reads the no_sync # flag during FORWARD, so the forward must run inside the context too. sync_ctx = (model.no_sync() if (is_ddp and micro < accum - 1) else contextlib.nullcontext()) with sync_ctx: with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): out_h = fwd(batch) loss, logs = compute_loss(out_h, batch, lam=cfg.get("lam", 0.1)) (loss / accum).backward() seen_chars += int((batch["seg_id"] > 0).sum()) * world # real (non-pad) chars gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.get("clip", 1.0)) # NaN-skip guard: one pathological batch must not poison the weights. # Check BEFORE opt.step() so non-finite grads are never applied. if torch.isfinite(gnorm): opt.step() consec_skips = 0 else: consec_skips += 1; total_skips += 1 if rank == 0: print(f" [skip] non-finite grad at step {step} (consecutive={consec_skips}, " f"total={total_skips})", flush=True) if consec_skips >= cfg.get("max_consec_skips", 25): if rank == 0: print(f" ABORT: {consec_skips} consecutive non-finite steps — real " f"divergence, lower LR.", flush=True) diverged = True break if rank == 0 and (step % cfg.get("log_every", 20) == 0): dt = time.time() - t0 mem = (round(torch.cuda.max_memory_allocated() / 2**30, 1) if device.type == "cuda" else 0) rec = dict(step=step, lr=round(lr, 6), gnorm=round(float(gnorm), 3), chps=round(seen_chars / dt / 1e6, 3), mem_gb=mem, anneal=anneal, **logs) print(" " + " ".join(f"{k}={v}" for k, v in rec.items()), flush=True) with open(metrics_f, "a") as mf: mf.write(json.dumps(rec) + "\n") if step > step0 and step % cfg.get("ckpt_every", 2000) == 0: if rank == 0: core = model.module if is_ddp else model # step+1: this step is DONE — resume must continue at the next one save_ckpt(dict(model=core.state_dict(), opt=opt.state_dict(), step=step + 1, cfg=cfg), ckpt) if cfg.get("eval_every_ckpt", True): try: from eval.intrinsic import held_out_records, evaluate, restore_demo if eval_shards: from eval.val_eval import val_records recs = val_records(eval_shards, cfg.get("eval_n", 256)) else: recs = held_out_records(f"{gdata}/shards/v1_punct", cfg.get("eval_n", 256)) m = evaluate(core, recs, device) m["demo"] = restore_demo(core, device) m["step"] = step m["split"] = eval_split print(" EVAL " + json.dumps(m, ensure_ascii=False), flush=True) with open(out / "eval.jsonl", "a") as ef: ef.write(json.dumps(m, ensure_ascii=False) + "\n") evs = [e for e in (json.loads(l) for l in open(out / "eval.jsonl") if "bits_per_char" in l) if e.get("split", "train") == eval_split] # best-checkpoint tracking: never lose the best held-out model, # whatever the anneal tail does (stateless: derived from eval.jsonl) if m["bits_per_char"] <= min(e["bits_per_char"] for e in evs): import shutil shutil.copyfile(ckpt, out / "best.pt") print(f" new best.pt (bpc {m['bits_per_char']})", flush=True) # anneal early-stop: gold-only + decaying LR can tip into overfitting; # if held-out hasn't improved for `anneal_patience` evals, stop — # best.pt is the product. # only judge the anneal from mid-decay on: WSD gains arrive when # the LR falls through ~50%, and the gold-switch at near-peak LR # causes a transient (observed +0.02 bpc at step 68k) that must # not be mistaken for a stalled anneal. pat = cfg.get("anneal_patience", 3) if anneal and pat > 0 and step >= a0 + decay_len // 2: ann = [e for e in evs if e["step"] > (anneal_start or 0)] if len(ann) > pat: bi = min(range(len(ann)), key=lambda i: ann[i]["bits_per_char"]) if len(ann) - 1 - bi >= pat: stop_early = True print(f" ANNEAL EARLY-STOP: no held-out improvement in " f"{pat} evals (best at step {ann[bi]['step']}); " f"best.pt is the final model", flush=True) except Exception as e: print(f" EVAL failed: {str(e)[:120]}", flush=True) # auto-anneal trigger: anneal the moment the held-out curve stalls if (auto and anneal_start is None and step >= min_start and eval_stalled(out / "eval.jsonl", cfg.get("stall_window", 8), cfg.get("stall_eps", 0.002), split=eval_split)): anneal_start = min(step + 1, max_start) marker.write_text(json.dumps(dict(step=anneal_start, reason="stall", triggered_at=step))) print(f" AUTO-ANNEAL: held-out bits/char stalled — annealing from step " f"{anneal_start} (planned {planned_anneal}, cap {max_start}); " f"run ends at {anneal_start + decay_len}", flush=True) elif auto and anneal_start is None and step == planned_anneal: print(f" EXTEND: still improving at planned anneal step {planned_anneal} " f"— stable phase continues (anneal by step {max_start} at latest)", flush=True) if is_ddp: import torch.distributed as dist t_a = torch.tensor([-1 if anneal_start is None else anneal_start, 1 if stop_early else 0], dtype=torch.long, device=device) dist.broadcast(t_a, 0) v = int(t_a[0].item()) anneal_start = None if v < 0 else v stop_early = bool(int(t_a[1].item())) if stop_early: break if rank == 0: if diverged: # do NOT write final.pt / advance last.pt: the run is not complete, and the # last good checkpoint is the thing to restart from (with a lower LR). print("ABORTED (divergence) — last good checkpoint left untouched") else: # step is stamped as end_step even on early stop so remaining chain jobs no-op; # stopped_at records the true last trained step. On early stop the model to USE # is best.pt (the early-stop criterion means the final weights are not the best). state = dict(model=(model.module if is_ddp else model).state_dict(), opt=opt.state_dict(), step=end_step, cfg=cfg, stopped_early=stop_early) save_ckpt(state, out / "final.pt") save_ckpt(state, ckpt) # last.pt at the end step: remaining chain jobs no-op print(f"DONE (end step {end_step}, nominal budget {total}" + (", EARLY-STOPPED — use best.pt" if stop_early else "") + ")") if is_ddp: import torch.distributed as dist dist.destroy_process_group() if diverged: sys.exit(1) if __name__ == "__main__": main()