27 lines
1.3 KiB
Python
27 lines
1.3 KiB
Python
"""Compare compile-time Swin diagnostic snapshots with the NumPy intermediate values."""
|
|
import argparse
|
|
import numpy as np
|
|
from m_replay_validate import me, sm, LAYERS, metrics
|
|
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--capture-dir', required=True)
|
|
parser.add_argument('--stages', required=True)
|
|
parser.add_argument('--block', default='0,4')
|
|
parser.add_argument('--layer', choices=list(LAYERS), default='enc1')
|
|
args = parser.parse_args()
|
|
p = me.Params(args.capture_dir)
|
|
c, heads, _, cin = LAYERS[args.layer]
|
|
bx, by = map(int, args.block.split(','))
|
|
pre = 4 * c * cin + 8 * c if cin else 0
|
|
x0 = sm.decoder_input(p, bx, by, c, cin, heads, p.p32) if cin else sm.load_tokens(p, bx, by, c // 32)
|
|
stages = {}
|
|
sm.swin_block(p, x0, c, heads, pre, stages=stages)
|
|
actual = np.fromfile(args.stages, np.float16).reshape(4, 64, 128).astype(np.float64)
|
|
for index, name in enumerate(('h1', 'x1', 'h2', 'O')):
|
|
expected = stages[name]
|
|
values = actual[index, :, :expected.shape[1]]
|
|
metrics(expected, values, f'SWIN_STAGE {name}')
|
|
different = np.argwhere(expected != values)
|
|
for row, channel in different[:8]:
|
|
print(f'DIFFERENT stage={name} token={row} channel={channel} reference={expected[row, channel]} native={values[row, channel]}')
|