diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 43 |
1 files changed, 43 insertions, 0 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index ed164ce..05c1e4e 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -53,6 +53,11 @@ ap.add_argument('--dgain_geo', type=float, default=0.0) # >0: per-layer geomet # gain = geo^l (layer 0 = x1), optionally capped ap.add_argument('--dgain_geo_cap', type=float, default=0.0) # >0: cap for the geometric profile ap.add_argument('--dgain_rand', type=float, default=0.0) +ap.add_argument('--gen', type=int, default=0) # >0: sample text from the resumed checkpoint and exit +ap.add_argument('--gen_prompts', default='') # '|'-separated prompts; empty = built-in set +ap.add_argument('--gen_temp', type=float, default=0.8) +ap.add_argument('--gen_topk', type=int, default=40) +ap.add_argument('--gen_new', type=int, default=120) ap.add_argument('--probe_dgspec', type=int, default=0) # >0: M1 spectroscopy, value = n batches; exits before training ap.add_argument('--probe_gains', default='1,2,4,8,16,32,64,128,256') ap.add_argument('--probe_f64', action='store_true') # fp64 states+model in the probe: the fp-floor decisive arm # >1: per-STEP log-uniform dgain_top in @@ -859,6 +864,44 @@ if args.ddp_grad_test: import sys sys.exit(0) +if args.gen > 0: + # Sampling from a resumed checkpoint. Inference here is an ordinary forward pass through the + # blocks, which is the point: EP appears only in training. The same code path loads an EP or a + # backprop checkpoint, since both scripts save the same keys, so the two can be sampled side by + # side under identical settings. + import sys + from tokenizers import Tokenizer as _Tok + _tk = _Tok.from_file(str(DD / 'tokenizer.json')) + prompts = [p for p in args.gen_prompts.split('|') if p] or [ + 'The main difference between a virus and a bacterium is', + 'To find the area of a circle, you', + 'In 1815, the eruption of Mount Tambora', + 'Photosynthesis is the process by which', + ] + @torch.no_grad() + def _sample(prompt, n_new, temp, topk, seed): + torch.manual_seed(seed) + ids = _tk.encode(prompt).ids[:args.T - n_new - 1] + idx = torch.zeros(1, args.T, dtype=torch.long, device=dev) + L = len(ids) + idx[0, :L] = torch.tensor(ids, device=dev) + for _ in range(n_new): + if L >= args.T: break + z = emb(idx) + for b in blocks: z = b(z, mask) + lg = readout(z)[0, L - 1].float() / max(temp, 1e-6) + v, _i = torch.topk(lg, topk) + lg[lg < v[-1]] = -float('inf') + nt = torch.multinomial(F.softmax(lg, -1), 1).item() + idx[0, L] = nt; L += 1 + return _tk.decode(idx[0, :L].tolist()) + print(f'# samples from {args.resume} (temp {args.gen_temp}, top-k {args.gen_topk})', flush=True) + for pi, p in enumerate(prompts): + for s in range(args.gen): + print(f'\n--- prompt {pi + 1}, sample {s + 1} ---', flush=True) + print(_sample(p, args.gen_new, args.gen_temp, args.gen_topk, 1000 * pi + s), flush=True) + sys.exit(0) + if args.probe_dgspec > 0: # M1 DGAIN SPECTROSCOPY: per-block leak vector vs uniform read-displacement gain. # Paired design: L_l(g) = mean_b[gEP_l(g) - gBP_l] on the SAME batch — batch-sampling |