1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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_every', 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_every == 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()
|
