File size: 6,469 Bytes
980b1b3 | 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 | """Small MLX inference implementation for Cactus Needle."""
from __future__ import annotations
import json
import math
from pathlib import Path
from typing import Any
import mlx.core as mx
class NeedleModel:
def __init__(self, weights: dict[str, mx.array], config: dict[str, Any]):
self.weights = weights
self.config = config
self.num_heads = int(config["num_heads"])
self.num_kv_heads = int(config["num_kv_heads"])
self.d_model = int(config["d_model"])
self.head_dim = self.d_model // self.num_heads
@classmethod
def from_pretrained(cls, path: str | Path) -> "NeedleModel":
path = Path(path)
config = json.loads((path / "config.json").read_text())
weights = mx.load(str(path / "model.safetensors"))
return cls(weights, config)
def _weight(self, name: str, layer: int | None = None) -> mx.array:
value = self.weights[name]
return value if layer is None else value[layer]
@staticmethod
def _linear(x: mx.array, kernel: mx.array, bias: mx.array | None = None) -> mx.array:
output = mx.matmul(x, kernel)
return output if bias is None else output + bias
@staticmethod
def _zcrms_norm(x: mx.array, scale: mx.array) -> mx.array:
rms = mx.sqrt(mx.mean(mx.square(x.astype(mx.float32)), axis=-1, keepdims=True) + 1e-6)
return ((1.0 + scale) * x / rms).astype(x.dtype)
def _rope(self, x: mx.array) -> mx.array:
sequence_length = x.shape[2]
half_dim = self.head_dim // 2
positions = mx.arange(sequence_length, dtype=mx.float32)
inv_freq = 1.0 / (
float(self.config["rope_theta"])
** (mx.arange(0, self.head_dim, 2, dtype=mx.float32) / self.head_dim)
)
angles = mx.outer(positions, inv_freq)
cos = mx.reshape(mx.cos(angles), (1, 1, sequence_length, half_dim))
sin = mx.reshape(mx.sin(angles), (1, 1, sequence_length, half_dim))
first, second = x[..., :half_dim], x[..., half_dim:]
return mx.concatenate([first * cos - second * sin, second * cos + first * sin], axis=-1)
def _attention(
self,
q_input: mx.array,
kv_input: mx.array,
prefix: str,
layer: int,
mask: mx.array | None,
rope: bool,
) -> mx.array:
q = self._linear(q_input, self._weight(f"{prefix}/q_proj/kernel", layer))
k = self._linear(kv_input, self._weight(f"{prefix}/k_proj/kernel", layer))
v = self._linear(kv_input, self._weight(f"{prefix}/v_proj/kernel", layer))
batch = q.shape[0]
q = mx.transpose(mx.reshape(q, (batch, -1, self.num_heads, self.head_dim)), (0, 2, 1, 3))
k = mx.transpose(mx.reshape(k, (batch, -1, self.num_kv_heads, self.head_dim)), (0, 2, 1, 3))
v = mx.transpose(mx.reshape(v, (batch, -1, self.num_kv_heads, self.head_dim)), (0, 2, 1, 3))
q = self._zcrms_norm(q, self._weight(f"{prefix}/q_norm/scale", layer))
k = self._zcrms_norm(k, self._weight(f"{prefix}/k_norm/scale", layer))
if self.num_heads != self.num_kv_heads:
repeats = self.num_heads // self.num_kv_heads
k = mx.repeat(k, repeats, axis=1)
v = mx.repeat(v, repeats, axis=1)
if rope:
q, k = self._rope(q), self._rope(k)
scores = mx.matmul(q.astype(mx.float32), mx.transpose(k.astype(mx.float32), (0, 1, 3, 2)))
scores = scores / math.sqrt(self.head_dim)
if mask is not None:
scores = mx.where(mask, scores, mx.array(-1e9, dtype=mx.float32))
output = mx.matmul(mx.softmax(scores, axis=-1).astype(v.dtype), v)
output = mx.reshape(mx.transpose(output, (0, 2, 1, 3)), (batch, -1, self.d_model))
return self._linear(output, self._weight(f"{prefix}/out_proj/kernel", layer))
def encode(self, source_tokens: mx.array) -> tuple[mx.array, mx.array]:
source_tokens = mx.array(source_tokens, dtype=mx.int32)
padding_mask = source_tokens != int(self.config["pad_token_id"])
mask = mx.reshape(padding_mask, (source_tokens.shape[0], 1, 1, source_tokens.shape[1]))
x = self.weights["embedding/embedding"][source_tokens] * math.sqrt(self.d_model)
for layer in range(int(self.config["num_encoder_layers"])):
prefix = "encoder/layers/EncoderBlock_0"
residual = x
x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_0/scale", layer))
x = self._attention(x, x, f"{prefix}/self_attn", layer, mask, rope=True)
gate = mx.sigmoid(self._weight(f"{prefix}/attn_gate", layer))
x = residual + gate * x
return self._zcrms_norm(x, self.weights["encoder/final_norm/scale"]), mask
def decode(self, target_tokens: mx.array, encoder_out: mx.array, cross_mask: mx.array) -> mx.array:
target_tokens = mx.array(target_tokens, dtype=mx.int32)
length = target_tokens.shape[1]
causal = mx.tril(mx.ones((length, length), dtype=mx.bool_))
target_valid = target_tokens != int(self.config["pad_token_id"])
self_mask = mx.reshape(causal, (1, 1, length, length)) & mx.reshape(
target_valid, (target_tokens.shape[0], 1, 1, length)
)
x = self.weights["embedding/embedding"][target_tokens] * math.sqrt(self.d_model)
prefix = "decoder/layers/DecoderBlock_0"
for layer in range(int(self.config["num_decoder_layers"])):
residual = x
x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_0/scale", layer))
x = self._attention(x, x, f"{prefix}/self_attn", layer, self_mask, rope=True)
x = residual + mx.sigmoid(self._weight(f"{prefix}/self_attn_gate", layer)) * x
residual = x
x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_1/scale", layer))
x = self._attention(x, encoder_out, f"{prefix}/cross_attn", layer, cross_mask, rope=False)
x = residual + mx.sigmoid(self._weight(f"{prefix}/cross_attn_gate", layer)) * x
x = self._zcrms_norm(x, self.weights["decoder/ZCRMSNorm_0/scale"])
return mx.matmul(x.astype(mx.float32), mx.transpose(self.weights["embedding/embedding"].astype(mx.float32)))
def forward(self, source_tokens: mx.array, target_tokens: mx.array) -> mx.array:
encoder_out, cross_mask = self.encode(source_tokens)
return self.decode(target_tokens, encoder_out, cross_mask)
|