File size: 3,514 Bytes
e72cbe1 | 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 | """Cortex-1 en MLX (Apple Silicon) — equivalente a la implementación NumPy/PyTorch.
Misma arquitectura: GPT pre-norm, RMSNorm, atención causal con QKV fusionado,
SwiGLU, embeddings de posición aprendidos, lm_head atado al embedding.
Paridad verificada: max|Δ logits| < 1e-4 vs el modelo NumPy original.
Uso: python generate_mlx.py "La capital de Francia es"
"""
from __future__ import annotations
import json
import os
import mlx.core as mx
import mlx.nn as nn
MODEL_FILE = "weights.safetensors"
class Attention(nn.Module):
def __init__(self, d: int, n_heads: int):
super().__init__()
self.n_heads = n_heads
self.d_head = d // n_heads
self.scale = self.d_head ** -0.5
self.qkv = nn.Linear(d, 3 * d, bias=False)
self.proj = nn.Linear(d, d, bias=False)
def __call__(self, x: mx.array) -> mx.array:
B, T, d = x.shape
qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.d_head)
q, k, v = (qkv[:, :, i].transpose(0, 2, 1, 3) for i in range(3))
att = (q @ k.transpose(0, 1, 3, 2)) * self.scale # (B,H,T,T)
mask = mx.triu(mx.full((T, T), -mx.inf, dtype=att.dtype), k=1)
att = mx.softmax(att + mask, axis=-1)
y = (att @ v).transpose(0, 2, 1, 3).reshape(B, T, d)
return self.proj(y)
class MLP(nn.Module):
def __init__(self, d: int, d_ff: int):
super().__init__()
self.gate_proj = nn.Linear(d, d_ff, bias=False)
self.up_proj = nn.Linear(d, d_ff, bias=False)
self.down_proj = nn.Linear(d_ff, d, bias=False)
def __call__(self, x: mx.array) -> mx.array:
return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
class Block(nn.Module):
def __init__(self, d: int, n_heads: int, d_ff: int, eps: float):
super().__init__()
self.input_layernorm = nn.RMSNorm(d, eps=eps)
self.self_attn = Attention(d, n_heads)
self.post_attention_layernorm = nn.RMSNorm(d, eps=eps)
self.mlp = MLP(d, d_ff)
def __call__(self, x):
x = x + self.self_attn(self.input_layernorm(x))
return x + self.mlp(self.post_attention_layernorm(x))
class Model(nn.Module):
def __init__(self, cfg: dict):
super().__init__()
d = cfg["hidden_size"]
self.cfg = cfg
self.embed_tokens = nn.Embedding(cfg["vocab_size"], d)
self.embed_positions = nn.Embedding(cfg["max_position_embeddings"], d)
self.layers = [Block(d, cfg["num_attention_heads"],
cfg["intermediate_size"], cfg["rms_norm_eps"])
for _ in range(cfg["num_hidden_layers"])]
self.norm = nn.RMSNorm(d, eps=cfg["rms_norm_eps"])
def __call__(self, input_ids: mx.array) -> mx.array:
T = input_ids.shape[1]
h = self.embed_tokens(input_ids) + self.embed_positions(mx.arange(T))
for layer in self.layers:
h = layer(h)
return self.norm(h)
def logits(self, input_ids: mx.array) -> mx.array:
"""(B,T) int -> (B,T,V); lm_head atado = embed_tokens."""
h = self(input_ids)
w = self.embed_tokens.weight # (V, d)
return h @ w.T
@classmethod
def load(cls, repo: str) -> "Model":
with open(os.path.join(repo, "config.json")) as f:
cfg = json.load(f)
model = cls(cfg)
weights = mx.load(os.path.join(repo, MODEL_FILE))
model.load_weights(list(weights.items()), strict=True)
return model
|