File size: 6,223 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 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Round 8: cmm + fused add+LN epilogue (tt_kernels.cmm_linear(addln=...)) vs the separate cmm and add_layernorm ops.
usage: cmm_addln_test.py shape=tiny|proj|fc2 [ws] [time] (ws: LN output width-sharded on 10 cores like qkv's input)
Expect bit-identical s and h."""
import sys
from pathlib import Path
from common import *
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
arg = lambda k, d: next((a.split("=", 1)[1] for a in sys.argv if a.startswith(k + "=")), d)
shape = arg("shape", "tiny")
GX, GY, NT, KT, kb = {"tiny": (2, 2, 1, 8, 4), "proj": (10, 6, 4, 48, 8), "fc2": (10, 6, 4, 160, 16)}[shape]
dev = open_dev()
try:
torch.manual_seed(0)
M, K, N = 32 * GY, 32 * KT, 32 * NT * GX
G = dev.compute_with_storage_grid_size()
res = ttnn.create_sharded_memory_config((M, N), core_grid=ttnn.CoreGrid(y=GY, x=GX), strategy=ttnn.ShardStrategy.BLOCK,
orientation=ttnn.ShardOrientation.ROW_MAJOR)
at = (torch.randn(1, 1, M, K) * 0.5).to(torch.bfloat16)
a = ttnn.from_torch(at, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.L1_MEMORY_CONFIG)
wt = (torch.randn(K, N) * K ** -0.5).to(torch.bfloat16)
wf = tk.CmmWeight(wt, dev, GX, NT, kb, ttnn.bfloat8_b)
b = ttnn.from_torch((torch.randn(N) * 0.1).reshape(1, 1, 1, N), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev)
xt = torch.randn(1, 1, M, N).to(torch.bfloat16)
x1 = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
x2 = ttnn.from_torch(xt, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
C = tk.addln_constants(torch.randn(N) * 0.2 + 1, torch.randn(N) * 0.1, dev)
Cf = C
if "noaff" in sys.argv: # round 9: reference = identity affine (gamma 1, beta 0); fused = ADDLN_NOAFF (both exact)
C = tk.addln_constants(torch.ones(N), torch.zeros(N), dev)
Cf = dict(tk.addln_constants(torch.ones(N), torch.zeros(N), dev), noaff=True)
if "nobeta" in sys.argv: # reference = gamma with beta 0; fused = ADDLN_NOBETA (beta add skipped)
gm = torch.randn(N) * 0.2 + 1
C = tk.addln_constants(gm, torch.zeros(N), dev)
Cf = dict(tk.addln_constants(gm, torch.zeros(N), dev), nobeta=True)
hmem = None
if "ws" in sys.argv:
ns = GX
sg = ttnn.num_cores_to_corerangeset(ns, G, row_wise=True)
hmem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1,
ttnn.ShardSpec(sg, [M, N // ns], ttnn.ShardOrientation.ROW_MAJOR))
if "xres" in sys.argv: # round 12: LN output straight into the resident in0 of the xmm fc1 (12 x 10 cores)
hmem = "xres"
defs = ("CMM_FLUSH", "CMM_LA=1")
d = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True)
s_ref, h_ref = tk.add_layernorm(x1, d, C, 1e-6, out_memory_config=None if hmem == "xres" else hmem)
print(f"[ca] {shape} G={GX}x{GY} NT={NT} KT={KT} kb={kb} ws={hmem is not None}", flush=True)
s, h = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True,
addln=dict(x=x2, consts=Cf, eps=1e-6, out_mem=hmem))
if "noaff" in sys.argv or "nobeta" in sys.argv: # standalone add+LN with NOAFF / NOBETA too
s3, h3 = tk.add_layernorm(x1, tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True), Cf, 1e-6,
out_memory_config=None if hmem == "xres" else hmem)
ttnn.synchronize_device(dev)
if hmem != "xres": print(f"[ca] standalone NOAFF h equal {bool(torch.equal(ttnn.to_torch(h_ref).float(), ttnn.to_torch(h3).float()))}", flush=True)
ttnn.synchronize_device(dev)
sr, hr = ttnn.to_torch(s_ref).float(), ttnn.to_torch(h_ref).float()
if hmem == "xres":
# rebuild [M, N] from the resident shards: core (x, y), parity p, shard tile (j, r, k) = row 2 i + r, K tile (2 j + p) 4 + k
rs = [ttnn.to_torch(t).float().reshape(10, 12, 5, 2, 4, 32, 32) for t in h]
hf = torch.full((M, N), float("nan"))
bad = 0
for y_ in range(10):
for x_ in range(12):
_, i_ = tk.xmm_place(x_, y_)
for p_ in range(2):
for j_ in range(5):
for r_ in range(2):
for k_ in range(4):
m_, kk = 2 * i_ + r_, (2 * j_ + p_) * 4 + k_
t_ = rs[p_][y_, x_, j_, r_, k_]
blk = hf[32 * m_:32 * m_ + 32, 32 * kk:32 * kk + 32]
if not torch.isnan(blk).any() and not torch.equal(blk, t_):
bad += 1
blk.copy_(t_)
print(f"[ca] xres: copies disagreeing {bad}", flush=True)
hf = hf.reshape(1, 1, M, N)
sf = ttnn.to_torch(s).float()
else:
sf, hf = ttnn.to_torch(s).float(), ttnn.to_torch(h).float()
print(f"[ca] first call OK s equal {bool(torch.equal(sr, sf))} max|ds| {float((sr - sf).abs().max()):.3e} "
f"h equal {bool(torch.equal(hr, hf))} max|dh| {float((hr - hf).abs().max()):.3e} h PCC {pcc(hf, hr):.7f}", flush=True)
if "noaff" in sys.argv: # vs torch LN (fp64) of the device residual s
sd = sf.double()
ht = torch.nn.functional.layer_norm(sd, (sd.shape[-1],), eps=1e-6)
print(f"[ca] h vs torch LN(s): max|d| {float((hf.double() - ht).abs().max()):.3e} PCC {pcc(hf, ht.float()):.7f}", flush=True)
if "time" in sys.argv:
def sep():
d_ = tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True)
r = tk.add_layernorm(x1, d_, C, 1e-6, out_memory_config=hmem)
ttnn.deallocate(d_)
return r
fus = lambda: tk.cmm_linear(a, wf, b, res, kb=kb, slots=3, defines=defs, v1=True,
addln=dict(x=x2, consts=Cf, eps=1e-6, out_mem=hmem))
print(f"[ca] separate {time_op(dev, sep, reps=5, iters=40):.2f} us fused {time_op(dev, fus, reps=5, iters=40):.2f} us", flush=True)
finally:
ttnn.close_device(dev)
|