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
|
"""Cross-env gradient fingerprint: three gradient paths on the SAME ckpt (s2000) + SAME seeded batch.
Run under torch 2.10 (local A6000) and torch 2.3.1 (timan107 GTX1080); diverging fingerprints pinpoint
which path is numerically broken on the pascal env (suspects: double-backward reg paths).
(a) plain CE backward through the feedforward block (bp-style backward)
(b) resreg path: tforce graph at z_T1 + autograd.grad (single backward through force graph)
(c) jacreg path: autograd.functional.jvp(create_graph) + backward (double-backward)
Prints total grad-norm per path + first-3 per-tensor norms."""
import torch, torch.nn.functional as F
import lt_ep_train as L
torch.manual_seed(0)
blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
ck = torch.load('runs/redx_traj/s2000.pt', map_location=L.dev)
with torch.no_grad():
for p, w in zip(blk.allp, ck['allp']):
p.copy_(w.to(L.dev))
torch.manual_seed(42)
idx, y = L.get_batch('train', 8, 256)
xin = blk.embed(idx).detach()
zs = L.relax(blk, xin.clone(), xin, 150, 0.1)
def report(tag, params, grads):
tot = sum(float(g.pow(2).sum()) for g in grads if g is not None) ** 0.5
firsts = " ".join(f"{float(g.norm()):.3e}" if g is not None else "None" for g in grads[:3])
fin = all(torch.isfinite(g).all() for g in grads if g is not None)
print(f"{tag}: total={tot:.4e} finite={fin} first3=[{firsts}]", flush=True)
# (a) plain BP backward
x = blk.embed(idx)
h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
loss = F.cross_entropy((h @ blk.Wh).reshape(-1, L.vocab), y.reshape(-1))
ga = torch.autograd.grad(loss, blk.allp, allow_unused=True)
report("(a) plain-BP ", blk.allp, list(ga))
# (b) resreg path
with torch.enable_grad():
Fz = blk.tforce(zs, xin)
Rr = (0.1 * Fz).pow(2).sum() / (zs.pow(2).sum() + 1e-9)
gb = torch.autograd.grad(Rr, blk.block, allow_unused=True)
report("(b) resreg ", blk.block, list(gb))
# (c) jacreg double-backward path
er = torch.randn_like(zs)
with torch.enable_grad():
Jv = torch.autograd.functional.jvp(blk.nc_force, zs.detach(), er, create_graph=True)[1]
R = 0.1 * (Jv ** 2).sum() / (er ** 2).sum()
gc = torch.autograd.grad(R, blk.block, allow_unused=True)
report("(c) jacreg dbb", blk.block, list(gc))
print("GRAD_ENV_PROBE_DONE", flush=True)
|
