applepie69/Nanbeige4.3-33B-MTP-head

MTP head for Nanbeige/Nanbeige4.3-33B

A small multi-token-prediction (MTP) head trained for Nanbeige/Nanbeige4.3-33B. It guesses the token after the next one, which a speculative-decoding runtime can use to produce several tokens per pass of the big model.

Status: experimental. The head is trained and its agreement with the base model is measured (below), but no speed-up has been measured and no inference engine can use it yet (see "What this is not"). This is not a finished model.

What this is

The base model has no built-in MTP weights, so this head was trained separately by self-distillation. It is not a full model: it needs the original Nanbeige/Nanbeige4.3-33B (it shares its embedding matrix and LM head).

Architecture (DeepSeek-V3 style, depth 1; the code is in mtp_head.py):

h_t      = final hidden state of the base model at position t (after model.norm)
e_{t+1}  = base embed_tokens(x_{t+1})
x        = eh_proj( concat[ enorm(e_{t+1}), hnorm(h_t) ] )     # (2*7168 -> 7168)
x        = LlamaDecoderLayer(x)                                # one extra decoder layer with its own attention + MLP
logits   = base lm_head( shared_head.norm(x) )                 # predicts token x_{t+2}

• Parameters in this file: 694 M (bf16, about 1.29 GB). The embedding matrix and LM head are not stored.
• Tensor names use the DeepSeek/GLM convention of an extra layer after the last block: model.layers.52.{enorm,hnorm,eh_proj}.weight, model.layers.52.shared_head.norm.weight and the usual Llama layer tensors (self_attn.*, mlp.*, input_layernorm, post_attention_layernorm).

Training

• Method: the base model is frozen. For every position the head sees (h_t, embedding of x_{t+1}) and is trained to match the base model's own output distribution for token x_{t+2} (KL divergence through the shared, frozen LM head). The head was initialised from a copy of the base model's last decoder layer.
• Data: chat, math reasoning and code text formatted with the base model's chat template and packed into rows of 1024 tokens:
• HuggingFaceH4/ultrachat_200k (split train_sft, about 40% of the tokens)
• open-r1/OpenR1-Math-220k (split train, about 40% of the tokens)
• bigcode/self-oss-instruct-sc2-exec-filter-50k (split train, about 20% of the tokens)
• Amount: 600,064 tokens of stored hidden states (576 rows of 1024 tokens for training, 10 held-out rows). 2 epochs, batch of 2 rows, peak learning rate 0.0002 (cosine schedule), 8-bit AdamW with fp16 autocast.
• Features come from the 4-bit (NF4) quantised base model, extracted on 2 x Tesla T4. The head therefore matches the 4-bit model best. With bf16 or 8-bit weights, expect somewhat lower agreement.

Results (held-out text, same kinds of data as training)

metric	value	meaning
top-1 agreement	0.9%	share of positions where the head's best guess for token t+2 equals the base model's best guess
top-3 contains top-1	1.9%	share where the base model's best token is among the head's three best
KL divergence	nan	average distance (nats) between the head's and the base model's distributions; lower is better
positions evaluated	10,220

Measured on held-out shards of the same text mixture, teacher-forced on dataset text. This is not measured on the base model's own generations and is not a speed benchmark. A common rule of thumb is that drafting needs roughly 50% acceptance to pay off. This is below that rule of thumb, so a real speed-up is unlikely without more training.

What this is not

• Not a standalone model and not a GGUF.
• No runtime support yet. llama.cpp's MTP speculative decoding (--spec-type draft-mtp) is implemented per model architecture, and as far as I could find the plain Llama architecture that Nanbeige uses has no MTP graph. vLLM and SGLang would need an equivalent model class. This repo gives the weights and a reference implementation to port from.
• No end-to-end speed or quality benchmark. In greedy speculative decoding the base model verifies every drafted token, so the output should match normal decoding; that has not been tested here.
• Short-context training. Training rows were at most 1024 tokens, so behaviour on much longer contexts is untested.
• The head only drafts one token ahead (depth 1). Chaining deeper drafts is untested.

Files

file	what
mtp_head.safetensors	the trained head (bf16)
mtp_config.json	training settings and the held-out numbers
mtp_head.py	stand-alone reference implementation (MTPHead, load_mtp_head)
for_merging/	config.json and model.safetensors.index.json patched to include the MTP layer, for people who have the base checkpoint locally

Using the reference implementation

import torch
from transformers import AutoConfig, AutoModelForCausalLM
from mtp_head import load_mtp_head

cfg = AutoConfig.from_pretrained("Nanbeige/Nanbeige4.3-33B")
head = load_mtp_head("mtp_head.safetensors", cfg, dtype=torch.float16, device="cuda")

# base = AutoModelForCausalLM.from_pretrained("Nanbeige/Nanbeige4.3-33B", ...)
# h = base.model(input_ids).last_hidden_state                  # (B, T, 7168), after the final norm
# emb = base.get_input_embeddings()(input_ids[:, 1:])           # embeddings of x_{t+1}
# z = head(h[:, :-1], emb)                                      # predicts x_{t+2} at each position
# draft_logits = base.lm_head(z)

License

Released under the same license as the base model (apache-2.0, see the base model card). Check the base model's terms before use.