| """Custom fused add+LayerNorm (tt_kernels.add_layernorm) vs torch fp32 and vs the library path |
| (ttnn.add + sharded ttnn.layer_norm + sharded_to_interleaved). usage: addln_test.py [eth|worker] [soak N]""" |
| import sys |
| from pathlib import Path |
| from common import * |
| from open_device_12x10 import open_device |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from hamer import tt_kernels as tk |
| disp = sys.argv[1] if len(sys.argv) > 1 else "eth" |
| soak = int(sys.argv[3]) if len(sys.argv) > 3 and sys.argv[2] == "soak" else 0 |
| timing = "time" in sys.argv |
| GM = int(next((a.split("=")[1] for a in sys.argv if a.startswith("mode=")), "2")) |
| NOY = "noy" in sys.argv |
| dev = open_device(grid="12x10", dispatch=disp, trace_region_size=64 << 20) |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| res = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK, |
| orientation=ttnn.ShardOrientation.ROW_MAJOR) |
| xt = (torch.randn(1, 1, 192, 1280) * 2 + torch.randn(1280) * 1.5).to(torch.bfloat16) |
| yt = (torch.randn(1, 1, 192, 1280) * 0.7).to(torch.bfloat16) |
| gt = 1 + 0.3 * torch.randn(1280); bt = 0.2 * torch.randn(1280) |
| eps = 1e-6 |
| x = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) |
| y = ttnn.from_torch(yt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res) |
| C = tk.addln_constants(gt, bt, dev) |
| if NOY: |
| _, h0 = tk.add_layernorm(x, None, C, eps) |
| ttnn.synchronize_device(dev) |
| h0r = torch.nn.functional.layer_norm(xt.float(), (1280,), gt, bt, eps) |
| hh = ttnn.to_torch(h0).float() |
| print(f"mode={GM} noy first call OK; h PCC {pcc(hh, h0r):.7f} nan={int(hh.isnan().sum())} inf={int(hh.isinf().sum())}", flush=True) |
| if GM == 0: |
| xb = xt.float().reshape(192, 10, 128) |
| mu = xb.sum(-1, keepdim=True) * 10 / 1280 |
| var = ((xb - mu) ** 2).sum(-1, keepdim=True) * 10 / 1280 |
| sr = (((xb - mu) / (var + eps).sqrt()).reshape(1, 1, 192, 1280) * gt + bt) |
| print(f"stub ref PCC {pcc(hh.nan_to_num(), sr):.7f}; per-core-col nan counts", hh.reshape(192, 10, 128).isnan().sum((0, 2)).tolist(), |
| "per row-tile", hh.reshape(6, 32, 1280).isnan().sum((1, 2)).tolist()) |
| print("sample", hh[0, 0, 0, :6].tolist(), sr[0, 0, 0, :6].tolist()) |
| z = ((hh - bt.to(torch.bfloat16).float()) / gt.to(torch.bfloat16).float()).reshape(192, 10, 128).double() |
| xx = xb.double() |
| for r in (0, 1, 31, 40): |
| for c in (0, 5): |
| A = torch.stack([xx[r, c], torch.ones(128, dtype=torch.float64)], 1) |
| sol = torch.linalg.lstsq(A, z[r, c].unsqueeze(1)).solution.squeeze() |
| rs, mu_dev = float(sol[0]), float(-sol[1] / sol[0]) |
| print(f"row {r} col {c}: dev rstd {rs:.4g} mean {mu_dev:.4g} | expect rstd {float((var[r, c] + eps).rsqrt()):.4g} mean {float(mu[r, c]):.4g} | var {float(var[r,c]):.4g}") |
| ttnn.close_device(dev); sys.exit(0) |
| s, h = tk.add_layernorm(x, y, C, eps) |
| ttnn.synchronize_device(dev) |
| if "wshard" in sys.argv: |
| G = dev.compute_with_storage_grid_size() |
| sg = ttnn.num_cores_to_corerangeset(8, ttnn.CoreCoord(G.x, G.y), row_wise=True) |
| wmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, |
| ttnn.ShardSpec(sg, [192, 160], ttnn.ShardOrientation.ROW_MAJOR)) |
| s3, h3 = tk.add_layernorm(x, y, C, eps, out_memory_config=wmc) |
| ttnn.synchronize_device(dev) |
| print("wshard h == interleaved h:", bool(torch.equal(ttnn.to_torch(h3).float(), ttnn.to_torch(h).float())), |
| "s equal:", bool(torch.equal(ttnn.to_torch(s3).float(), ttnn.to_torch(s).float())), flush=True) |
| if "time" in sys.argv: |
| print(f"wshard addln {time_op(dev, lambda: list(tk.add_layernorm(x, y, C, eps, out_memory_config=wmc))):.2f} us", flush=True) |
| print(f"mode={GM} first call OK", flush=True) |
| st = ttnn.to_torch(s).float(); ht = ttnn.to_torch(h).float() |
| sref = (xt.float() + yt.float()).to(torch.bfloat16).float() |
| href = torch.nn.functional.layer_norm(sref, (1280,), gt, bt, eps) |
| print("s exact:", bool(torch.equal(st, sref)), "max|ds|", float((st - sref).abs().max())) |
| print(f"h PCC {pcc(ht, href):.7f} max|dh| {float((ht - href).abs().max()):.4f} mean|dh| {float((ht - href).abs().mean()):.5f}") |
| |
| ga = ttnn.from_torch(gt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| be = ttnn.from_torch(bt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev) |
| lnpc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(10, 6), subblock_w=4, block_h=1, block_w=4, inplace=False) |
| def lib(): |
| s2 = ttnn.add(x, y, memory_config=res) |
| hh = ttnn.layer_norm(s2, weight=ga, bias=be, epsilon=eps, memory_config=res, program_config=lnpc) |
| hi = ttnn.sharded_to_interleaved(hh, L1) |
| ttnn.deallocate(hh) |
| return [s2, hi] |
| s2, hl = lib() |
| s2t = ttnn.to_torch(s2).float() |
| print("lib s exact:", bool(torch.equal(s2t, sref)), "max|ds|", float((s2t - sref).abs().max()), "custom==lib s:", bool(torch.equal(st, s2t)), |
| "n diff", int((st != sref).sum())) |
| hl = ttnn.to_torch(hl).float() |
| href2 = torch.nn.functional.layer_norm(sref, (1280,), gt.to(torch.bfloat16).float(), bt.to(torch.bfloat16).float(), eps) |
| print(f"lib h PCC {pcc(hl, href):.7f} max|dh| {float((hl - href).abs().max()):.4f} mean|dh| {float((hl - href).abs().mean()):.5f}") |
| print(f"custom vs bf16-gamma ref PCC {pcc(ht, href2):.7f} mean|dh| {float((ht - href2).abs().mean()):.5f}; lib {float((hl - href2).abs().mean()):.5f}") |
| |
| _, h0 = tk.add_layernorm(x, None, C, eps) |
| h0r = torch.nn.functional.layer_norm(xt.float(), (1280,), gt, bt, eps) |
| print(f"noy h PCC {pcc(ttnn.to_torch(h0).float(), h0r):.7f}") |
| if timing: |
| print(f"custom addln {time_op(dev, lambda: list(tk.add_layernorm(x, y, C, eps))):.2f} us; library add+ln+s2i {time_op(dev, lib):.2f} us", flush=True) |
| if soak: |
| tid = ttnn.begin_trace_capture(dev, cq_id=0) |
| try: |
| outs = [tk.add_layernorm(x, y, C, eps) for _ in range(10)] |
| finally: |
| ttnn.end_trace_capture(dev, tid, cq_id=0) |
| for i in range(soak // 10): |
| ttnn.execute_trace(dev, tid, cq_id=0, blocking=False) |
| if i % 50 == 0: |
| ttnn.synchronize_device(dev) |
| ttnn.synchronize_device(dev) |
| bad = sum(int(not torch.equal(ttnn.to_torch(o[1]).float(), ht)) for o in outs) |
| print(f"soak {soak} replays OK, outputs differing from first call: {bad}/10") |
| ttnn.release_trace(dev, tid) |
| ttnn.close_device(dev) |
|
|
