From 819cdb344946dc39318b6e48d122c065223542a5 Mon Sep 17 00:00:00 2001 From: Brian Wilfley Date: Wed, 22 Jul 2026 16:44:31 -0700 Subject: [PATCH 1/4] Fix transformers 5.14+ compatibility and remove dead import - torch_bitnet.py: change _tied_weights_keys from a list to a dict mapping lm_head.weight to model.embed_tokens.weight, required by the updated API in transformers 5.14+ - test_interop.py: remove unused import of training.bit_linear which does not exist in the repo and caused an immediate ModuleNotFoundError Closes #2 Co-Authored-By: Claude --- test_interop.py | 1 - torch_bitnet.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/test_interop.py b/test_interop.py index ee0e1bc..67de7dd 100644 --- a/test_interop.py +++ b/test_interop.py @@ -26,7 +26,6 @@ from torch_bitnet import BitnetForCausalLM as TorchBitnetForCausalLM from torch_bitnet import BitnetDecoderLayer as TorchBitnetDecoderLayer from transformers.activations import silu as torch_silu -from training.bit_linear import weight_quant as bit_linear_weight_quant class TestBitLinearInterop(unittest.TestCase): def setUp(self): diff --git a/torch_bitnet.py b/torch_bitnet.py index ef58099..a77456c 100644 --- a/torch_bitnet.py +++ b/torch_bitnet.py @@ -585,7 +585,7 @@ def _update_causal_mask(self, attention_mask, input_tensor, cache_position): class BitnetForCausalLM(BitnetPreTrainedModel): - _tied_weights_keys = ["lm_head.weight"] + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} def __init__(self, config): super().__init__(config) From 60d4d5913b3064b631259907b78ff10a24a80f6c Mon Sep 17 00:00:00 2001 From: Brian Wilfley Date: Wed, 22 Jul 2026 17:35:00 -0700 Subject: [PATCH 2/4] Fix lm_head weight key in convert.py and add MLX inference script MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removed the incorrect ("lm_head.", "lm_head.linear.") replacement pattern from convert.py — it was a leftover from another model's converter and caused load_causal_model to fail with "no parameter named linear". Added run_mlx.py to run text generation with the converted MLX model, with streaming output and correct SentencePiece token decoding. Co-Authored-By: Claude --- convert.py | 1 - run_mlx.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 run_mlx.py diff --git a/convert.py b/convert.py index bfc600f..e34a4d0 100644 --- a/convert.py +++ b/convert.py @@ -8,7 +8,6 @@ (".q.", ".query_proj."), (".v.", ".value_proj."), ("shared.", "wte."), - ("lm_head.", "lm_head.linear."), (".layer.0.layer_norm.", ".ln1."), (".layer.1.layer_norm.", ".ln2."), (".layer.2.layer_norm.", ".ln3."), diff --git a/run_mlx.py b/run_mlx.py new file mode 100644 index 0000000..ec61ab7 --- /dev/null +++ b/run_mlx.py @@ -0,0 +1,31 @@ +import argparse +import mlx.core as mx +from mlx_bitnet import load_causal_model + +parser = argparse.ArgumentParser() +parser.add_argument("--model", default="1bitLLM/bitnet_b1_58-xl", type=str) +parser.add_argument("--prompt", default="The meaning of life is", type=str) +parser.add_argument("--max_tokens", default=50, type=int) +parser.add_argument("--temp", default=0.0, type=float) +args = parser.parse_args() + +print(f"Loading model {args.model} ...") +model, tokenizer = load_causal_model(args.model) + +tokens = tokenizer.encode(args.prompt) +input_ids = mx.array([tokens]) +attention_mask = mx.ones_like(input_ids) + +generated = list(tokens) +decoded_so_far = args.prompt +print(decoded_so_far, end="", flush=True) + +for token in model.generate(input_ids, attention_mask, temp=args.temp): + generated.append(token.item()) + new_text = tokenizer.decode(generated) + print(new_text[len(decoded_so_far):], end="", flush=True) + decoded_so_far = new_text + if len(generated) - len(tokens) >= args.max_tokens: + break + +print() From acd781bd2c79d5accaff1de0ba23fe3aec7e3756 Mon Sep 17 00:00:00 2001 From: Brian Wilfley Date: Wed, 22 Jul 2026 17:39:31 -0700 Subject: [PATCH 3/4] Fix sanitize_config to pass num_hidden_layers from model config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit num_hidden_layers was missing from sanitize_config, so it always defaulted to 24 — causing load_causal_model to fail for models with a different layer count (e.g. the 3B model has 26 layers). Co-Authored-By: Claude --- mlx_bitnet.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mlx_bitnet.py b/mlx_bitnet.py index 6eea92f..b02c482 100644 --- a/mlx_bitnet.py +++ b/mlx_bitnet.py @@ -846,6 +846,7 @@ def sanitize_config(_config: BitnetConfig) -> MinimalBitnetConfig: intermediate_size=_config.intermediate_size, max_position_embeddings=_config.max_position_embeddings, num_attention_heads=_config.num_attention_heads, + num_hidden_layers=_config.num_hidden_layers, num_key_value_heads=_config.num_key_value_heads, pad_token_id=_config.pad_token_id, rms_norm_eps=_config.rms_norm_eps, From 911db207e2a59fb6b3558b01f63217d9a08f1765 Mon Sep 17 00:00:00 2001 From: Brian Wilfley Date: Wed, 22 Jul 2026 17:42:55 -0700 Subject: [PATCH 4/4] Fix first-token spacing glitch in run_mlx.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Used tokenizer.decode(tokens) as the baseline for delta computation instead of the raw prompt string — the two can differ at token boundaries in SentencePiece, causing the last word of the prompt to appear duplicated in the first generated token. Co-Authored-By: Claude --- run_mlx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/run_mlx.py b/run_mlx.py index ec61ab7..93543aa 100644 --- a/run_mlx.py +++ b/run_mlx.py @@ -17,8 +17,8 @@ attention_mask = mx.ones_like(input_ids) generated = list(tokens) -decoded_so_far = args.prompt -print(decoded_so_far, end="", flush=True) +decoded_so_far = tokenizer.decode(tokens) +print(args.prompt, end="", flush=True) for token in model.generate(input_ids, attention_mask, temp=args.temp): generated.append(token.item())