| """Two stock Beam decoder layers executing a compiled CPU transition. |
| |
| Layer 0's SwiGLU MLP decodes the current PC into a memory address. Layer 1's |
| attention reads RAM, and its SwiGLU MLP computes the next CPU state. The host |
| serializes state/write events but does not look up ROM instructions during a |
| step. Weights and event format are constructed; this is not pretrained Beam. |
| """ |
| 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, from_bits |
| from .hf_decoder_cpu import DecoderCPU |
| from .large_cpu import EventResult |
|
|
|
|
| class TwoLayerCPU(DecoderCPU): |
| def __init__(self, program, max_writes=128): |
| super().__init__(program, max_writes) |
| config = BeamConfig(hidden_size=128, intermediate_size=256, num_hidden_layers=2, |
| num_attention_heads=6, num_key_value_heads=1, head_dim=128, |
| layer_types=['full_attention', 'full_attention'], |
| mlp_layer_types=['dense', 'dense'], attention_dropout=0.) |
| config._attn_implementation = 'eager' |
| self.first = BeamDecoderLayer(config, 0).eval().requires_grad_(False) |
| self.layer.self_attn.layer_idx = 1 |
| self.layer.attn_residual_add_idx = 4 |
| self.layer.mlp_residual_add_idx = 5 |
| self.layer.self_attn.config = config |
| |
| ids = torch.arange(256) |
| patterns = (((ids[:, None] >> torch.arange(8)) & 1)*2-1).float() |
| addresses = [] |
| for pc in range(256): |
| op, arg = program[pc] if pc < len(program) else ('HALT', 0) |
| addr = arg if op in {'LOAD', 'STORE', 'ADDM', 'SUBM'} else 0 |
| addresses.append([2*b-1 for b in bits(addr, 4)]) |
| with torch.no_grad(): |
| for proj in (self.first.self_attn.q_proj, self.first.self_attn.k_proj, |
| self.first.self_attn.v_proj, self.first.self_attn.o_proj, |
| self.first.self_attn.gate_proj, self.first.mlp.gate_proj, |
| self.first.mlp.up_proj, self.first.mlp.down_proj): |
| proj.weight.zero_() |
| self.first.self_attn.gate_bias.zero_() |
| self.first.mlp.gate_proj.weight[:, 48:56] = 32 * patterns |
| self.first.mlp.gate_proj.weight[:, 5] = -7.5 * 32 |
| self.first.mlp.up_proj.weight[:, 5] = 1 |
| self.first.mlp.down_proj.weight[:4, :] = torch.tensor(addresses, dtype=torch.float32).T |
| |
| |
| for head in range(6): |
| for i in range(4): self.layer.self_attn.q_proj.weight[head*128+i, i] *= 1/(4*math.sqrt(2)) |
| |
| |
| old_alpha, new_alpha = 1/math.sqrt(2), 1/math.sqrt(3) |
| old_read, new_read = 8., math.sqrt(128/3) |
| self.layer.mlp.gate_proj.weight[:, 48:60] *= old_alpha/new_alpha |
| self.layer.mlp.gate_proj.weight[:, 5] *= old_alpha/new_alpha |
| self.layer.mlp.gate_proj.weight[:, 32:48] *= old_read/new_read |
| self.layer.mlp.up_proj.weight[:, 5] *= old_alpha/new_alpha |
| self.cache = DynamicCache(config=config) |
| self.pc = self.acc = self.writes = 0 |
| self.halted = False |
| rows = torch.stack([self._event(a, 0, 0, True, 0, 0) for a in range(16)]).unsqueeze(0) |
| with torch.inference_mode(): self._forward(rows) |
|
|
| def _event(self, address, value, position, valid, pc, acc): |
| |
| x = torch.zeros(128) |
| x[4:6] = 1 |
| x[6:10] = torch.tensor([2*b-1 for b in bits(address, 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 |
| x[48:60] = torch.tensor([2*b-1 for b in bits(pc, 8)+bits(acc, 4)]) |
| x[-1] = math.sqrt(128-float(x.square().sum())) |
| return x |
|
|
| def _forward(self, rows): |
| if not hasattr(self, 'first'): |
| return super()._forward(rows) |
| old = self.cache.get_seq_length(layer_idx=1) |
| 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) |
| kwargs = dict(position_embeddings=None, attention_mask=mask, |
| attention_scaling=compute_attention_scaling(positions), |
| past_key_values=self.cache, use_cache=True) |
| return self.layer(self.first(rows, **kwargs), **kwargs) |
|
|
| 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._event(address, value, self.writes, True, self.pc, self.acc).reshape(1, 1, -1)) |
|
|
| def step(self): |
| if self.halted: raise RuntimeError('machine is halted') |
| old_acc = self.acc |
| with torch.inference_mode(): |
| output = self._forward(self._event(0, 0, self.writes, False, self.pc, self.acc).reshape(1, 1, -1))[0, 0] |
| if not torch.isfinite(output).all(): raise RuntimeError('non-finite output') |
| address = from_bits((output[:4] > 0).int().tolist()) |
| decoded = output[60:90].reshape(15, 2).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) |
|
|
