diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 3 |
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): |