diff options
Diffstat (limited to 'ep_run')
| -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 |
3 files changed, 62 insertions, 1 deletions
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) |