summaryrefslogtreecommitdiff
path: root/ep_run/casc_eq_train.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-13 05:58:48 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-13 05:58:48 -0500
commit1d638edd83dc888cb392e902eb86ad57b22949a1 (patch)
tree4438279fd15d977acd1d678f47cd54cfd46cfa52 /ep_run/casc_eq_train.py
parent27cecb987fff1d6acd2f2e7bfaa7f8a01670bde5 (diff)
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
Diffstat (limited to 'ep_run/casc_eq_train.py')
-rw-r--r--ep_run/casc_eq_train.py3
1 files changed, 2 insertions, 1 deletions
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):