File size: 3,941 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 71 72 73 74 | """Binary-clock memory probe through Transformers' BeamAttention implementation.
This is a small constructed layer, not pretrained Beam. A host provides write
positions as bits. The eager Beam backend materializes attention scores in the
layer's dtype, unlike the separate SDPA probe's implementation.
"""
import math
import torch
from transformers import BeamConfig
from transformers.models.beam.modeling_beam import BeamAttention, compute_attention_scaling
from .attention_cpu import bits, from_bits
class HFBinaryMemory:
def __init__(self, max_writes=128, dtype=torch.float32):
if not 1 <= max_writes <= 2047: raise ValueError('max_writes must be 1..2047')
self.maximum, self.dtype, self.writes = max_writes, dtype, 0
self.width = max_writes.bit_length()
self.clock = [2**i/(2*(max_writes+1)) for i in range(self.width)]
config = BeamConfig(hidden_size=128, num_hidden_layers=1, num_attention_heads=6,
num_key_value_heads=1, head_dim=128, layer_types=['full_attention'],
mlp_layer_types=['dense'], attention_dropout=0.)
config._attn_implementation = 'eager'
self.layer = BeamAttention(config, 0).eval().to(dtype).requires_grad_(False)
key_start = 6
value_start = key_start + 5 + self.width
with torch.no_grad():
for m in (self.layer.q_proj, self.layer.k_proj, self.layer.v_proj,
self.layer.o_proj, self.layer.gate_proj): m.weight.zero_()
for head in range(6):
off = head*128
for i in range(4): self.layer.q_proj.weight[off+i, i] = 1
for i, coefficient in enumerate(self.clock): self.layer.q_proj.weight[off+4+i, 4] = coefficient
self.layer.q_proj.weight[off+4+self.width, 4] = 2
for i in range(5+self.width): self.layer.k_proj.weight[i, key_start+i] = 1
for i in range(4):
self.layer.v_proj.weight[i, value_start+i] = 1
self.layer.o_proj.weight[i, i] = 2
self.layer.gate_bias.zero_()
q2 = 8 + sum(x*x for x in self.clock)
k2 = 5 + self.width
gain = 16*(max_writes+1)*math.sqrt(q2*k2)/math.sqrt(128)
self.layer.q_norm.weight.fill_(gain)
self.layer.k_norm.weight.fill_(1)
self.rows = [self._row(a, 0, 0, True) for a in range(16)]
def _row(self, address, value, position, valid):
row = torch.zeros(128, dtype=self.dtype)
row[:4] = torch.tensor([2*b-1 for b in bits(address, 4)], dtype=self.dtype)
row[4] = 1
key = [*(2*b-1 for b in bits(address, 4)),
*(2*b-1 for b in bits(position, self.width)), 1 if valid else -1]
row[6:6+len(key)] = torch.tensor(key, dtype=self.dtype)
value_start = 6 + 5 + self.width
row[value_start:value_start+4] = torch.tensor(bits(value, 4), dtype=self.dtype)
return row
def write(self, address, value):
if not 0 <= address < 16 or not 0 <= value < 16: raise ValueError('expected nibbles')
if self.writes >= self.maximum: raise ValueError('write budget exhausted')
self.writes += 1
self.rows.append(self._row(address, value, self.writes, True))
def read(self, address):
if not 0 <= address < 16: raise ValueError('expected nibble address')
inputs = torch.stack(self.rows + [self._row(address, 0, self.writes, False)]).unsqueeze(0)
length = inputs.shape[1]
mask = torch.full((1, 1, length, length), float('-inf'), dtype=self.dtype).triu(1)
positions = torch.arange(length).unsqueeze(0)
with torch.inference_mode():
out, _ = self.layer(inputs, position_embeddings=None, attention_mask=mask,
attention_scaling=compute_attention_scaling(positions))
return from_bits((out[0, -1, :4] > .5).int().tolist())
|