File size: 4,018 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 75 76 77 78 79 80 81 82 | """Attention-backed nibble memory through a stock BeamDecoderLayer.
This probe uses constructed weights, a compact one-group Beam config and
host-encoded event vectors. Only memory reads are implemented in this decoder;
the layer's MLP is zeroed. It does not use pretrained Beam weights.
"""
import math
import torch
from transformers import BeamConfig
from transformers.cache_utils import DynamicCache
from transformers.models.beam.modeling_beam import BeamDecoderLayer, compute_attention_scaling
from .attention_cpu import bits
class DecoderMemory:
def __init__(self, max_writes=128):
if not 1 <= max_writes <= 1024: raise ValueError('max_writes must be 1..1024')
self.maximum = max_writes
self.writes = 0
config = BeamConfig(hidden_size=128, intermediate_size=16, 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 = BeamDecoderLayer(config, 0).eval().requires_grad_(False)
attn = self.layer.self_attn
with torch.no_grad():
for proj in (attn.q_proj, attn.k_proj, attn.v_proj, attn.o_proj, attn.gate_proj,
self.layer.mlp.gate_proj, self.layer.mlp.up_proj, self.layer.mlp.down_proj):
proj.weight.zero_()
for head in range(6):
offset = 128*head
for i in range(5): attn.q_proj.weight[offset+i, i] = 1
attn.q_proj.weight[offset+5, 5] = 5
for i in range(7): attn.k_proj.weight[i, 6+i] = 1
for i in range(16):
attn.v_proj.weight[i, 13+i] = 1
attn.o_proj.weight[32+i, i] = 1
attn.gate_bias.zero_()
attn.q_norm.weight.fill_(16*max_writes*math.sqrt(30*6)/math.sqrt(128))
attn.k_norm.weight.fill_(1)
self.cache = DynamicCache(config=config)
# Prefill zero-initialized RAM; earlier events cannot inspect future ones.
rows = torch.stack([self._row(a, 0, 0, True) for a in range(16)]).unsqueeze(0)
with torch.inference_mode(): self._forward(rows)
def _row(self, address, value, position, valid):
x = torch.zeros(128)
x[:4] = torch.tensor([2*b-1 for b in bits(address, 4)])
x[4:6] = 1
x[6:10] = x[:4]
p = position / self.maximum
x[10], x[11], x[12] = p, 1 if valid else -1, math.sqrt(max(0, 1-p*p))
x[13+value] = 1
# The input layer's RMS normalization must not change V magnitudes.
x[-1] = math.sqrt(128-float(x.square().sum()))
return x
def _forward(self, rows):
old = self.cache.get_seq_length()
length = rows.shape[1]
mask = None
if length > 1:
mask = rows.new_zeros((1, 1, length, old+length))
mask[0, 0, :, old:].masked_fill_(torch.ones(length, length, dtype=torch.bool).triu(1), float('-inf'))
positions = torch.arange(old, old+length).unsqueeze(0)
return self.layer(rows, position_embeddings=None, attention_mask=mask,
attention_scaling=compute_attention_scaling(positions),
past_key_values=self.cache, use_cache=True)
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
with torch.inference_mode(): self._forward(self._row(address, value, self.writes, True).reshape(1, 1, -1))
def read(self, address):
if not 0 <= address < 16: raise ValueError('expected nibble')
with torch.inference_mode():
output = self._forward(self._row(address, 0, self.writes, False).reshape(1, 1, -1))
return int(output[0, 0, 32:48].argmax())
|