fw72m crown run launched (NCCL DDP prod), exact resume (opt state), RESULT 17, shuffle+prep scripts

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

4 files changed, 79 insertions, 1 deletions

diff --git a/docs/campaign/CASCADE_ABLATION_PLAN.md b/docs/campaign/CASCADE_ABLATION_PLAN.md
index 2710582..c044fb6 100644
--- a/docs/campaign/CASCADE_ABLATION_PLAN.md
+++ b/docs/campaign/CASCADE_ABLATION_PLAN.md
@@ -563,6 +563,23 @@ direction), not training-under-fault; wave-2 = co-training with faults injected
 (expectation from the literature and from (c): tolerances IMPROVE). Feeds COMPONENT_HW_MAP.md
 (per-row status updated) + UIUC outreach dossier.
 
+
+### RESULT 17 (2026-07-13): fw72m LAUNCHED (the 62.7M-crown run) + size ladder complete + resume upgraded.
+- **fw72m**: 72.11M (L12 C512 @32k vocab) x 1.44B FineWeb-Edu tokens (Chinchilla), --amp, first NCCL
+  2-GPU DDP production run (GPU0+3, B12/rank = global B24 preserving recipe semantics). Startup:
+  world=2, cos 0.9999, zero skips, **2.92 it/s -> ETA ~22h**. Crown context: KHS VGG10 = 62.7M
+  (verified from their Table 6: convs 9.2M + dense 25088->2048 = 51.4M + head 2.05M = 62.65M).
+  BP twin queued for GPU1 (waiting on user's phasescan). Corpus: 9.99B tokens, doc-level shuffled
+  (9.67M docs, seed 1234, val re-drawn 20M disjoint) per user directive.
+- **nsize ladder DONE** (TinyStories 4k-vocab, 3k steps, fp32, seed 1): c256 1.9265 / c512 1.7920 /
+  c768 1.7913 / c1024 1.7876 — at FIXED 3k steps quality saturates with width (data/steps-bound,
+  expected); purpose = noise-robustness probe ckpts (runs/nsize_*_s3000.pt x4). Probe script queued.
+- **Resume upgraded to exact**: MultiOpt gains state_dict/load_state_dict; trainer saves 'opt' in
+  every ckpt and restores it on --resume (chunked HPC jobs no longer lose Adam/Muon state).
+- **Delta storage recon**: /scratch 1.5T/1.5T FULL, /work/hdd/bfqt over quota -> 300M data transfer
+  BLOCKED until space found (own old-ept footprint = first cleanup candidate). A100x4 queue ~8.5d,
+  A40 same-day (48h cap -> needs exact resume, now DONE).
+
 ### RESULT 16 (2026-07-13): STAGE-2 DATA PIPELINE LIVE + FINEWEB SMOKE PASSED.
 `prepare_fineweb.py`: FineWeb-Edu sample-10BT -> 32k ByteLevel BPE (<|eot|> id 0) -> uint16 bins,
 tinystories_bpe format, `--data` flag added to both trainers. SMOKE_READY in 11 min (download 5min
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 05273df..573f030 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -412,6 +412,13 @@ def evaluate(nb=6):
         tot += F.cross_entropy(readout(z).reshape(-1, vocab), y.reshape(-1)).item()
     return tot / nb
 
+if args.resume and _ck.get('opt') is not None:
+    try:
+        opt.load_state_dict(_ck['opt'])
+        print('[resume] optimizer state restored (exact chunked-resume)', flush=True)
+    except Exception as e:
+        print(f'[resume] optimizer state NOT restored ({e}) — cold optimizer', flush=True)
+
 if DDP:   # belt & suspenders on top of identical init seeds: rank0's params are law
     with torch.no_grad():
         for p in all_params:
@@ -509,6 +516,7 @@ for step in range(start_step, args.steps + 1):
         torch.save({'tok': tok.state_dict(), 'pos': pos.state_dict(), 'blocks': blocks.state_dict(),
                     'wout': (W_out.detach().cpu() if args.untie else None),
                     'lnf': (ln_f.state_dict() if not isinstance(ln_f, nn.Identity) else None),
+                    'opt': opt.state_dict(),   # full optimizer state -> exact resume for chunked HPC jobs
                     'step': step, 'val': best, 'config': vars(args)}, Path('runs') / f'{args.tag}_s{step}.pt')
 if RANK == 0:
     print(f'[{args.tag}] DONE best val CE {best:.4f}', flush=True)
diff --git a/ep_run/muon.py b/ep_run/muon.py
index 3116030..7281811 100644
--- a/ep_run/muon.py
+++ b/ep_run/muon.py
@@ -38,12 +38,16 @@ class Muon(torch.optim.Optimizer):
 
 
 class MultiOpt:
-    """duck-typed bundle of optimizers (step/zero_grad API-compatible)."""
+    """duck-typed bundle of optimizers (step/zero_grad/state_dict API-compatible)."""
     def __init__(self, opts): self.optimizers = opts
     def step(self):
         for o in self.optimizers: o.step()
     def zero_grad(self, set_to_none=True):
         for o in self.optimizers: o.zero_grad(set_to_none=set_to_none)
+    def state_dict(self):
+        return [o.state_dict() for o in self.optimizers]
+    def load_state_dict(self, sds):
+        for o, sd in zip(self.optimizers, sds): o.load_state_dict(sd)
 
 
 class MultiSched:
diff --git a/ep_run/shuffle_fineweb.py b/ep_run/shuffle_fineweb.py
new file mode 100644
index 0000000..686c3b9
--- /dev/null
+++ b/ep_run/shuffle_fineweb.py
@@ -0,0 +1,49 @@
+"""Document-level order randomization for the fineweb_edu bins (user directive 2026-07-13:
+防止顺序相似文本导致 eval OOD). Two things it fixes:
+  1. val was the HEAD of shard 0 (a narrow, possibly topically-clustered slice) -> val becomes a
+     seeded random document sample from the WHOLE corpus.
+  2. train.bin document order was shard-sequential -> becomes a seeded global permutation
+     (future-proofs sequential/streaming readers; note the current trainer's get_batch already
+     samples uniform-random offsets, so batch order was never sequential).
+Method: full corpus fits RAM (~21GB, box has ~450G free). concat(val, train) restores the exact
+original stream (heals the doc split at the old 20M cut), split on <|eot|>(id 0), permute docs,
+re-draw val (~20M tokens), rewrite via tmp + atomic replace. Marker: DONE_SHUFFLE.
+"""
+import os, time
+import numpy as np
+from pathlib import Path
+
+D = Path('/home/yurenh2/ept/ep_run/data/fineweb_edu')
+EOT, VAL_TOKENS, SEED = 0, 20_000_000, 1234
+t0 = time.time()
+
+val = np.fromfile(D / 'val.bin', dtype=np.uint16)
+train = np.fromfile(D / 'train.bin', dtype=np.uint16)
+full = np.concatenate([val, train]); del val, train
+print(f"loaded {len(full)/1e9:.3f}B tokens {time.time()-t0:.0f}s", flush=True)
+
+ends = np.flatnonzero(full == EOT) + 1          # each doc ends right after its <|eot|>
+starts = np.concatenate([[0], ends[:-1]])
+if ends[-1] != len(full):                        # trailing partial doc (no eot) -> keep as one doc
+    starts = np.append(starts, ends[-1]); ends = np.append(ends, len(full))
+ndoc = len(ends)
+print(f"{ndoc/1e6:.2f}M docs, mean {len(full)/ndoc:.0f} tok/doc {time.time()-t0:.0f}s", flush=True)
+
+rng = np.random.default_rng(SEED)
+perm = rng.permutation(ndoc)
+lens = (ends - starts)[perm]
+cut = int(np.searchsorted(np.cumsum(lens), VAL_TOKENS)) + 1   # first `cut` permuted docs -> val
+
+def write(idx, path):
+    out = np.empty(int((ends[idx] - starts[idx]).sum()), dtype=np.uint16)
+    o = 0
+    for i in idx:
+        n = ends[i] - starts[i]
+        out[o:o + n] = full[starts[i]:ends[i]]; o += n
+    tmp = path.with_suffix('.tmp')
+    out.tofile(tmp); os.replace(tmp, path)
+    return len(out)
+
+nv = write(perm[:cut], D / 'val.bin')
+nt = write(perm[cut:], D / 'train.bin')
+print(f"DONE_SHUFFLE train {nt/1e9:.3f}B val {nv/1e6:.1f}M (seed {SEED}, doc-level) {time.time()-t0:.0f}s", flush=True)