1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
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)
|
