diff options
| -rw-r--r-- | docs/campaign/CASCADE_ABLATION_PLAN.md | 17 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 8 | ||||
| -rw-r--r-- | ep_run/muon.py | 6 | ||||
| -rw-r--r-- | ep_run/shuffle_fineweb.py | 49 |
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) |