| """ÇABA v0-core: recurrent three-bank causal LM for Hugging Face Transformers. |
| |
| The recurrent update here is an explicit delta-rule prototype, not a claim of |
| bit-for-bit equivalence with the official Gated DeltaNet-2 implementation. |
| Training on TPU uses PyTorch/XLA's scan operator; cached inference keeps a |
| constant-size recurrent state plus the unfinished fixed-size promotion window. |
| """ |
|
|
| from dataclasses import dataclass |
| import math |
| from typing import Any, Dict, List, Optional, Tuple |
|
|
| import torch |
| import transformers as _transformers |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.utils.checkpoint import checkpoint as _torch_checkpoint |
| try: |
| from torch_xla.utils.checkpoint import checkpoint as _xla_checkpoint |
| except ImportError: |
| _xla_checkpoint = None |
| from transformers import PreTrainedModel |
| from transformers.generation import GenerationMixin |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
| from .configuration_caba import CabaConfig |
|
|
| try: |
| from torch_xla.experimental.scan import scan as _xla_scan |
| except Exception: |
| _xla_scan = None |
|
|
|
|
| @dataclass |
| class CabaCausalLMOutput(CausalLMOutputWithPast): |
| """Causal LM output with small scalar diagnostics for training logs.""" |
|
|
| memory_metrics: Optional[Dict[str, torch.Tensor]] = None |
|
|
|
|
| def _inverse_softplus(x: float) -> float: |
| return math.log(math.expm1(x)) |
|
|
|
|
| def _scan(fn, carry: torch.Tensor, xs: torch.Tensor): |
| """Run a scan on TPU, with an eager reference path for CPU smoke use.""" |
| if carry.device.type == "xla": |
| if _xla_scan is None: |
| raise RuntimeError( |
| "TPU training needs torch_xla.experimental.scan. Install a matched " |
| "PyTorch/PyTorch-XLA pair from the first notebook cell and restart runtime." |
| ) |
| return _xla_scan(fn, carry, xs) |
|
|
| outputs = [] |
| for index in range(xs.shape[0]): |
| carry, output = fn(carry, xs[index]) |
| outputs.append(output) |
| if not outputs: |
| raise ValueError("scan input must contain at least one step") |
| return carry, torch.stack(outputs, dim=0) |
|
|
|
|
| def _read(memory: torch.Tensor, query: torch.Tensor) -> torch.Tensor: |
| |
| return torch.einsum("bhkv,bhk->bhv", memory, query.float()) |
|
|
|
|
| def _predict_error(memory: torch.Tensor, key: torch.Tensor, value: torch.Tensor): |
| prediction = _read(memory, key) |
| error = value.float() - prediction |
| score = ( |
| torch.linalg.vector_norm(error, dim=-1) |
| / (torch.linalg.vector_norm(value.float(), dim=-1) + 1e-6) |
| ) |
| return error, score.detach() |
|
|
|
|
| def _outer(error: torch.Tensor, key: torch.Tensor) -> torch.Tensor: |
| return torch.einsum("bhv,bhk->bhkv", error, key.float()) |
|
|
|
|
| def _mark_data_sharding(tensor: torch.Tensor, mesh, spec): |
| if mesh is None: |
| return tensor |
| import torch_xla.distributed.spmd as xs |
| return xs.mark_sharding(tensor, mesh, spec) |
|
|
|
|
| class CabaRMSNorm(nn.Module): |
| def __init__(self, hidden_size: int, eps: float): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.eps = eps |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| scale = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) |
| return (x.float() * scale * self.weight.float()).to(x.dtype) |
|
|
|
|
| class CabaBlock(nn.Module): |
| def __init__(self, config: CabaConfig): |
| super().__init__() |
| self.hidden = config.hidden_size |
| self.heads = config.num_heads |
| self.head_dim = config.head_dim |
| self.medium_clock = config.medium_clock |
| self.slow_clock = config.slow_clock |
| self.slow_span = self.medium_clock * self.slow_clock |
| self.conv_kernel = config.conv_kernel |
| self.eps = config.rms_norm_eps |
| self.data_mesh = None |
|
|
| self.input_norm = CabaRMSNorm(self.hidden, self.eps) |
| self.local_conv = nn.Conv1d( |
| self.hidden, |
| self.hidden, |
| kernel_size=self.conv_kernel, |
| groups=self.hidden, |
| bias=False, |
| ) |
| self.qkv = nn.Linear(self.hidden, 3 * self.hidden, bias=False) |
| self.fast_write = nn.Linear(self.hidden, self.heads, bias=False) |
| self.merge_gate = nn.Linear(self.hidden, 3 * self.heads, bias=False) |
| self.out_proj = nn.Linear(self.hidden, self.hidden, bias=False) |
|
|
| self.fast_score_alpha = nn.Parameter( |
| torch.full((self.heads,), _inverse_softplus(1.0)) |
| ) |
| self.fast_score_tau = nn.Parameter(torch.ones(self.heads)) |
| self.medium_score_alpha = nn.Parameter( |
| torch.full((self.heads,), _inverse_softplus(1.0)) |
| ) |
| self.medium_score_tau = nn.Parameter(torch.ones(self.heads)) |
| self.medium_admit_threshold = nn.Parameter(torch.full((self.heads,), 2.0)) |
| self.slow_admit_threshold = nn.Parameter(torch.full((self.heads,), 3.0)) |
|
|
| |
| self.raw_fast_timescale = nn.Parameter( |
| torch.full((self.heads,), _inverse_softplus(7.0)) |
| ) |
| self.raw_medium_gap = nn.Parameter( |
| torch.full((self.heads,), _inverse_softplus(56.0)) |
| ) |
| self.raw_slow_gap = nn.Parameter( |
| torch.full((self.heads,), _inverse_softplus(448.0)) |
| ) |
|
|
| self.ffn_norm = CabaRMSNorm(self.hidden, self.eps) |
| self.ffn_gate_up = nn.Linear(self.hidden, 2 * config.intermediate_size, bias=False) |
| self.ffn_down = nn.Linear(config.intermediate_size, self.hidden, bias=False) |
|
|
| def _decays(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| t_fast = F.softplus(self.raw_fast_timescale.float()) + 1.0 |
| t_medium = t_fast + F.softplus(self.raw_medium_gap.float()) |
| t_slow = t_medium + F.softplus(self.raw_slow_gap.float()) |
| return torch.exp(-1.0 / t_fast), torch.exp(-1.0 / t_medium), torch.exp(-1.0 / t_slow) |
|
|
| def _project(self, x: torch.Tensor): |
| batch, length, _ = x.shape |
| normalized = self.input_norm(x) |
| conv_in = F.pad(normalized.transpose(1, 2), (self.conv_kernel - 1, 0)) |
| h = F.silu(self.local_conv(conv_in).transpose(1, 2)) |
| q, k, v = self.qkv(h).view( |
| batch, length, 3, self.heads, self.head_dim |
| ).unbind(dim=2) |
| q = F.normalize(q.float(), dim=-1, eps=1e-6) |
| k = F.normalize(k.float(), dim=-1, eps=1e-6) |
| v = v.float() |
| write = torch.sigmoid(self.fast_write(h).float()).unsqueeze(-1) |
| return h, q, k, v, write |
|
|
| def _fast_sequence(self, q, k, v, write): |
| batch, length, heads, dim = q.shape |
| decay_fast, _, _ = self._decays() |
| decay_fast = decay_fast.view(1, heads, 1, 1) |
| xs = torch.cat((q, k, v, write), dim=-1).permute(1, 0, 2, 3).contiguous() |
| decay_steps = decay_fast.reshape(1, 1, heads, 1).expand(length, batch, heads, 1) |
| xs = torch.cat((xs, decay_steps), dim=-1).contiguous() |
| xs = _mark_data_sharding(xs, self.data_mesh, (None, "data", None, None)) |
| initial = torch.zeros( |
| batch, heads, dim, dim, device=q.device, dtype=torch.float32 |
| ) |
| initial = _mark_data_sharding(initial, self.data_mesh, ("data", None, None, None)) |
|
|
| def step(memory, item): |
| query, key, value, gate, decay = torch.split( |
| item, (dim, dim, dim, 1, 1), dim=-1 |
| ) |
| key = F.normalize(key.float(), dim=-1, eps=1e-6) |
| error, score = _predict_error(memory, key.squeeze(-1) if key.shape[-1] == 1 else key, value) |
| next_memory = decay.unsqueeze(-1) * memory + gate.float().unsqueeze(-1) * _outer(error, key) |
| read = _read(next_memory, query) |
| packed = torch.cat((read, error, score.unsqueeze(-1)), dim=-1) |
| return next_memory, packed |
|
|
| final, packed = _scan(step, initial, xs) |
| final = _mark_data_sharding(final, self.data_mesh, ("data", None, None, None)) |
| packed = _mark_data_sharding(packed, self.data_mesh, (None, "data", None, None)) |
| packed = packed.permute(1, 0, 2, 3).contiguous() |
| read = _mark_data_sharding(packed[..., :dim], self.data_mesh, ("data", None, None, None)) |
| error = _mark_data_sharding(packed[..., dim : 2 * dim], self.data_mesh, ("data", None, None, None)) |
| score = _mark_data_sharding(packed[..., -1], self.data_mesh, ("data", None, None)) |
| return final, read, error, score |
|
|
| def _medium_sequence(self, q, k, v, scores): |
| batch, length, heads, dim = q.shape |
| clock = self.medium_clock |
| if length % clock: |
| raise ValueError( |
| f"Training sequence length {length} must be divisible by medium_clock={clock}." |
| ) |
| windows = length // clock |
| q_windows = q.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4) |
| k_windows = k.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4) |
| v_windows = v.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4) |
| s_windows = scores.reshape(batch, windows, clock, heads).permute(0, 1, 3, 2) |
|
|
| selection = torch.softmax(s_windows, dim=-1) |
| summary_k = (selection.unsqueeze(-1) * k_windows).sum(dim=-2) |
| summary_v = (selection.unsqueeze(-1) * v_windows).sum(dim=-2) |
| log_mean_exp = torch.logsumexp(s_windows, dim=-1) - math.log(clock) |
| admission = torch.sigmoid( |
| log_mean_exp - self.medium_admit_threshold.float().view(1, 1, heads) |
| ).unsqueeze(-1) |
|
|
| q_flat = q_windows.reshape(batch, windows, heads, clock * dim).permute(1, 0, 2, 3) |
| xs = torch.cat( |
| ( |
| q_flat, |
| summary_k.permute(1, 0, 2, 3), |
| summary_v.permute(1, 0, 2, 3), |
| admission.permute(1, 0, 2, 3), |
| ), |
| dim=-1, |
| ).contiguous() |
| initial = torch.zeros(batch, heads, dim, dim, device=q.device, dtype=torch.float32) |
| _, decay_medium, _ = self._decays() |
| decay_medium = decay_medium.pow(clock).view(1, heads, 1, 1) |
| decay_steps = decay_medium.reshape(1, 1, heads, 1).expand(windows, batch, heads, 1) |
| xs = torch.cat((xs, decay_steps), dim=-1).contiguous() |
| xs = _mark_data_sharding(xs, self.data_mesh, (None, "data", None, None)) |
| initial = _mark_data_sharding(initial, self.data_mesh, ("data", None, None, None)) |
|
|
| def step(memory, item): |
| offset = clock * dim |
| query = item[..., :offset].reshape(batch, heads, clock, dim).permute(0, 2, 1, 3) |
| key = item[..., offset : offset + dim] |
| value = item[..., offset + dim : offset + 2 * dim] |
| gate = item[..., -2:-1] |
| decay = item[..., -1:] |
| read = torch.einsum("bhkv,brhk->brhv", memory, query.float()) |
| error, score = _predict_error(memory, key, value) |
| next_memory = ( |
| decay.unsqueeze(-1) * memory + gate.float().unsqueeze(-1) * _outer(error, key) |
| ) |
| repeated_error = error.unsqueeze(2).expand(-1, -1, clock, -1) |
| repeated_score = score.unsqueeze(-1).unsqueeze(2).expand(-1, -1, clock, -1) |
| packed = torch.cat( |
| (read.permute(0, 2, 1, 3), repeated_error, repeated_score), dim=-1 |
| ).reshape(batch, heads, clock * (2 * dim + 1)) |
| return next_memory, packed |
|
|
| final, packed = _scan(step, initial, xs) |
| final = _mark_data_sharding(final, self.data_mesh, ("data", None, None, None)) |
| packed = _mark_data_sharding(packed, self.data_mesh, (None, "data", None, None)) |
| packed = packed.reshape(windows, batch, heads, clock, 2 * dim + 1) |
| packed = packed.permute(1, 0, 3, 2, 4).contiguous() |
| read = _mark_data_sharding(packed[..., :dim].reshape(batch, length, heads, dim), self.data_mesh, ("data", None, None, None)) |
| |
| error = _mark_data_sharding(packed[:, :, 0, :, dim : 2 * dim], self.data_mesh, ("data", None, None, None)) |
| score = _mark_data_sharding(packed[:, :, 0, :, -1], self.data_mesh, ("data", None, None)) |
| summary_k = _mark_data_sharding(summary_k, self.data_mesh, ("data", None, None, None)) |
| summary_v = _mark_data_sharding(summary_v, self.data_mesh, ("data", None, None, None)) |
| admission = _mark_data_sharding(admission.squeeze(-1), self.data_mesh, ("data", None, None)) |
| return final, read, summary_k, summary_v, error, score, admission |
|
|
| def _slow_sequence(self, q, medium_k, medium_v, medium_scores): |
| batch, length, heads, dim = q.shape |
| clock = self.slow_span |
| if length % clock: |
| raise ValueError( |
| f"Training sequence length {length} must be divisible by slow span={clock}." |
| ) |
| windows = length // clock |
| medium_windows_per_slow = self.slow_clock |
| score_windows = medium_scores.reshape( |
| batch, windows, medium_windows_per_slow, heads |
| ).permute(0, 1, 3, 2) |
| mk = medium_k.reshape(batch, windows, medium_windows_per_slow, heads, dim).permute( |
| 0, 1, 3, 2, 4 |
| ) |
| mv = medium_v.reshape(batch, windows, medium_windows_per_slow, heads, dim).permute( |
| 0, 1, 3, 2, 4 |
| ) |
| selection = torch.softmax(score_windows, dim=-1) |
| summary_k = (selection.unsqueeze(-1) * mk).sum(dim=-2) |
| summary_v = (selection.unsqueeze(-1) * mv).sum(dim=-2) |
| log_mean_exp = torch.logsumexp(score_windows, dim=-1) - math.log( |
| medium_windows_per_slow |
| ) |
| admission = torch.sigmoid( |
| log_mean_exp - self.slow_admit_threshold.float().view(1, 1, heads) |
| ).unsqueeze(-1) |
|
|
| q_windows = q.reshape(batch, windows, clock, heads, dim).permute(0, 1, 3, 2, 4) |
| q_flat = q_windows.reshape(batch, windows, heads, clock * dim).permute(1, 0, 2, 3) |
| xs = torch.cat( |
| ( |
| q_flat, |
| summary_k.permute(1, 0, 2, 3), |
| summary_v.permute(1, 0, 2, 3), |
| admission.permute(1, 0, 2, 3), |
| ), |
| dim=-1, |
| ).contiguous() |
| initial = torch.zeros(batch, heads, dim, dim, device=q.device, dtype=torch.float32) |
| _, _, decay_slow = self._decays() |
| decay_slow = decay_slow.pow(clock).view(1, heads, 1, 1) |
| decay_steps = decay_slow.reshape(1, 1, heads, 1).expand(windows, batch, heads, 1) |
| xs = torch.cat((xs, decay_steps), dim=-1).contiguous() |
| xs = _mark_data_sharding(xs, self.data_mesh, (None, "data", None, None)) |
| initial = _mark_data_sharding(initial, self.data_mesh, ("data", None, None, None)) |
|
|
| def step(memory, item): |
| offset = clock * dim |
| query = item[..., :offset].reshape(batch, heads, clock, dim).permute(0, 2, 1, 3) |
| key = item[..., offset : offset + dim] |
| value = item[..., offset + dim : offset + 2 * dim] |
| gate = item[..., -2:-1] |
| decay = item[..., -1:] |
| read = torch.einsum("bhkv,brhk->brhv", memory, query.float()) |
| error, _ = _predict_error(memory, key, value) |
| next_memory = ( |
| decay.unsqueeze(-1) * memory + gate.float().unsqueeze(-1) * _outer(error, key) |
| ) |
| return next_memory, read.permute(0, 2, 1, 3).reshape( |
| batch, heads, clock * dim |
| ) |
|
|
| final, read_flat = _scan(step, initial, xs) |
| final = _mark_data_sharding(final, self.data_mesh, ("data", None, None, None)) |
| read_flat = _mark_data_sharding(read_flat, self.data_mesh, (None, "data", None, None)) |
| read = read_flat.reshape(windows, batch, heads, clock, dim) |
| read = _mark_data_sharding(read.permute(1, 0, 3, 2, 4).reshape(batch, length, heads, dim), self.data_mesh, ("data", None, None, None)) |
| admission = _mark_data_sharding(admission.squeeze(-1), self.data_mesh, ("data", None, None)) |
| return final, read, admission |
|
|
| def forward_sequence(self, x: torch.Tensor): |
| batch, length, _ = x.shape |
| if length % (self.medium_clock * self.slow_clock): |
| raise ValueError( |
| "use_cache=False training/evaluation sequence length must be divisible " |
| f"by medium_clock*slow_clock={self.medium_clock * self.slow_clock}; " |
| "the notebook default is 512." |
| ) |
| h, q, k, v, write = self._project(x) |
| h = _mark_data_sharding(h, self.data_mesh, ("data", None, None)) |
| q = _mark_data_sharding(q, self.data_mesh, ("data", None, None, None)) |
| k = _mark_data_sharding(k, self.data_mesh, ("data", None, None, None)) |
| v = _mark_data_sharding(v, self.data_mesh, ("data", None, None, None)) |
| write = _mark_data_sharding(write, self.data_mesh, ("data", None, None, None)) |
| _, fast_read, fast_error, fast_score = self._fast_sequence(q, k, v, write) |
| unresolvedness = fast_score.detach().mean() |
| fast_score = F.softplus(self.fast_score_alpha.float()).view(1, 1, -1) * ( |
| fast_score - self.fast_score_tau.float().view(1, 1, -1) |
| ) |
| _, medium_read, medium_k, medium_v, medium_error, medium_score, medium_admission = ( |
| self._medium_sequence(q, k, v, fast_score) |
| ) |
| medium_score = F.softplus(self.medium_score_alpha.float()).view(1, 1, -1) * ( |
| medium_score - self.medium_score_tau.float().view(1, 1, -1) |
| ) |
| _, slow_read, slow_admission = self._slow_sequence( |
| q, medium_k, medium_v, medium_score |
| ) |
|
|
| gates = torch.softmax( |
| self.merge_gate(h).view(batch, length, self.heads, 3).float(), dim=-1 |
| ) |
| reads = torch.stack((fast_read, medium_read, slow_read), dim=-2) |
| merged = (reads * gates.unsqueeze(-1)).sum(dim=-2).reshape(batch, length, self.hidden) |
| x = x + self.out_proj(merged.to(x.dtype)) |
| gate_up = self.ffn_gate_up(self.ffn_norm(x)) |
| gate, up = gate_up.chunk(2, dim=-1) |
| x = x + self.ffn_down(F.silu(gate) * up) |
| metrics = { |
| "unresolvedness": unresolvedness, |
| "medium_admission": medium_admission.detach().mean(), |
| "slow_admission": slow_admission.detach().mean(), |
| } |
| return x, metrics |
|
|
| def _new_cache(self, batch: int, device: torch.device) -> Dict[str, Any]: |
| h, d = self.heads, self.head_dim |
| return { |
| "fast": torch.zeros(batch, h, d, d, device=device, dtype=torch.float32), |
| "medium": torch.zeros(batch, h, d, d, device=device, dtype=torch.float32), |
| "slow": torch.zeros(batch, h, d, d, device=device, dtype=torch.float32), |
| "conv": torch.zeros( |
| batch, max(self.conv_kernel - 1, 0), self.hidden, device=device |
| ), |
| "medium_pending_k": [], |
| "medium_pending_v": [], |
| "medium_pending_score": [], |
| "slow_pending_k": [], |
| "slow_pending_v": [], |
| "slow_pending_score": [], |
| "position": 0, |
| } |
|
|
| def step(self, x: torch.Tensor, cache: Dict[str, Any], position: int): |
| batch = x.shape[0] |
| normalized = self.input_norm(x) |
| context = torch.cat((cache["conv"], normalized), dim=1) |
| conv = self.local_conv(context.transpose(1, 2)).transpose(1, 2) |
| h = F.silu(conv) |
| conv_state = ( |
| context[:, -(self.conv_kernel - 1) :, :] |
| if self.conv_kernel > 1 |
| else context[:, :0, :] |
| ) |
|
|
| q, k, v = self.qkv(h).view(batch, 1, 3, self.heads, self.head_dim).unbind(dim=2) |
| q = F.normalize(q[:, 0].float(), dim=-1, eps=1e-6) |
| k = F.normalize(k[:, 0].float(), dim=-1, eps=1e-6) |
| v = v[:, 0].float() |
| write = torch.sigmoid(self.fast_write(h).float()[:, 0]).unsqueeze(-1) |
|
|
| decay_fast, decay_medium, decay_slow = self._decays() |
| fast_error, fast_unresolvedness = _predict_error(cache["fast"], k, v) |
| fast_next = ( |
| decay_fast.view(1, self.heads, 1, 1) * cache["fast"] |
| + write.float().unsqueeze(-1) * _outer(fast_error, k) |
| ) |
| fast_read = _read(fast_next, q) |
| medium_read = _read(cache["medium"], q) |
| slow_read = _read(cache["slow"], q) |
|
|
| score_fast = F.softplus(self.fast_score_alpha.float()).view(1, -1) * ( |
| fast_unresolvedness - self.fast_score_tau.float().view(1, -1) |
| ) |
| pending_k = cache["medium_pending_k"] + [k] |
| pending_v = cache["medium_pending_v"] + [v] |
| pending_score = cache["medium_pending_score"] + [score_fast] |
| medium_next = cache["medium"] |
| slow_next = cache["slow"] |
| slow_pending_k = cache["slow_pending_k"] |
| slow_pending_v = cache["slow_pending_v"] |
| slow_pending_score = cache["slow_pending_score"] |
| medium_admission_value = torch.zeros((), device=x.device) |
| slow_admission_value = torch.zeros((), device=x.device) |
|
|
| if (position + 1) % self.medium_clock == 0: |
| k_stack = torch.stack(pending_k, dim=2) |
| v_stack = torch.stack(pending_v, dim=2) |
| s_stack = torch.stack(pending_score, dim=2) |
| selected = torch.softmax(s_stack, dim=2) |
| summary_k = (selected.unsqueeze(-1) * k_stack).sum(dim=2) |
| summary_v = (selected.unsqueeze(-1) * v_stack).sum(dim=2) |
| log_mean_exp = torch.logsumexp(s_stack, dim=2) - math.log(self.medium_clock) |
| medium_admission = torch.sigmoid( |
| log_mean_exp |
| - self.medium_admit_threshold.float().view(1, self.heads) |
| ).unsqueeze(-1) |
| medium_error, medium_unresolvedness = _predict_error( |
| cache["medium"], summary_k, summary_v |
| ) |
| medium_next = ( |
| decay_medium.pow(self.medium_clock).view(1, self.heads, 1, 1) |
| * cache["medium"] |
| + medium_admission.unsqueeze(-1) * _outer(medium_error, summary_k) |
| ) |
| score_medium = F.softplus(self.medium_score_alpha.float()).view(1, -1) * ( |
| medium_unresolvedness - self.medium_score_tau.float().view(1, -1) |
| ) |
| slow_pending_k = slow_pending_k + [summary_k] |
| slow_pending_v = slow_pending_v + [summary_v] |
| slow_pending_score = slow_pending_score + [score_medium] |
| pending_k, pending_v, pending_score = [], [], [] |
| medium_admission_value = medium_admission.mean() |
|
|
| if len(slow_pending_k) == self.slow_clock: |
| sk = torch.stack(slow_pending_k, dim=2) |
| sv = torch.stack(slow_pending_v, dim=2) |
| ss = torch.stack(slow_pending_score, dim=2) |
| slow_selected = torch.softmax(ss, dim=2) |
| slow_key = (slow_selected.unsqueeze(-1) * sk).sum(dim=2) |
| slow_value = (slow_selected.unsqueeze(-1) * sv).sum(dim=2) |
| slow_log_mean_exp = torch.logsumexp(ss, dim=2) - math.log(self.slow_clock) |
| slow_admission = torch.sigmoid( |
| slow_log_mean_exp |
| - self.slow_admit_threshold.float().view(1, self.heads) |
| ).unsqueeze(-1) |
| slow_error, _ = _predict_error(cache["slow"], slow_key, slow_value) |
| slow_next = ( |
| decay_slow.pow(self.slow_span).view(1, self.heads, 1, 1) |
| * cache["slow"] |
| + slow_admission.unsqueeze(-1) * _outer(slow_error, slow_key) |
| ) |
| slow_pending_k, slow_pending_v, slow_pending_score = [], [], [] |
| slow_admission_value = slow_admission.mean() |
|
|
| gates = torch.softmax( |
| self.merge_gate(h).view(batch, 1, self.heads, 3).float(), dim=-1 |
| )[:, 0] |
| reads = torch.stack((fast_read, medium_read, slow_read), dim=-2) |
| merged = (reads * gates.unsqueeze(-1)).sum(dim=-2).reshape(batch, 1, self.hidden) |
| x = x + self.out_proj(merged.to(x.dtype)) |
| gate_up = self.ffn_gate_up(self.ffn_norm(x)) |
| gate, up = gate_up.chunk(2, dim=-1) |
| x = x + self.ffn_down(F.silu(gate) * up) |
|
|
| cache = { |
| "fast": fast_next, |
| "medium": medium_next, |
| "slow": slow_next, |
| "conv": conv_state, |
| "medium_pending_k": pending_k, |
| "medium_pending_v": pending_v, |
| "medium_pending_score": pending_score, |
| "slow_pending_k": slow_pending_k, |
| "slow_pending_v": slow_pending_v, |
| "slow_pending_score": slow_pending_score, |
| "position": position + 1, |
| } |
| metrics = { |
| "unresolvedness": fast_unresolvedness.detach().mean(), |
| "medium_admission": medium_admission_value.detach(), |
| "slow_admission": slow_admission_value.detach(), |
| } |
| return x, cache, metrics |
|
|
|
|
| class CabaPreTrainedModel(PreTrainedModel): |
| config_class = CabaConfig |
| base_model_prefix = "caba" |
| supports_gradient_checkpointing = False |
| _supports_cache_class = False |
|
|
| def _init_weights(self, module): |
| std = self.config.initializer_range |
| if isinstance(module, (nn.Linear, nn.Embedding, nn.Conv1d)): |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
| if getattr(module, "bias", None) is not None: |
| nn.init.zeros_(module.bias) |
| if isinstance(module, nn.Embedding) and module.padding_idx is not None: |
| with torch.no_grad(): |
| module.weight[module.padding_idx].zero_() |
| elif isinstance(module, CabaRMSNorm): |
| nn.init.ones_(module.weight) |
|
|
|
|
| class CabaModel(CabaPreTrainedModel): |
| def __init__(self, config: CabaConfig): |
| super().__init__(config) |
| self.data_mesh = None |
| self.embed_tokens = nn.Embedding( |
| config.vocab_size, |
| config.hidden_size, |
| padding_idx=config.pad_token_id, |
| ) |
| self.layers = nn.ModuleList( |
| [CabaBlock(config) for _ in range(config.num_hidden_layers)] |
| ) |
| self.norm = CabaRMSNorm(config.hidden_size, config.rms_norm_eps) |
|
|
| def get_input_embeddings(self): |
| return self.embed_tokens |
|
|
| def set_input_embeddings(self, value): |
| self.embed_tokens = value |
|
|
| def set_data_mesh(self, mesh): |
| self.data_mesh = mesh |
| for layer in self.layers: |
| layer.data_mesh = mesh |
|
|
| def forward( |
| self, |
| input_ids: torch.LongTensor, |
| past_key_values=None, |
| use_cache: bool = False, |
| **kwargs, |
| ): |
| if input_ids.ndim != 2: |
| raise ValueError("input_ids must have shape [batch, sequence]") |
| x = self.embed_tokens(input_ids) |
| x = _mark_data_sharding(x, self.data_mesh, ("data", None, None)) |
| metrics_by_layer = [] |
|
|
| if use_cache or past_key_values is not None: |
| caches = ( |
| list(past_key_values) |
| if past_key_values is not None |
| else [ |
| layer._new_cache(input_ids.shape[0], input_ids.device) |
| for layer in self.layers |
| ] |
| ) |
| position = caches[0]["position"] if caches else 0 |
| pieces = [] |
| step_metrics = [] |
| for offset in range(input_ids.shape[1]): |
| hidden = x[:, offset : offset + 1] |
| new_caches = [] |
| current_metrics = [] |
| for layer, cache in zip(self.layers, caches): |
| hidden, cache, metrics = layer.step(hidden, cache, position + offset) |
| new_caches.append(cache) |
| current_metrics.append(metrics) |
| caches = new_caches |
| pieces.append(hidden) |
| step_metrics.append(current_metrics) |
| x = torch.cat(pieces, dim=1) |
| for layer_index in range(len(self.layers)): |
| metrics_by_layer.append( |
| { |
| key: torch.stack( |
| [m[layer_index][key] for m in step_metrics] |
| ).mean() |
| for key in ("unresolvedness", "medium_admission", "slow_admission") |
| } |
| ) |
| cache_output = tuple(caches) |
| else: |
| for layer in self.layers: |
| if self.config.gradient_checkpointing and self.training: |
| def _run_checkpointed(hidden, _layer=layer): |
| hidden_out, layer_metrics = _layer.forward_sequence(hidden) |
| values = torch.stack([ |
| layer_metrics["unresolvedness"].detach().float(), |
| layer_metrics["medium_admission"].detach().float(), |
| layer_metrics["slow_admission"].detach().float(), |
| ]) |
| return hidden_out, values |
| checkpoint_fn = _xla_checkpoint if x.device.type == "xla" else _torch_checkpoint |
| if checkpoint_fn is None: |
| raise RuntimeError("TPU gradient checkpointing requires torch_xla.utils.checkpoint") |
| x, values = checkpoint_fn(_run_checkpointed, x, use_reentrant=True) |
| metrics_by_layer.append({ |
| "unresolvedness": values[0], |
| "medium_admission": values[1], |
| "slow_admission": values[2], |
| }) |
| else: |
| x, metrics = layer.forward_sequence(x) |
| metrics_by_layer.append(metrics) |
| x = _mark_data_sharding(x, self.data_mesh, ("data", None, None)) |
| cache_output = None |
|
|
| x = self.norm(x) |
| x = _mark_data_sharding(x, self.data_mesh, ("data", None, None)) |
| keys = ("unresolvedness", "medium_admission", "slow_admission") |
| metrics = { |
| key: torch.stack([entry[key] for entry in metrics_by_layer]).mean() |
| for key in keys |
| } |
| return x, cache_output, metrics |
|
|
|
|
| class CabaForCausalLM(CabaPreTrainedModel, GenerationMixin): |
| """Hugging Face causal-LM wrapper with recurrent cached generation.""" |
|
|
| |
| _tied_weights_keys = ( |
| {"lm_head.weight": "model.embed_tokens.weight"} |
| if int(_transformers.__version__.split(".", 1)[0]) >= 5 |
| else ["lm_head.weight"] |
| ) |
| _supports_cache_class = False |
| main_input_name = "input_ids" |
|
|
| def __init__(self, config: CabaConfig): |
| super().__init__(config) |
| self.model = CabaModel(config) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.post_init() |
|
|
| if config.tie_word_embeddings: |
| self.lm_head.weight = self.model.embed_tokens.weight |
|
|
| def get_input_embeddings(self): |
| return self.model.embed_tokens |
|
|
| def set_input_embeddings(self, value): |
| self.model.set_input_embeddings(value) |
|
|
| def get_output_embeddings(self): |
| return self.lm_head |
|
|
| def set_output_embeddings(self, value): |
| self.lm_head = value |
|
|
| def forward( |
| self, |
| input_ids: torch.LongTensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| labels: Optional[torch.LongTensor] = None, |
| past_key_values=None, |
| use_cache: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| **kwargs, |
| ): |
| if use_cache is None: |
| use_cache = bool(self.config.use_cache and not self.training) |
| hidden, cache, metrics = self.model( |
| input_ids, |
| past_key_values=past_key_values, |
| use_cache=bool(use_cache), |
| ) |
| logits = self.lm_head(hidden).float() |
| loss = None |
| if labels is not None: |
| shift_logits = logits[:, :-1, :].contiguous() |
| shift_labels = labels[:, 1:].contiguous() |
| loss = F.cross_entropy( |
| shift_logits.view(-1, self.config.vocab_size), |
| shift_labels.view(-1), |
| ignore_index=-100, |
| ) |
| output = CabaCausalLMOutput( |
| loss=loss, |
| logits=logits, |
| past_key_values=cache, |
| memory_metrics=metrics, |
| ) |
| if return_dict is False: |
| items = (logits, cache, metrics) |
| return ((loss,) + items) if loss is not None else items |
| return output |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| input_ids: torch.LongTensor, |
| max_new_tokens: int = 64, |
| do_sample: bool = False, |
| temperature: float = 1.0, |
| top_k: int = 0, |
| top_p: float = 1.0, |
| eos_token_id: Optional[int] = None, |
| pad_token_id: Optional[int] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| **kwargs, |
| ) -> torch.LongTensor: |
| """Generate with the recurrent state; supports greedy and top-k/top-p sampling.""" |
| if input_ids.ndim != 2 or input_ids.shape[1] == 0: |
| raise ValueError("input_ids must be a non-empty [batch, sequence] tensor") |
| if temperature <= 0: |
| raise ValueError("temperature must be > 0") |
| if kwargs: |
| unsupported = ", ".join(sorted(kwargs)) |
| raise TypeError(f"Unsupported generation options in ÇABA v0-core: {unsupported}") |
| self.eval() |
| eos = self.config.eos_token_id if eos_token_id is None else eos_token_id |
| pad = self.config.pad_token_id if pad_token_id is None else pad_token_id |
| sequences = input_ids |
| output = self(input_ids, use_cache=True) |
| cache = output.past_key_values |
| logits = output.logits[:, -1, :] |
| finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device) |
|
|
| for _ in range(max_new_tokens): |
| scores = logits.float() |
| if do_sample: |
| scores = scores / temperature |
| if top_k > 0: |
| k = min(int(top_k), scores.shape[-1]) |
| threshold = torch.topk(scores, k, dim=-1).values[:, -1, None] |
| scores = scores.masked_fill(scores < threshold, -float("inf")) |
| if 0.0 < top_p < 1.0: |
| sorted_scores, sorted_ids = torch.sort(scores, descending=True, dim=-1) |
| cumulative = torch.softmax(sorted_scores, dim=-1).cumsum(dim=-1) |
| remove = cumulative > top_p |
| remove[:, 1:] = remove[:, :-1].clone() |
| remove[:, 0] = False |
| sorted_scores = sorted_scores.masked_fill(remove, -float("inf")) |
| scores = torch.full_like(scores, -float("inf")).scatter( |
| 1, sorted_ids, sorted_scores |
| ) |
| next_token = torch.multinomial(torch.softmax(scores, dim=-1), num_samples=1) |
| else: |
| next_token = scores.argmax(dim=-1, keepdim=True) |
| if eos is not None: |
| next_token = torch.where( |
| finished[:, None], |
| torch.full_like(next_token, int(pad)), |
| next_token, |
| ) |
| sequences = torch.cat((sequences, next_token), dim=1) |
| if eos is not None: |
| finished |= next_token[:, 0].eq(int(eos)) |
| if bool(finished.all()): |
| break |
| output = self(next_token, past_key_values=cache, use_cache=True) |
| cache = output.past_key_values |
| logits = output.logits[:, -1, :] |
| return sequences |
|
|
