File size: 3,005 Bytes
324ecd2 | 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 | """Run a small, deterministic backend/precision audit on CPU.
This is an engineering smoke test, not a statistically general benchmark.
"""
import argparse
import json
import platform
import random
from pathlib import Path
import torch
import transformers
from beamboot.hf_binary_memory import HFBinaryMemory
from beamboot.causal_lm import CausalMachine
from beamboot.large_cpu import LargeCPU, LargeRuntime
def attention_case(dtype, backend, seed=417, horizon=128):
m = HFBinaryMemory(horizon, dtype)
m.layer.config._attn_implementation = backend
rng = random.Random(seed)
failures = []
for i in range(horizon):
address, value = rng.randrange(16), rng.randrange(16)
m.write(address, value)
actual = m.read(address)
if actual != value:
failures.append({'write': i+1, 'address': address, 'expected': value, 'actual': actual})
return {'dtype': str(dtype), 'attention': backend, 'horizon': horizon,
'seed': seed, 'errors': len(failures), 'failures': failures}
def machine_case(dtype, backend):
program = [('LDI', 7), ('STORE', 3), ('LDI', 2), ('ADDM', 3), ('OUT', 0), ('HALT', 0)]
# Instantiate in safe fp32, then deliberately override for the negative
# control. The public runtime rejects BF16+eager because it is known bad.
machine = CausalMachine(program)
machine.model.to(dtype)
machine.model.config._attn_implementation = backend
machine._reset(128)
ref = LargeRuntime(LargeCPU(program))
failures = []
for i in range(len(program)):
actual, expected = machine.step(), ref.step()
if actual != expected:
failures.append({'step': i+1, 'expected_acc': expected.acc, 'actual_acc': actual.acc})
if actual.halt: break
return {'dtype': str(dtype), 'attention': backend, 'steps': i+1,
'mismatches': len(failures), 'failures': failures}
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--output', default='results/precision_smoke.json')
args = parser.parse_args()
cases = [(torch.float32, 'eager'), (torch.bfloat16, 'eager'), (torch.bfloat16, 'sdpa')]
results = {'scope': 'CPU smoke test; constructed weights, not pretrained Beam; one deterministic seed',
'python': platform.python_version(), 'torch': torch.__version__, 'transformers': transformers.__version__,
'attention_read': [], 'machine': []}
for dtype, backend in cases:
results['attention_read'].append(attention_case(dtype, backend))
print('attention', dtype, backend, results['attention_read'][-1]['errors'], flush=True)
results['machine'].append(machine_case(dtype, backend))
print('machine', dtype, backend, results['machine'][-1]['mismatches'], flush=True)
path = Path(args.output)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(results, indent=2)+'\n')
print('saved', path, flush=True)
if __name__ == '__main__': main()
|