File size: 4,069 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 | """Experimental CPU transition inside a stock BeamDecoderLayer.
Both RAM retrieval and ROM-specific transitions execute in stock Beam decoder
attention and SwiGLU MLP weights. The host still selects the ROM operand/address
from the current PC and serializes discrete state and write events. This is not
an autoregressive LM, a trained checkpoint, or an embedded 501B Beam model.
"""
import math
import torch
from transformers.models.beam.modeling_beam import BeamMLP
from .attention_cpu import bits, from_bits
from .hf_decoder_memory import DecoderMemory
from .large_cpu import validate, transition, EventResult
class DecoderCPU(DecoderMemory):
def __init__(self, program, max_writes=128):
validate(program)
self.program = list(program)
super().__init__(max_writes)
# Each neuron recognizes one of 256*16*16 (PC, ACC, RAM) states.
# This is a large, exponentially scaling truth table, not a practical CPU.
n = 65536
mlp = BeamMLP(self.layer.self_attn.config, intermediate_size=n)
ids = torch.arange(n, dtype=torch.int64)
pc = ids & 255
acc = (ids >> 8) & 15
mem = (ids >> 12) & 15
state = (((ids[:, None] >> torch.arange(12)) & 1)*2-1).float()
self.layer.mlp = mlp.eval().requires_grad_(False)
# After the attention residual, a bipolar input bit has magnitude 1/sqrt(2)
# and the one-hot memory response has magnitude 8. Layer norms preserve
# these ratios because the two subspaces are orthogonal and fixed-norm.
alpha = 1 / math.sqrt(2)
temperature = 32.0
transitions = []
for i in range(n):
npc, nacc, write, output, halt = transition(program, int(pc[i]), int(acc[i]), int(mem[i]))
transitions.append(bits(npc, 8) + bits(nacc, 4) + [write, output, halt])
transitions = torch.tensor(transitions, dtype=torch.int64)
with torch.no_grad():
mlp.gate_proj.weight.zero_(); mlp.up_proj.weight.zero_(); mlp.down_proj.weight.zero_()
mlp.gate_proj.weight[:, 48:60] = state * (temperature / alpha)
mlp.gate_proj.weight[ids, 32 + mem] = temperature / 8
mlp.gate_proj.weight[:, 5] = -12.5 * temperature / alpha
mlp.up_proj.weight[:, 5] = 1 / alpha
# Two output channels per bit; exactly one is active per transition.
for i in range(15):
mlp.down_proj.weight[60 + 2*i, ids] = (transitions[:, i] == 0).float()
mlp.down_proj.weight[61 + 2*i, ids] = (transitions[:, i] == 1).float()
self.pc = 0
self.acc = 0
self.halted = False
def _cpu_row(self, address, pc, acc):
x = self._row(address, 0, self.writes, False)
state = bits(pc, 8) + bits(acc, 4)
x[48:60] = torch.tensor([2*b-1 for b in state], dtype=torch.float32)
x[-1] = 0
x[-1] = math.sqrt(128-float(x.square().sum()))
return x
def device_write(self, address, value):
self.write(address, value)
def step(self):
if self.halted: raise RuntimeError('machine is halted')
# The address decoder still runs on the host: the model's transition
# is compiled, but this protocol is not yet a self-contained decoder.
op, operand = self.program[self.pc] if self.pc < len(self.program) else ('HALT', 0)
address = operand if op in {'LOAD', 'STORE', 'ADDM', 'SUBM'} else 0
old_acc = self.acc
with torch.inference_mode():
output = self._forward(self._cpu_row(address, self.pc, self.acc).reshape(1, 1, -1))[0, 0]
pairs = output[60:90].reshape(15, 2)
if not torch.isfinite(pairs).all(): raise RuntimeError('non-finite output')
decoded = pairs.argmax(dim=-1).tolist()
self.pc, self.acc = from_bits(decoded[:8]), from_bits(decoded[8:12])
write, out, self.halted = map(bool, decoded[12:])
if write: self.write(address, old_acc)
return EventResult(self.pc, self.acc, write, out, self.halted, address, old_acc)
|