File size: 8,255 Bytes
d5147ba | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | # SPDX-License-Identifier: Apache-2.0
# Original project extension; not part of the upstream LOOM snapshot.
# Licensed under the Apache License, Version 2.0.
# See source/LICENSE in this distribution for the license terms.
# This distribution adds this license notice only; executable Python
# syntax is unchanged from the original trained-model artifact.
"""Differentiable per-unit MELT memory and causal history selection."""
from typing import NamedTuple
import torch
import torch.nn.functional as F
from torch import Tensor, nn
class MeltGate(nn.Module):
def __init__(self, width: int, *, device=None, dtype=None):
super().__init__()
self.input_weight = nn.Parameter(torch.zeros(width, width, device=device, dtype=dtype))
self.memory_weight = nn.Parameter(torch.zeros(width, width, device=device, dtype=dtype))
self.bias = nn.Parameter(torch.full((width,), -4.0, device=device, dtype=dtype))
def forward(self, current: Tensor, previous: Tensor) -> Tensor:
retention = torch.sigmoid(
F.linear(current, self.input_weight, self.bias)
+ F.linear(previous, self.memory_weight)
)
return retention * previous + (1.0 - retention) * current
class KV(NamedTuple):
keys: Tensor
values: Tensor
class MeltHistory:
"""One dense equal-length LM unit; tensors stay connected across chunks.
Container assignments commit immutable tensors, never in-place tensor writes.
Final-loop history survives for the whole unit. SHARED's earlier-loop Python
references retain only the recent window; autograd may retain older tensors.
"""
def __init__(self, layers: int, loops: int, policy: str, window_size: int,
*, gate_enabled: bool = True):
if policy not in ("full", "shared"):
raise ValueError("Policy must be full or shared")
if type(window_size) is not int or window_size < 0:
raise ValueError("Window size must be a nonnegative integer")
if layers < 1 or loops < 1:
raise ValueError("Layers and loops must be positive")
self.layers, self.loops = layers, loops
self.policy, self.window_size = policy, window_size
self.gate_enabled = gate_enabled
self.history: list[list[KV | None]] = [[None] * loops for _ in range(layers)]
self.length = 0
self.failed = False
self._active = False
self._batch_size: int | None = None
def begin_chunk(self, batch_size: int, size: int) -> None:
if self.failed or self._active:
raise RuntimeError("History is failed or has an unfinished chunk")
if size < 1 or size > self.window_size + 1:
raise ValueError("Exact chunks require 1 <= C <= W+1")
if self._batch_size is not None and batch_size != self._batch_size:
raise ValueError("Unit batch size cannot change")
self._batch_size = batch_size
self._size = size
self._next_visit = 0
self._memory: list[Tensor | None] = [None] * self.layers
self._chunk: list[list[KV | None]] = [[None] * self.loops for _ in range(self.layers)]
self._active = True
def _check_visit(self, layer: int, loop: int) -> None:
if self.failed or not self._active:
raise RuntimeError("Begin a healthy chunk before layer visits")
if self._next_visit != loop * self.layers + layer:
raise RuntimeError("MELT layer/loop visits must follow the native loop order")
def memory(self, layer: int, loop: int, current: Tensor, gate: MeltGate) -> Tensor:
self._check_visit(layer, loop)
if current.ndim != 3 or current.shape[:2] != (self._batch_size, self._size):
raise ValueError("MELT requires a dense synchronous B x C x width input")
previous = self._memory[layer]
if loop == 0 or not self.gate_enabled:
result = current
else:
if previous is None:
raise RuntimeError("Previous-loop latent memory is missing")
result = gate(current, previous)
self._memory[layer] = result
return result
def attend(self, layer: int, loop: int, query: Tensor, key: Tensor,
value: Tensor) -> Tensor:
self._check_visit(layer, loop)
current = KV(key, value)
prefix = self.history[layer][loop]
query_pos = torch.arange(self.length, self.length + self._size, device=query.device)
current_pos = query_pos
parts = []
masks = []
if self.policy == "full" or loop == self.loops - 1:
if prefix is not None:
parts.append(prefix)
masks.append(torch.ones(self._size, prefix.keys.shape[1],
device=query.device, dtype=torch.bool))
else:
boundary = (query_pos - self.window_size).clamp_min(0)
final = self.history[layer][-1]
if final is not None:
parts.append(final)
positions = torch.arange(self.length, device=query.device)
masks.append(positions[None, :] < boundary[:, None])
if prefix is not None:
parts.append(prefix)
positions = torch.arange(self.length - prefix.keys.shape[1], self.length,
device=query.device)
masks.append(positions[None, :] >= boundary[:, None])
parts.append(current)
masks.append(current_pos[None, :] <= query_pos[:, None])
keys = torch.cat([part.keys for part in parts], dim=1).transpose(1, 2)
values = torch.cat([part.values for part in parts], dim=1).transpose(1, 2)
heads = query.shape[2]
if heads % keys.shape[1]:
raise ValueError("Query heads must be divisible by KV heads")
repeats = heads // keys.shape[1]
if repeats != 1:
keys = keys.repeat_interleave(repeats, dim=1)
values = values.repeat_interleave(repeats, dim=1)
mask = torch.cat(masks, dim=-1)[None, None]
output = F.scaled_dot_product_attention(
query.transpose(1, 2), keys, values, attn_mask=mask,
dropout_p=0.0, is_causal=False,
).transpose(1, 2)
self._chunk[layer][loop] = current
self._next_visit += 1
return output
def commit_chunk(self) -> None:
if self.failed or not self._active or self._next_visit != self.layers * self.loops:
raise RuntimeError("Complete every native layer/loop visit before commit")
for layer in range(self.layers):
for loop in range(self.loops):
previous, current = self.history[layer][loop], self._chunk[layer][loop]
if current is None:
raise RuntimeError("Missing current K/V")
if previous is None:
combined = current
else:
combined = KV(torch.cat((previous.keys, current.keys), dim=1),
torch.cat((previous.values, current.values), dim=1))
if self.policy == "shared" and loop != self.loops - 1:
start = max(0, combined.keys.shape[1] - self.window_size)
combined = KV(combined.keys[:, start:].clone(),
combined.values[:, start:].clone())
self.history[layer][loop] = combined
self.length += self._size
self._memory = []
self._chunk = []
self._active = False
@property
def live_kv_bytes(self) -> int:
return sum(tensor.numel() * tensor.element_size()
for layer in self.history for entry in layer if entry is not None
for tensor in entry)
@property
def referenced_kv_storage_bytes(self) -> int:
storages = {}
for layer in self.history:
for entry in layer:
if entry is not None:
for tensor in entry:
storage = tensor.untyped_storage()
storages[(str(tensor.device), storage.data_ptr())] = storage.nbytes()
return sum(storages.values())
|