1 files changed, 71 insertions, 0 deletions
diff --git a/ep_run/dp_ep.py b/ep_run/dp_ep.py
new file mode 100644
index 0000000..0d76911
--- /dev/null
+++ b/ep_run/dp_ep.py
@@ -0,0 +1,71 @@
+"""dp_ep.py — data-parallel EP (task #14's last item). EP has no autograd backward, so DDP does not
+apply; this is the minimal manual-allreduce loop: identical init on every rank (same seed), per-rank
+batches, ep_step -> all_reduce(mean) over the canonical param order (zeros for absent grads so the
+collective stays aligned), all_reduce the residual so the jr controller stays in lockstep, identical
+AdamW steps -> weights never diverge. Launch:
+  torchrun --nproc_per_node=N dp_ep.py --steps 300 [recipe flags...]
+Effective batch = B * world (lr scaling deliberately NOT applied for the smoke; tune later)."""
+import argparse, os, time, torch
+import torch.distributed as dist
+import lt_ep_train as L
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--steps', type=int, default=300)
+    ap.add_argument('--B', type=int, default=24)
+    ap.add_argument('--lr', type=float, default=6e-4)
+    ap.add_argument('--seed', type=int, default=0)
+    ap.add_argument('--reg_delay', type=int, default=10 ** 9)  # smoke default: reg-free (Pascal-safe)
+    ap.add_argument('--resreg', type=float, default=0.2)
+    ap.add_argument('--jacreg', type=float, default=0.1)
+    ap.add_argument('--t2sel', type=int, default=40)
+    ap.add_argument('--hr', type=float, default=0.02)
+    ap.add_argument('--holofast', action='store_true')
+    ap.add_argument('--sdpa', action='store_true')
+    ap.add_argument('--log', type=int, default=50)
+    cfg = ap.parse_args()
+
+    dist.init_process_group('nccl')
+    rank, world = dist.get_rank(), dist.get_world_size()
+    torch.cuda.set_device(rank)
+    L.dev = f'cuda:{rank}'
+
+    torch.manual_seed(cfg.seed)                            # identical init on every rank
+    blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick')
+    blk.qknorm = True; blk.track = True
+    blk.holofast, blk.sdpa = cfg.holofast, cfg.sdpa
+    for p in blk.allp:
+        dist.broadcast(p.data, 0)                          # belt & suspenders: exact bitwise identical start
+    torch.manual_seed(cfg.seed * 1009 + rank + 1)          # per-rank data stream
+    opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=1e-4)
+
+    jr = cfg.jacreg
+    t0 = time.time()
+    for step in range(1, cfg.steps + 1):
+        idx, y = L.get_batch('train', cfg.B, blk.T)
+        dly = step < cfg.reg_delay
+        grads, res = L.ep_step(blk, idx, y, 150, 20, 0.1, 0.02, 0.0 if dly else jr, holo=2, hr=cfg.hr,
+                               t1max=300, res_est=1e-4, t2sel=cfg.t2sel, corr_every=1, res_gate=0.0,
+                               resreg=0.0 if dly else cfg.resreg)
+        rt = torch.tensor([res], device=L.dev)
+        dist.all_reduce(rt); res = float(rt) / world       # controller signal identical across ranks
+        for p in blk.allp:                                 # canonical order: aligned collectives
+            g = grads.get(id(p))
+            if g is None:
+                g = torch.zeros_like(p)
+            dist.all_reduce(g)
+            g /= world
+            p.grad = g
+        opt.step(); opt.zero_grad(set_to_none=True)
+        if step % cfg.log == 0 and rank == 0:
+            v = L.evaluate(blk, 150, 0.1, nb=2)
+            print(f"[dp{world}] step {step}/{cfg.steps} | val {v:.4f} | res {res:.1e} | "
+                  f"{step / (time.time() - t0):.3f} it/s(x{world}B)", flush=True)
+    if rank == 0:
+        print(f"[dp{world}] DONE {cfg.steps} steps in {time.time() - t0:.0f}s", flush=True)
+    dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    main()