diff options
| author | Yuren Hao <yurenh2@illinois.edu> | 2026-07-06 09:42:52 -0500 |
|---|---|---|
| committer | Yuren Hao <yurenh2@illinois.edu> | 2026-07-06 09:42:52 -0500 |
| commit | 40be67d4f5b5a6b46c662c70b759e585e136d70e (patch) | |
| tree | bbe02d5164c9b86a0af6a3634fb7556f065c8672 /ep_run/dp_ep.py | |
| parent | 9a8b2796ca12e4d4c24717485a635a301aa6d07f (diff) | |
Tier-3 gates: Anderson math-yes/impl-no (parked for v2); bf16polish UNSAFE near-edge (eval-only); dp_ep.py ready
Anderson: res 25-35x deeper per budget but 5.8x slower (naive history stacks
+ per-iter safeguard eval) — v2 = ring buffers + periodic safeguard, est +1.3x
on the speed tier. bf16+20polish: 1.41x free phase, res parity, BUT z-diff
1.2e-3 — near-marginal operators contract too slowly for a 20-step polish
(0.998^20≈0.96), same magnitude as the TF32 kill verdict and the estimator's
50%-sensitivity input. Predicted by our own depth/noise theory. Flags kept
with warnings; neither ships for training. dp_ep.py: manual-allreduce EP DP
(controller in lockstep, aligned collectives), smoke pending freed 1080s.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/dp_ep.py')
| -rw-r--r-- | ep_run/dp_ep.py | 71 |
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() |
