Stage-2 data pipeline: FineWeb-Edu sample-10BT -> 32k BPE -> bins; --data flag on both trainers

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

3 files changed, 100 insertions, 2 deletions

diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py
index d45327b..ebd7536 100644
--- a/ep_run/casc_bp_train.py
+++ b/ep_run/casc_bp_train.py
@@ -25,12 +25,13 @@ ap.add_argument('--resume', default='')                  # path to a ckpt (tok/p
 ap.add_argument('--olmo2', action='store_true')          # OLMo2-standard block (see casc_eq_train.py)
 ap.add_argument('--wd', type=float, default=-1.0)        # >=0: grouped weight decay; <0 = legacy uniform 1e-4
 ap.add_argument('--zloss', type=float, default=0.0)      # z-loss coefficient; 0 = off
+ap.add_argument('--data', default='tinystories_bpe')     # dataset dir under ep_run/data
 args = ap.parse_args()
 if args.olmo2 and args.tok_init <= 0: args.tok_init = 0.02
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe')
+DD = Path('/home/yurenh2/ept/ep_run/data') / args.data
 vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size']
 
 def get_batch(split):
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 6a05cf5..3f34daa 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -46,6 +46,7 @@ ap.add_argument('--amp', action='store_true')            # PROPER mixed precisio
 ap.add_argument('--dtop_every', type=int, default=1)    # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr)
 ap.add_argument('--gate_every', type=int, default=200)  # in-training cos(EP,BP) telemetry; <=0 = fully BP-free (no bp_gate at all)
 ap.add_argument('--gate_govern', action='store_true')   # let gate cos adjust K/bscale (default: observe-only => training control is BP-free)
+ap.add_argument('--data', default='tinystories_bpe')     # dataset dir under ep_run/data (train.bin/val.bin/meta.pkl)
 args = ap.parse_args()
 if args.olmo2:
     args.untie = True
@@ -53,7 +54,7 @@ if args.olmo2:
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe')
+DD = Path('/home/yurenh2/ept/ep_run/data') / args.data
 vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size']
 
 def get_batch(split):
diff --git a/ep_run/prepare_fineweb.py b/ep_run/prepare_fineweb.py
new file mode 100644
index 0000000..33fce59
--- /dev/null
+++ b/ep_run/prepare_fineweb.py
@@ -0,0 +1,96 @@
+"""FineWeb-Edu sample-10BT -> 32k ByteLevel BPE -> train.bin/val.bin (uint16) + meta.pkl + tokenizer.json.
+Same bin format as tinystories_bpe; trainers consume it via --data fineweb_edu.
+Stage-2 corpus (recipe order: TinyStories epoch -> FineWeb-Edu -> OLMo2-corpus class).
+
+Phases (all resumable-ish, markers for the watcher):
+  1. snapshot_download of sample/10BT parquets (~28GB)          -> DOWNLOAD_DONE
+  2. train 32k BPE on ~1.5GB sample from shard 0               -> TOKENIZER_DONE
+  3. tokenize shard 0: first 20M tokens -> val.bin, rest ->
+     train.bin (disjoint)                                       -> SMOKE_READY
+  4. tokenize remaining shards, append to train.bin             -> DONE_FINEWEB
+Docs are joined with a <|eot|> separator (id 0). vocab 32768 fits uint16.
+NFS note: peak disk = raw parquet ~28GB + bins ~20GB; keep raw/ for tokenizer reruns.
+"""
+import pickle, time
+from pathlib import Path
+import numpy as np
+import pyarrow.parquet as pq
+from huggingface_hub import snapshot_download
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import ByteLevel
+from tokenizers.decoders import ByteLevel as ByteLevelDec
+
+D = Path('/home/yurenh2/ept/ep_run/data/fineweb_edu')
+RAW = D / 'raw'
+D.mkdir(parents=True, exist_ok=True)
+VOCAB = 32768
+EOT = '<|eot|>'
+SAMPLE_BYTES = 1_500_000_000   # tokenizer-training text cap
+VAL_TOKENS = 20_000_000
+
+t0 = time.time()
+snapshot_download(repo_id='HuggingFaceFW/fineweb-edu', repo_type='dataset',
+                  allow_patterns=['sample/10BT/*'], local_dir=str(RAW))
+shards = sorted((RAW / 'sample' / '10BT').glob('*.parquet'))
+print(f"DOWNLOAD_DONE {len(shards)} shards {time.time()-t0:.0f}s", flush=True)
+
+def docs(shard, batch_rows=2000):
+    for b in pq.ParquetFile(shard).iter_batches(batch_size=batch_rows, columns=['text']):
+        for t in b.column('text').to_pylist():
+            if t: yield t
+
+tok_path = D / 'tokenizer.json'
+if tok_path.exists():
+    tok = Tokenizer.from_file(str(tok_path))
+    print(f"TOKENIZER_DONE (cached) vocab={tok.get_vocab_size()}", flush=True)
+else:
+    def sample_iter():
+        n = 0
+        for t in docs(shards[0]):
+            yield t
+            n += len(t)
+            if n > SAMPLE_BYTES: return
+    tok = Tokenizer(BPE(unk_token=None))
+    tok.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    tok.decoder = ByteLevelDec()
+    tok.train_from_iterator(sample_iter(), BpeTrainer(vocab_size=VOCAB, special_tokens=[EOT], show_progress=False))
+    tok.save(str(tok_path))
+    print(f"TOKENIZER_DONE vocab={tok.get_vocab_size()} {time.time()-t0:.0f}s", flush=True)
+eot_id = tok.token_to_id(EOT)
+assert tok.get_vocab_size() <= 65535 and eot_id is not None
+pickle.dump({'vocab_size': tok.get_vocab_size()}, open(D / 'meta.pkl', 'wb'))
+
+def encode_shard(shard, sink, head_sink=None, head_n=0):
+    """tokenize one shard; first head_n tokens to head_sink (if given), rest to sink. Returns tokens written."""
+    wrote_head, n, buf = 0, 0, []
+    def flush(buf):
+        nonlocal wrote_head, n
+        ids = []
+        for e in tok.encode_batch(buf):
+            ids.extend(e.ids); ids.append(eot_id)
+        a = np.array(ids, dtype=np.uint16)
+        if head_sink is not None and wrote_head < head_n:
+            k = min(head_n - wrote_head, len(a))
+            a[:k].tofile(head_sink); wrote_head += k; a = a[k:]
+        a.tofile(sink); n += len(a)
+    for t in docs(shard):
+        buf.append(t)
+        if len(buf) >= 1000: flush(buf); buf = []
+    if buf: flush(buf)
+    return n + wrote_head
+
+train_f = open(D / 'train.bin', 'wb')
+val_f = open(D / 'val.bin', 'wb')
+n0 = encode_shard(shards[0], train_f, head_sink=val_f, head_n=VAL_TOKENS)
+val_f.close(); train_f.flush()
+print(f"SMOKE_READY shard0 {n0/1e6:.0f}M tokens (val {VAL_TOKENS/1e6:.0f}M) {time.time()-t0:.0f}s", flush=True)
+
+total = n0
+for i, sh in enumerate(shards[1:], 1):
+    ns = encode_shard(sh, train_f)
+    total += ns
+    print(f"[shard {i}/{len(shards)-1}] +{ns/1e6:.0f}M -> total {total/1e9:.2f}B tokens {time.time()-t0:.0f}s", flush=True)
+train_f.close()
+print(f"DONE_FINEWEB total {total/1e9:.2f}B tokens (train {(total-VAL_TOKENS)/1e9:.2f}B / val {VAL_TOKENS/1e6:.0f}M) vocab {tok.get_vocab_size()} {time.time()-t0:.0f}s", flush=True)