File size: 992 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 | """Reproduce bounded Beam-style attention precision sweep on CPU."""
import argparse
import random
import torch
from beamboot.beam_probe import BeamAttentionProbe, calibrated_gain
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--horizons', nargs='+', type=int, default=[128, 512, 1024])
parser.add_argument('--seed', type=int, default=417)
args = parser.parse_args()
for dtype in (torch.float32, torch.float16, torch.bfloat16):
for horizon in args.horizons:
rng = random.Random(args.seed)
memory = BeamAttentionProbe(horizon, dtype, calibrated_gain(horizon))
errors = 0
for _ in range(horizon):
address, value = rng.randrange(16), rng.randrange(16)
memory.write(address, value)
errors += memory.read(address) != value
print(f'{dtype} horizon={horizon} immediate-read-errors={errors}', flush=True)
if __name__ == '__main__': main()
|