allow models to be pickled + test (#261)

mlx_lm/models/dbrx.py

@@ -105,10 +105,9 @@ class MLP(nn.Module):
         self.v1 = nn.Linear(d_model, ffn_dim, bias=False)
         self.w1 = nn.Linear(d_model, ffn_dim, bias=False)
         self.w2 = nn.Linear(ffn_dim, d_model, bias=False)
-        self.act_fn = nn.silu
 
     def __call__(self, x: mx.array) -> mx.array:
-        current_hidden_states = self.act_fn(self.w1(x)) * self.v1(x)
+        current_hidden_states = nn.silu(self.w1(x)) * self.v1(x)
         current_hidden_states = self.w2(current_hidden_states)
         return current_hidden_states
 

mlx_lm/models/deepseek.py

@@ -118,10 +118,9 @@ class DeepseekMLP(nn.Module):
         self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = nn.silu
 
     def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
 
 
 class MoEGate(nn.Module):

mlx_lm/models/deepseek_v3.py

@@ -130,6 +130,14 @@ def clipped_silu(x):
     return mx.clip(x * mx.sigmoid(x), a_min=-100, a_max=100)
 
 
+class ClippedSilu(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x):
+        return clipped_silu(x)
+
+
 class DeepseekV3Attention(nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
@@ -344,7 +352,7 @@ class DeepseekV3MoE(nn.Module):
             config.hidden_size,
             config.moe_intermediate_size,
             config.n_routed_experts,
-            activation=clipped_silu,
+            activation=ClippedSilu(),
         )
 
         self.gate = MoEGate(config)

mlx_lm/models/phi.py

@@ -111,10 +111,9 @@ class PhiMLP(nn.Module):
         super().__init__()
         self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
         self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.act = nn.GELU(approx="precise")
 
     def __call__(self, x) -> mx.array:
-        return self.fc2(self.act(self.fc1(x)))
+        return self.fc2(nn.gelu_approx(self.fc1(x)))
 
 
 class PhiDecoderLayer(nn.Module):

tests/test_models.py

@@ -1,4 +1,5 @@
 # Copyright © 2024 Apple Inc.
+import copy
 import unittest
 
 import mlx.core as mx
@@ -230,6 +231,9 @@ class TestModels(unittest.TestCase):
             self.assertEqual(outputs.shape, (1, 1, vocab_size))
             self.assertEqual(outputs.dtype, t)
 
+        # Make sure the model can be copied / pickled
+        copy.deepcopy(model)
+
     def test_llama(self):
         from mlx_lm.models import llama
 
@@ -256,15 +260,6 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
-    def test_phixtral(self):
-        from mlx_lm.models import phixtral
-
-        args = phixtral.ModelArgs(
-            "phixtral", num_vocab=1000, num_layers=4, model_dim=1024
-        )
-        model = phixtral.Model(args)
-        self.model_test_runner(model, args.model_type, args.num_vocab, args.num_layers)
-
     def test_phi3(self):
         from mlx_lm.models import phi3