File size: 5,486 Bytes
338bb9c | 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 | from __future__ import annotations
import json
from pathlib import Path
from typing import Any
try:
import mlx.core as mx
import mlx.nn as nn
except ImportError: # Linux export validation intentionally has no MLX runtime.
mx = None
nn = None
if nn is not None and mx is not None:
class RMSNorm(nn.Module):
def __init__(self, dimension: int, epsilon: float = 1e-6) -> None:
super().__init__()
self.weight = mx.ones((dimension,))
self.epsilon = epsilon
def __call__(self, values: Any) -> Any:
normalized = values * mx.rsqrt(
mx.mean(mx.square(values), axis=-1, keepdims=True) + self.epsilon
)
return normalized * self.weight
def _rotate_half(values: Any) -> Any:
first, second = mx.split(values, 2, axis=-1)
return mx.concatenate((-second, first), axis=-1)
def apply_rope(query: Any, key: Any) -> tuple[Any, Any]:
sequence_length = query.shape[-2]
dimension = query.shape[-1]
positions = mx.arange(sequence_length, dtype=mx.float32)
frequencies = 1.0 / (10000 ** (mx.arange(0, dimension, 2, dtype=mx.float32) / dimension))
angles = positions[:, None] * frequencies[None, :]
angles = mx.concatenate((angles, angles), axis=-1)[None, None, :, :]
cosine = mx.cos(angles)
sine = mx.sin(angles)
return query * cosine + _rotate_half(query) * sine, key * cosine + _rotate_half(key) * sine
class CausalSelfAttention(nn.Module):
def __init__(self, config: dict[str, Any]) -> None:
super().__init__()
self.heads = int(config["attention_heads"])
self.head_dimension = int(config["model_width"]) // self.heads
width = int(config["model_width"])
self.qkv = nn.Linear(width, 3 * width, bias=False)
self.output = nn.Linear(width, width, bias=False)
def __call__(self, values: Any) -> Any:
batch, sequence, width = values.shape
qkv = self.qkv(values).reshape(batch, sequence, 3, self.heads, self.head_dimension)
qkv = qkv.transpose(2, 0, 3, 1, 4)
query, key, value = qkv[0], qkv[1], qkv[2]
query, key = apply_rope(query, key)
mask = nn.MultiHeadAttention.create_additive_causal_mask(sequence)
attended = mx.fast.scaled_dot_product_attention(
query,
key,
value,
scale=self.head_dimension**-0.5,
mask=mask,
)
return self.output(attended.transpose(0, 2, 1, 3).reshape(batch, sequence, width))
class SwiGLU(nn.Module):
def __init__(self, config: dict[str, Any]) -> None:
super().__init__()
width = int(config["model_width"])
ffn_width = int(config["ffn_width"])
self.gate = nn.Linear(width, ffn_width, bias=False)
self.up = nn.Linear(width, ffn_width, bias=False)
self.down = nn.Linear(ffn_width, width, bias=False)
def __call__(self, values: Any) -> Any:
return self.down(nn.silu(self.gate(values)) * self.up(values))
class TransformerBlock(nn.Module):
def __init__(self, config: dict[str, Any]) -> None:
super().__init__()
width = int(config["model_width"])
self.attention_norm = RMSNorm(width)
self.attention = CausalSelfAttention(config)
self.ffn_norm = RMSNorm(width)
self.ffn = SwiGLU(config)
def __call__(self, values: Any) -> Any:
values = values + self.attention(self.attention_norm(values))
return values + self.ffn(self.ffn_norm(values))
class TransformerRouterMLX(nn.Module):
def __init__(self, config: dict[str, Any]) -> None:
super().__init__()
width = int(config["model_width"])
self.config = config
self.embedding = nn.Embedding(int(config["vocab_size"]), width)
self.blocks = [TransformerBlock(config) for _ in range(int(config["layers"]))]
self.norm = RMSNorm(width)
self.pass_head = nn.Linear(width, int(config["candidate_count"]))
self.score_head = nn.Linear(width, int(config["candidate_count"]))
def __call__(self, input_ids: Any) -> dict[str, Any]:
hidden = self.embedding(input_ids)
for block in self.blocks:
hidden = block(hidden)
hidden = self.norm(hidden)
route_state = hidden[:, -1]
return {
"pass_logits": self.pass_head(route_state),
"scores": mx.sigmoid(self.score_head(route_state)),
}
def load_mlx_router(directory: str | Path) -> Any:
root = Path(directory)
config = json.loads((root / "config.json").read_text(encoding="utf-8"))["model"]
model = TransformerRouterMLX(config)
weights = mx.load(str(root / "model.safetensors"))
model.load_weights(list(weights.items()), strict=True)
mx.eval(model.parameters())
return model
else:
class TransformerRouterMLX: # type: ignore[no-redef]
def __init__(self, config: dict[str, Any]) -> None:
raise RuntimeError("MLX requires macOS on Apple Silicon")
def load_mlx_router(directory: str | Path) -> Any:
raise RuntimeError("MLX requires macOS on Apple Silicon")
|