Port of bitnet1.58 with custom metal kernel (#331)

* Port mlx-lm bitnet1.58 ml-explore/mlx-lm#219

* update: Update relu2 function to use compile for shapeless input

* refactor: Update BitLinear & Format code

* update: Add quantization parameters to BaseConfiguration

* update: Improve error handling during weight pre-loading in ContentView

* update: Add bitnet_b1_58_2b_4t_4bit model configuration to LLMModelFactory

* update: Rename relu2 function to reluSquared and refactor implementation

* update: ACKNOWLEDGMENTS.md

* Improve the bitnet kernel

* remove: eliminate reluSquared function from Bitnet.swift

* refactor: update kernel

Author: johnmai-dev

ACKNOWLEDGMENTS.md

@@ -1,14 +1,16 @@
 # Individual Contributors
 
-If you wish to be acknowledged for your contributions, please list your name
-with a short description of your contribution(s) below. For example:
-
-- Jane Smith: Added the `foo` and `bar` ops.
+> If you wish to be acknowledged for your contributions, please list your name
+> with a short description of your contribution(s) below. For example:
+> - Jane Smith: Added the `foo` and `bar` ops.
 
 MLX Swift was developed with contributions from the following individuals:
 
+- [John Mai](https://github.com/johnmai-dev): Added support for multiple models (Qwen2, Starcoder2, InternLM2, Qwen3, Qwen3 MoE, GLM-4, MiMo, BitNet).
+
+
 <a href="https://github.com/ml-explore/mlx-swift-examples/graphs/contributors">
-  <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/-swift-examples&anon=0&columns=20&max=100&r=true" />
+  <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx-swift-examples&anon=0&columns=20&max=100&r=true" />
 </a>
 
 

Applications/LLMEval/ContentView.swift

@@ -133,8 +133,12 @@ struct ContentView: View {
 
         }
         .task {
-            // pre-load the weights on launch to speed up the first generation
-            _ = try? await llm.load()
+            do {
+                // pre-load the weights on launch to speed up the first generation
+                _ = try await llm.load()
+            } catch {
+                llm.output = "Failed: \(error)"
+            }
         }
     }
 

Libraries/MLXLLM/LLMModelFactory.swift

@@ -48,6 +48,7 @@ public class LLMTypeRegistry: ModelTypeRegistry, @unchecked Sendable {
             "mimo": create(MiMoConfiguration.self, MiMoModel.init),
             "glm4": create(GLM4Configuration.self, GLM4Model.init),
             "acereason": create(Qwen2Configuration.self, Qwen2Model.init),
+            "bitnet": create(BitnetConfiguration.self, BitnetModel.init),
         ]
     }
 
@@ -225,6 +226,11 @@ public class LLMRegistry: AbstractModelRegistry, @unchecked Sendable {
         defaultPrompt: ""
     )
 
+    static public let bitnet_b1_58_2b_4t_4bit = ModelConfiguration(
+        id: "mlx-community/bitnet-b1.58-2B-4T-4bit",
+        defaultPrompt: "Why is the sky blue?"
+    )
+
     private static func all() -> [ModelConfiguration] {
         [
             codeLlama13b4bit,
@@ -256,6 +262,7 @@ public class LLMRegistry: AbstractModelRegistry, @unchecked Sendable {
             mimo_7b_sft_4bit,
             glm4_9b_4bit,
             acereason_7b_4bit,
+            bitnet_b1_58_2b_4t_4bit,
         ]
     }