File size: 1,604 Bytes
58ba6df | 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 | """custom GELU (tt_kernels.gelu) soak: one eager call, then N trace replays of 50 calls on the fc1-output shape;
every output must stay bit-identical. usage: gelu_soak.py N"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels
n = int(sys.argv[1]) if len(sys.argv) > 1 else 1000
dev = open_dev()
try:
L1 = ttnn.L1_MEMORY_CONFIG
torch.manual_seed(0)
xt = torch.randn(1, 1, 192, 5120) * 3.0
x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
o0 = ttnn.to_torch(tt_kernels.gelu(x, memory_config=L1)).float()
ref = torch.nn.functional.gelu(xt.bfloat16().float())
print(f"[gsoak] first call OK pcc {pcc(o0, ref):.8f} max|err| {float((o0 - ref).abs().max()):.3e}", flush=True)
reps = 50
tid = ttnn.begin_trace_capture(dev, cq_id=0)
outs = []
try:
for _ in range(reps):
outs.append(tt_kernels.gelu(x, memory_config=L1))
finally:
ttnn.end_trace_capture(dev, tid, cq_id=0)
try:
for i in range(max(1, n // reps)):
ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
if i % 40 == 39:
ttnn.synchronize_device(dev)
ttnn.synchronize_device(dev)
bad = sum(int(not torch.equal(ttnn.to_torch(t).float(), o0)) for t in outs)
print(f"[gsoak] {max(1, n // reps) * reps} calls OK, outputs differing from first: {bad}/{reps}", flush=True)
finally:
ttnn.release_trace(dev, tid)
finally:
ttnn.close_device(dev)
|