| """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 |
| 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 |
| 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 |
|
|
