File size: 1,414 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 | import random
import unittest
import torch
from beamboot.binary_time_probe import BinaryTimeProbe
from beamboot.hf_binary_memory import HFBinaryMemory
from beamboot.causal_lm import CausalMachine
class BinaryPrecisionTests(unittest.TestCase):
def test_bf16_binary_clock_and_retention(self):
memory = BinaryTimeProbe(512, torch.bfloat16)
rng = random.Random(417)
reference = [0]*16
for i in range(512):
address, value = rng.randrange(16), rng.randrange(16)
memory.write(address, value)
reference[address] = value
self.assertEqual(memory.read(address), value, f'write {i+1}')
for address, value in enumerate(reference):
self.assertEqual(memory.read(address), value)
def test_hf_beam_attention_sdpa_bf16(self):
memory = HFBinaryMemory(64, torch.bfloat16)
memory.layer.config._attn_implementation = 'sdpa'
rng = random.Random(417)
for _ in range(64):
address, value = rng.randrange(16), rng.randrange(16)
memory.write(address, value)
self.assertEqual(memory.read(address), value)
def test_known_unsafe_runtime_is_rejected(self):
with self.assertRaisesRegex(ValueError, 'known incorrect'):
CausalMachine._validate_precision(torch.bfloat16, 'eager')
CausalMachine._validate_precision(torch.bfloat16, 'sdpa')
|