1 files changed, 49 insertions, 0 deletions
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)