Feat: add falcon-e support for bitnet models (#268)

* add falcon-e support for bitnet models

* add comments for clarity

* aaddress offline comments

* Update mlx_lm/models/bitlinear_layers.py

Co-authored-by: Awni Hannun <[email protected]>

* address comments

* nits

---------

Co-authored-by: Awni Hannun <[email protected]>
Co-authored-by: Awni Hannun <[email protected]>

mlx_lm/models/bitlinear_layers.py

@@ -3,6 +3,33 @@
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.nn.layers.quantized import QuantizedLinear
+from mlx.utils import tree_flatten, tree_unflatten
+
+
+def bitnet_quantize(model, quantization_config: dict):
+    quantize_layers = []
+    modules_to_not_convert = quantization_config.get("modules_to_not_convert", [])
+    invert_weight_scales = (
+        quantization_config.get("linear_class", "") != "autobitlinear"
+    )
+
+    for name, module in tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module):
+
+        # Replace nn.Linear layers, but skip any layer from the `modules_to_not_convert` list
+        if name not in modules_to_not_convert and isinstance(module, nn.Linear):
+            old_weight = module.weight
+            out_features, in_features = old_weight.shape
+            bias = "bias" in module
+            new_layer = BitLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                invert_weight_scales=invert_weight_scales,
+            )
+            quantize_layers.append((name, new_layer))
+    if len(quantize_layers) > 0:
+        model.update_modules(tree_unflatten(quantize_layers))
+    return model
 
 
 def make_bitlinear_kernel():

mlx_lm/models/bitnet.py

@@ -112,6 +112,7 @@ class MLP(nn.Module):
         self.gate_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
         self.down_proj = BitLinear(hidden_dim, dim, bias=mlp_bias)
         self.up_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
+
         self.ffn_sub_norm = nn.RMSNorm(args.intermediate_size, eps=args.rms_norm_eps)
 
     def __call__(self, x) -> mx.array:

mlx_lm/utils.py

@@ -202,6 +202,15 @@ def load_model(
             bits=quantization["bits"],
             class_predicate=class_predicate,
         )
+    elif quantization_config := config.get("quantization_config", False):
+        # Handle legacy quantization config
+        quant_method = quantization_config["quant_method"]
+        if quant_method == "bitnet":
+            from .models.bitlinear_layers import bitnet_quantize
+
+            model = bitnet_quantize(model, quantization_config)
+        else:
+            raise ValueError(f"Unsupported quantization method {quant_method}")
 
     model.load_weights(list(weights.items()), strict=strict)