| """Backward reproducibility: repeated runs must be bitwise identical.""" |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT)) |
| import load_local as mamba3 |
|
|
| dev = "cuda" |
| NAMES = ["q", "k", "v", "z", "q_bias", "k_bias", "mimo_v", "mimo_o", "mimo_z", "D"] |
| fails = [] |
|
|
| for S, H, G, P, N, R, C in [(512, 16, 1, 64, 128, 4, 16), (256, 8, 2, 32, 64, 2, 32)]: |
| B, Na = 1, N // 2 |
| torch.manual_seed(0) |
| f = lambda *s: torch.randn(*s, device=dev) |
| shp = {"q": (B, S, R, G, N), "k": (B, S, R, G, N), "v": (B, S, H, P), |
| "z": (B, S, H, P), "q_bias": (H, R, N), "k_bias": (H, R, N), |
| "mimo_v": (H, R, P), "mimo_o": (H, R, P), "mimo_z": (H, R, P), "D": (H,)} |
| base = {n: f(*shp[n]) for n in NAMES} |
| dtv = F.softplus(-3.0 + f(B, H, S)) |
| adtv = -F.softplus(f(B, H, S)).clamp(max=-1e-4) * dtv |
| trv = torch.rand(B, H, S, device=dev) * 0.5 |
| thv = torch.rand(B, S, H, Na, device=dev) * 0.4 |
| dy = torch.randn(B, S, H, P, device=dev) |
|
|
| def once(): |
| t = {n: base[n].detach().clone().requires_grad_(True) for n in NAMES} |
| sch = {n: x.clone().requires_grad_(True) for n, x in |
| {"dt": dtv, "adt": adtv, "trap": trv, "th": thv}.items()} |
| y = mamba3.forward(t["q"], t["k"], t["v"], t["q_bias"], t["k_bias"], t["mimo_v"], |
| t["mimo_o"], sch["th"], sch["adt"], sch["dt"], sch["trap"], |
| z=t["z"], mimo_z=t["mimo_z"], D=t["D"], chunk_size=C) |
| y.backward(dy) |
| g = {n: t[n].grad.clone() for n in NAMES} |
| g.update({n: sch[n].grad.clone() for n in sch}) |
| return y.detach().clone(), g |
|
|
| y1, g1 = once() |
| print(f"S={S} H={H} G={G} P={P} N={N} R={R} C={C}") |
| for trial in range(3): |
| y2, g2 = once() |
| if not torch.equal(y1, y2): |
| fails.append(f"forward@S{S}") |
| for n in g1: |
| if not torch.equal(g1[n], g2[n]): |
| fails.append(f"d{n}@S{S}") |
| bad = sorted(set(x for x in fails if f"@S{S}" in x)) |
| print(f" forward + 14 gradients over 3 repeats: " |
| f"{'all bitwise identical' if not bad else 'DIFFER: ' + ', '.join(bad)}") |
|
|
| print("\nPASS" if not fails else "\nFAIL: " + ", ".join(sorted(set(fails)))) |
| sys.exit(0 if not fails else 1) |
|
|
