1 files changed, 22 insertions, 0 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index f115a91..fb9dca5 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -54,6 +54,9 @@ ap.add_argument('--ddp_grad_test', action='store_true')  # one-step grad equival
 ap.add_argument('--beta_ride', type=float, default=1.0)   # cap ceiling: >1 lets the governor RAISE beta
                                                           # above the schedule, up to ride x schedule
 ap.add_argument('--beta_ride_up', type=float, default=1.02)  # per-step climb rate in the calm branch
+ap.add_argument('--bpmix', default='')                    # ABLATION: overwrite EP grads with TRUE BP grads
+                                                          # for selected groups: 'blocks:0-5' | 'blocks:6-11'
+                                                          # | 'attn' | 'ffn' | 'head' (comma-separated)
 ap.add_argument('--beta_sync', type=int, default=0)       # >0: SYNCHRONOUS acceptance — this step's own
                                                           # relax telemetry gates the commit; on reject,
                                                           # halve beta and retry same batch (N halvings max)
@@ -729,6 +732,25 @@ for step in range(start_step, args.steps + 1):
                 GOV['K'] = min(GOV['K'] + 2, args.kmax); GOV['bscale'] = max(GOV['bscale'] * 0.7, 0.05)
             elif gcos > 0.995 and GOV['K'] > args.K:
                 GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0)
+    if args.bpmix and ok:
+        with torch.enable_grad():
+            _bg = bp_gate(x, y)                      # true BP grads, same batch
+        sel = set()
+        for spec in args.bpmix.split(','):
+            spec = spec.strip()
+            if spec.startswith('blocks:'):
+                a_, b_ = spec.split(':')[1].split('-')
+                for bi in range(int(a_), int(b_) + 1):
+                    sel.update(id(p) for p in blocks[bi].parameters())
+            elif spec == 'attn':
+                for blk in blocks: sel.update(id(p) for m in (blk.attn, blk.na) for p in m.parameters())
+            elif spec == 'ffn':
+                for blk in blocks: sel.update(id(p) for m in (blk.ff, blk.nf) for p in m.parameters())
+            elif spec == 'head':
+                sel.update(id(p) for p in (list(tok.parameters()) + ([W_out] if isinstance(W_out, torch.nn.Parameter) else []) + list(ln_f.parameters())))
+        for p, g in zip(all_params, _bg):
+            if id(p) in sel and g is not None:
+                p.grad = g.detach().clone()
     torch.nn.utils.clip_grad_norm_(all_params, 1.0)
     opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
     if args.qup_bits > 0: