diff --git a/mlx_lm/models/bitlinear_layers.py b/mlx_lm/models/bitlinear_layers.py index 043b489..94842c6 100644 --- a/mlx_lm/models/bitlinear_layers.py +++ b/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(): diff --git a/mlx_lm/models/bitnet.py b/mlx_lm/models/bitnet.py index 6020b38..0b2ba86 100644 --- a/mlx_lm/models/bitnet.py +++ b/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: diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index 862e7ca..8002c8d 100644 --- a/mlx_lm/utils.py +++ b/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)