openlm

llms/mlx_lm/models/openlm.py

@@ -0,0 +1,183 @@
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from .base import BaseModelArgs, create_additive_causal_mask
+
+
+@dataclass
+class ParamsArgs(BaseModelArgs):
+    dim: int
+    ffn_type: str
+    n_heads: int
+    n_layers: int
+    norm_eps: float
+    positional_embedding_type: str
+    post_embed_norm: bool
+    qk_norm: bool
+    vocab_size: int
+    weight_tying: bool
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    params_args_dict: ParamsArgs
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.dim = args.dim
+        self.n_heads = args.n_heads
+        self.head_dim = self.dim // self.n_heads
+        self.qk_norm = args.qk_norm
+        self.scale = self.head_dim**-0.5
+
+        self.in_proj = nn.Linear(self.dim, 3 * self.dim, bias=False)
+        self.out_proj = nn.Linear(self.dim, self.dim, bias=False)
+        if self.qk_norm:
+            self.q_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
+            self.k_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
+        self.rope = nn.RoPE(
+            self.head_dim,
+            traditional=False,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.in_proj(x).split(3, axis=-1)
+
+        if self.qk_norm:
+            queries = self.q_norm(queries)
+            keys = self.q_norm(keys)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.out_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        # https://github.com/mlfoundations/open_lm/blob/c65b43042ff31c0fe26f930decf1ccab1b03ab4b/open_lm/model.py#L254C2-L254C3
+        hidden_dim = 256 * ((int(2 * 4 * args.dim / 3) + 256 - 1) // 256)
+        self.w12 = nn.Linear(args.dim, 2 * hidden_dim, bias=False)
+        self.w3 = nn.Linear(hidden_dim, args.dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        gate, x = self.w12(x).split(2, axis=-1)
+        return self.w3(nn.silu(gate) * x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.attention = Attention(args)
+        self.feed_forward = MLP(args)
+        self.ffn_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
+        self.attention_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+        r = self.attention(self.attention_norm(x), mask, cache)
+        h = x + r
+        r = self.feed_forward(self.ffn_norm(h))
+        out = h + r
+        return out
+
+
+class OpenLM(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)
+        self.layers = [TransformerBlock(args=args) for _ in range(args.n_layers)]
+        self.norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
+        self.output = nn.Linear(args.dim, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        _, L = inputs.shape
+
+        h = self.tok_embeddings(inputs)
+
+        mask = None
+        if h.shape[1] > 1:
+            mask = create_additive_causal_mask(
+                h.shape[1], cache[0].offset if cache is not None else 0
+            )
+            mask = mask.astype(h.dtype)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.output(self.norm(h))
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        args.params_args_dict = ParamsArgs.from_dict(args.params_args_dict)
+        self.args = args.params_args_dict
+        self.model_type = args.model_type
+        self.model = OpenLM(self.args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, cache)
+        return out
+
+    def sanitize(self, weights):
+        # Remove unused precomputed rotary freqs
+        return {k: v for k, v in weights.items() if "inv_freq" not in k}
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def head_dim(self):
+        return self.args.dim // self.args.n_heads
+
+    @property
+    def n_kv_heads(self):
+        return self.args.n_heads

llms/mlx_lm/tuner/utils.py

@@ -122,8 +122,10 @@ def linear_to_lora_layers(
         keys = set(["norm_attn_norm.attn.Wqkv", "ffn.router.layer"])
     elif model.model_type == "internlm2":
         keys = set(["attention.wqkv", "attention.wo"])
+    elif model.model_type == "openlm":
+        keys = set(["attention.in_proj", "attention.out_proj"])
     else:
-        raise ValueError(f"Lora does not support {model.model_type}")
+        raise ValueError(f"LoRA does not support {model.model_type}")
 
     for l in model.layers[num_layers - num_lora_layers :]:
         lora_layers = [(k, to_lora(m)) for k, m in l.named_modules() if k in keys]