johnhalloran/Nanbeige4.2-3B-mps-fix

Nanbeige4.2-3B — Apple Silicon / transformers compatibility fix

This is Nanbeige/Nanbeige4.2-3B (base revision 5d54321e9e01e0d026f8e371046678fc384dca39) with five bugs fixed in its custom modeling code and baked into the checkpoint's weights. Independent project — not affiliated with or endorsed by the Nanbeige team. All credit for the architecture, training, and base weights belongs to them; see their model card and technical report.

Independent project — not affiliated with or endorsed by the Nanbeige team. All credit for the architecture, training, and base weights belongs to them; see their model card and technical report.

Paper: arXiv:2608.13987

Full harness, paper, and evaluation scripts: github.com/johnhalloran321/Nanbeige4.2-3B-mps-fix — covers two further issues (a memory ceiling and a system-prompt regression) that this checkpoint alone does not fix; see below.

The short version

Loading the stock checkpoint through generic HF transformers (trust_remote_code=True) crashes hard on Apple Silicon (MPS) a few generation steps in, and — even patched around — produces structurally incoherent output (looping </think> tokens, character-level word salad) regardless of dtype, sampling settings, or enable_thinking. Five distinct transformers-version compatibility bugs are responsible; none are security issues, and none are the model's fault in any deep sense — its custom code was written against an older transformers API surface that has since moved on. Full root-cause writeup, exact evidence, and the elimination methodology for each: MPS_FIX_NOTES.md.

The dominant one: NanbeigeRotaryEmbedding's inv_freq buffer was being silently zeroed on load (persistent=False + meta-device init means it's never restored from the checkpoint) — meaning RoPE was contributing zero positional information regardless of any other setting. That single bug explains nearly all of the incoherence; the other four (a RoPE-config dispatch KeyError, a Cache API sentinel mismatch, the position_ids re-trim that caused the actual MPS crash, and a _tied_weights_keys format issue that breaks save_pretrained) are all real but secondary.

Also worth knowing before citing this architecture: the released checkpoint's config.json disables every one of the model's more novel advertised features (LoopSplit, manifold-constrained hyper-connections/mHC, depth attention, n-gram embeddings) — what actually runs is standard GQA attention plus a weight-shared loop over 22 layers, executed twice. Nanbeige's own official Ollama/MLX serving path (Nanbeige/ollama, nanbeige42 branch) independently confirms this: it implements only plain attention and the loop-repeat, nothing else. Details in MPS_FIX_NOTES.md.

What this checkpoint alone does not fix

The five bugs above are baked into these weights and this repo's config files — nothing further is needed for them. Two other issues, covered in the paper and fixed in the GitHub repo's harness, are not fixed here, because neither one is something that can be baked into a checkpoint:

• A memory ceiling from the Looped Transformer's doubled attention cost. Naive prefill hard-crashes by roughly 9,000 tokens even at batch size 1. The fix (chunked prefill) is a serving-time strategy — see harness/nanbeige_harness_server.py in the repo above.
• A chat-template regression that silently discards its own tool-use system prompt. Supply any system message — which most agent frameworks do — and multi-tool-call output breaks. The fix splices the caller's system content into the rendered template output after the fact, so it can't live inside chat_template.jinja itself; this repo's copy is unmodified from the base checkpoint. Fixed in the same harness above.

If you use this checkpoint directly via plain transformers (below) with a system message or a long context, you'll hit both. Use the harness in the GitHub repo, or port the fix yourself — exact code and a full write-up are there.

Usage

Works exactly like the base model — no runtime patching required, no transformers version pin:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "johnhalloran/Nanbeige4.2-3B-mps-fix"
tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True,
)

messages = [{"role": "user", "content": "Which number is bigger, 9.11 or 9.8?"}]
text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tok([text], return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=256)
print(tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Verified on Apple Silicon (MPS) and CPU. Not yet validated under vLLM's Transformers backend — see the note at the end of MPS_FIX_NOTES.md.

Acknowledgments

• Nanbeige — the model, training, and base weights.
• jishnuvenugopal/nanbeige-mlx — an independent MLX port whose documentation of the inv_freq buffer's persistence semantics was the lead that found this repo's dominant bug. Unaffiliated with Nanbeige or this repo.

License

Base model and weights: Apache-2.0, per upstream. Code changes here: MIT. Changes from the base revision are fully documented in MPS_FIX_NOTES.md per Apache-2.0 §4.