Skip to content

Add Nanbeige (Nanbeige4.2 looped transformer) model support - #1597

Open
MercuriusDream wants to merge 3 commits into
ml-explore:mainfrom
MercuriusDream:add-nanbeige-model
Open

Add Nanbeige (Nanbeige4.2 looped transformer) model support#1597
MercuriusDream wants to merge 3 commits into
ml-explore:mainfrom
MercuriusDream:add-nanbeige-model

Conversation

@MercuriusDream

Copy link
Copy Markdown

Adds the nanbeige model type used by Nanbeige/Nanbeige4.2-3B.

The architecture is a llama-style decoder executed num_loops times with shared weights ("Looped Transformer"). Each loop gets its own KV cache entries and the final RMSNorm is applied at the end of each loop (or once at the end when skip_loop_final_norm is set).

Correctness:

  • Logits match the reference modeling_nanbeige.py (fp32, CPU) to a max abs diff of 9.5e-05 over the full 166k vocab.
  • Cached token-by-token decoding matches single-shot prefill.
  • Verified generation for bf16 and 2/3/4/5/6/8-bit quants on an M-series Mac.

Config flags of the reference implementation not used by the released checkpoint (enable_double_loop_split, loop_share_kv, enable_depth_attention) raise NotImplementedError instead of loading silently. Checkpoints using n-gram embeddings or mHC fail naturally on unexpected weight keys.

Added a test_nanbeige case to tests/test_models.py; the suite passes apart from the pre-existing test_ssm failure on main.

@MercuriusDream

Copy link
Copy Markdown
Author

Is anyone there?

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would mind rewriting the PR description following our new template?
https://github.com/ml-explore/mlx-lm/blob/main/.github/PULL_REQUEST_TEMPLATE/new_model.md?plain=1

It would help me verify the implementation much quicker.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants