File size: 3,082 Bytes
3f6e26d | 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 72 73 74 75 76 77 78 79 | """
Unit-test the three NEW CUDA backward primitives against the proven torch vjps
in ttt_manual_backward.py (which match autograd to 1e-12). Localizes any
Plan-A CUDA backward bug to a single kernel.
Run:
TORCHDYNAMO_DISABLE=1 PYTHONPATH="$PWD" CUDA_VISIBLE_DEVICES=5 \
/mnt/afs-h200/yuyangcheng/venvs/fla_triton32/bin/python -m src.models.test_ttt_cuda_primitives
"""
import sys
import torch
from src.models import ttt_cuda
from src.models.ttt_manual_backward import (
silu_p, silu_pp, frobnorm_fwd, frobnorm_bwd, weightnorm_fwd, weightnorm_bwd)
def main():
if not torch.cuda.is_available():
print("SKIP: no CUDA"); return 0
ext = ttt_cuda._load_extension()
if ext is None:
print("FAIL: extension not built"); return 1
dev = "cuda"
g = torch.Generator(device=dev).manual_seed(7)
results = []
for dt, tol in [(torch.float32, 2e-5), (torch.float64, 1e-10)]:
# skip fp64 if kernel doesn't dispatch it (AT_DISPATCH_FLOATING_TYPES covers double)
x = torch.randn(4, 32, 16, generator=g, device=dev, dtype=dt)
# ---- silu_derivs ----
try:
sp, spp = ext.silu_derivs(x)
e1 = (sp - silu_p(x)).abs().max().item()
e2 = (spp - silu_pp(x)).abs().max().item()
ok = e1 < tol and e2 < tol
print(f"[{'OK ' if ok else 'BAD'}] silu_derivs {str(dt):14s} | sp_err={e1:.2e} spp_err={e2:.2e}")
results.append(ok)
except RuntimeError as e:
print(f"[skip] silu_derivs {dt}: {str(e)[:60]}");
# ---- frobnorm_bwd ([B,A,C], Frobenius over (1,2)) ----
xm = torch.randn(4, 16, 8, generator=g, device=dev, dtype=dt)
gy = torch.randn(4, 16, 8, generator=g, device=dev, dtype=dt)
_, cache = frobnorm_fwd(xm)
ref = frobnorm_bwd(gy, cache)
try:
cu = ext.frobnorm_bwd(gy, xm, 1e-7)
e = (ref - cu).abs().max().item()
ok = e < tol
print(f"[{'OK ' if ok else 'BAD'}] frobnorm_bwd {str(dt):14s} | err={e:.2e}")
results.append(ok)
except RuntimeError as e:
print(f"[skip] frobnorm_bwd {dt}: {str(e)[:60]}")
# ---- weightnorm_bwd ([B,A,C], per-column over dim=1) ----
wpre = torch.randn(4, 16, 8, generator=g, device=dev, dtype=dt)
wn = wpre.norm(2, 1, keepdim=True).detach()
gyw = torch.randn(4, 16, 8, generator=g, device=dev, dtype=dt)
_, wc = weightnorm_fwd(wpre, wn)
refw = weightnorm_bwd(gyw, wc)
try:
cuw = ext.weightnorm_bwd(gyw, wpre, wn, 1e-5)
e = (refw - cuw).abs().max().item()
ok = e < tol
print(f"[{'OK ' if ok else 'BAD'}] weightnorm_bwd {str(dt):12s} | err={e:.2e}")
results.append(ok)
except RuntimeError as e:
print(f"[skip] weightnorm_bwd {dt}: {str(e)[:60]}")
allok = all(results) and len(results) > 0
print("\n=== PRIMITIVES", "PASS ===" if allok else "FAIL ===")
return 0 if allok else 1
if __name__ == "__main__":
sys.exit(main())
|