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 | from common import *
dev = open_dev()
torch.manual_seed(0)
x_t = torch.randn(1, 1, 192, 1280); g_t = 1 + .1 * torch.randn(1280); b_t = .1 * torch.randn(1280)
ref = torch.nn.functional.layer_norm(x_t, (1280,), g_t, b_t, 1e-6)
gw = ttnn.from_torch(g_t.reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev)
bw = ttnn.from_torch(b_t.reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev)
for gx, gy in [(10, 6), (8, 6), (5, 6), (4, 6), (2, 6), (10, 3), (10, 2), (5, 3)]:
mc = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=gy, x=gx), strategy=ttnn.ShardStrategy.BLOCK, orientation=ttnn.ShardOrientation.ROW_MAJOR)
xs = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=mc)
bwd = 40 // gx; bh = 6 // gy
for sbw in sorted({d for d in range(1, min(bwd, 4) + 1) if bwd % d == 0}):
pc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(gx, gy), subblock_w=sbw, block_h=bh, block_w=bwd, inplace=False)
try:
o = ttnn.layer_norm(xs, weight=gw, bias=bw, epsilon=1e-6, memory_config=mc, program_config=pc)
p = pcc(ttnn.to_torch(o), ref)
us = time_op(dev, lambda pc=pc, xs=xs, mc=mc: ttnn.layer_norm(xs, weight=gw, bias=bw, epsilon=1e-6, memory_config=mc, program_config=pc))
print(f"LN block {gx}x{gy} sbw{sbw}: {us:.1f} us pcc={p:.5f}")
except Exception as e:
print(f"LN block {gx}x{gy} sbw{sbw}: FAIL {str(e).splitlines()[0][:120]}")
ttnn.close_device(dev)
|