"""Hybrid mamba2/attention model from NVIDIA"""
model_arch = gguf.MODEL_ARCH.NEMOTRON_H
is_moe: bool = False
+ supports_mtp_export = True
def __init__(self, *args, **kwargs):
# We have to determine the correct model architecture (MoE vs non-MoE) before
self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"]
self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"]
+ # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head
+ self._mtp_bid: int | None = None
+ if self.is_moe and not self.no_mtp:
+ n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0
+ if n_nextn > 0:
+ assert n_nextn == 1, (
+ "NemotronH MTP conversion currently supports num_nextn_predict_layers == 1"
+ )
+ self._mtp_bid = self.block_count
+ self.block_count += 1
+ # The folded MTP block carries both an attention sub-layer and a
+ # MoE sub-layer, so register it as both so the per-layer metadata arrays cover it
+ self._attn_layers.append(self._mtp_bid)
+ self._mlp_layers.append(self._mtp_bid)
+ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+ if self.mtp_only and self._mtp_bid is None:
+ raise ValueError("--mtp was requested, but this model does not contain a supported MTP head")
+
def get_attn_layers(self):
pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")
if pattern is None:
return [i for i, val in enumerate(pattern) if val == "attention"]
+ @classmethod
+ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+ name, gen = item
+ if name.startswith("mtp."):
+ # --no-mtp: drop the MTP head entirely
+ if cls.no_mtp:
+ return None
+ elif cls.mtp_only:
+ # --mtp: export the MTP head plus the tensors it shares with the target model
+ keep = name in (
+ "backbone.embeddings.weight",
+ "backbone.norm_f.weight",
+ "lm_head.weight",
+ )
+ if not keep:
+ return None
+ return super().filter_tensors((name, gen))
+
+ def prepare_metadata(self, vocab_only: bool):
+ from_dir = self.fname_out.is_dir()
+ super().prepare_metadata(vocab_only=vocab_only)
+
+ if not self.mtp_only or not from_dir:
+ return
+ output_type: str = self.ftype.name.partition("_")[2]
+ fname_default: str = gguf.naming_convention(
+ self.metadata.name, self.metadata.basename, self.metadata.finetune,
+ self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+ self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
def set_gguf_parameters(self):
super().set_gguf_parameters()
if (latent_size := self.hparams.get("moe_latent_size")) is not None:
self.gguf_writer.add_moe_latent_size(latent_size)
+ # MTP head: number of trailing NextN blocks
+ if self._mtp_bid is not None:
+ self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"])
+
def set_vocab(self):
# The NemotronH config uses pattern characters (e.g. '-') that may not
# be supported by the installed transformers version. AutoTokenizer
if not self.is_moe:
self.gguf_writer.add_add_bos_token(True)
+ _MTP_SPECIAL_RENAMES = {
+ "mtp.layers.0.enorm.weight": "model.layers.{bid}.enorm.weight",
+ "mtp.layers.0.hnorm.weight": "model.layers.{bid}.hnorm.weight",
+ "mtp.layers.0.eh_proj.weight": "model.layers.{bid}.eh_proj.weight",
+ "mtp.layers.1.norm.weight": "model.layers.{bid}.post_attention_layernorm.weight",
+ "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight",
+ }
+
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
- if self.is_moe and bid is not None:
- # Skip Multi-Token Prediction (MTP) tensors. These are used for
- # for speculative decoding but we don't include them in this model
- # conversion. See https://github.com/ggml-org/llama.cpp/pull/18886
- if name.startswith("mtp."):
- logger.info(f"gguf: Skipping MTP (Speculative) layer: {name}")
- return
+ # mtp.layers.0: NextN input fusion + attention
+ # mtp.layers.1: MoE + final head norm
+ if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")):
+ suffix = name.split(".", 3)[3]
+ bid = self._mtp_bid
+ renamed = self._MTP_SPECIAL_RENAMES.get(name)
+ name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}"
+ if self.is_moe and bid is not None:
if name.endswith("mixer.gate.e_score_correction.bias"):
yield from ModelBase.modify_tensors(self, data_torch, name, bid)
return
MODEL_TENSOR.FFN_DOWN_SHEXP,
MODEL_TENSOR.FFN_UP_SHEXP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
+ # NextN/MTP (draft head)
+ MODEL_TENSOR.ATTN_POST_NORM,
+ MODEL_TENSOR.NEXTN_EH_PROJ,
+ MODEL_TENSOR.NEXTN_ENORM,
+ MODEL_TENSOR.NEXTN_HNORM,
+ MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
],
MODEL_ARCH.EXAONE: [
MODEL_TENSOR.TOKEN_EMBD,
params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
(arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
+ const bool mtp_on_hybrid_nemotron =
+ params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE;
+
if (llm_arch_is_recurrent(arch)) {
res = new llama_memory_recurrent(
*this,
cparams.n_seq_max,
cparams.n_rs_seq,
nullptr);
- } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) {
+ } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen && !mtp_on_hybrid_nemotron) {
// The main difference between hybrid architectures is the
// layer filters, so pick the right one here
llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
};
}
- if (mtp_on_hybrid_qwen) {
+ if (mtp_on_hybrid_qwen || mtp_on_hybrid_nemotron) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
using graph = llama_model_nemotron_h::graph;
+ struct graph_mtp : public llm_graph_context {
+ graph_mtp(const llama_model & model, const llm_graph_params & params);
+ };
+
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
#include "models.h"
std::unique_ptr<llm_graph_context> llama_model_nemotron_h_moe::build_arch_graph(const llm_graph_params & params) const {
+ if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
+ return std::make_unique<graph_mtp>(*this, params);
+ }
return std::make_unique<graph>(*this, params);
}
+// MTP draft head for Nemotron-H MoE
+llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+ : llm_graph_context(params) {
+ GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block");
+
+ const int64_t n_embd_head = hparams.n_embd_head_v();
+ GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+ const int il = hparams.n_layer();
+ const auto & layer = model.layers[il];
+
+ GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm);
+ GGML_ASSERT(layer.ffn_gate_inp);
+
+ // token embedding weights
+ ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+ GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings");
+
+ auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->tokens);
+
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
+ ggml_set_input(inp->embd);
+
+ ggml_tensor * tok_embd;
+ if (ubatch.token) {
+ tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+ } else {
+ tok_embd = inp->embd;
+ }
+ cb(tok_embd, "mtp_tok_embd", il);
+
+ inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
+ ggml_set_input(inp->h);
+ ggml_set_name(inp->h, "mtp_h_input");
+
+ ggml_tensor * h_embd = inp->h;
+
+ res->add_input(std::move(inp));
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ // attention fills KV over all tokens, but the MoE is position-wise: gather output rows before
+ // it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state)
+ const bool emit_h_nextn = cparams.embeddings_nextn;
+ const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
+
+ auto * inp_attn = build_attn_inp_kv();
+
+ ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+ cb(h_norm, "mtp_hnorm", il);
+ ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+ cb(e_norm, "mtp_enorm", il);
+
+ ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
+ cb(concat, "mtp_concat", il);
+
+ ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+ cb(cur, "mtp_eh_proj", il);
+
+ // dense NoPE attention sub-layer (mtp.layers.0)
+ ggml_tensor * inpSA = cur;
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_norm", il);
+
+ {
+ auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
+ const float kq_scale = hparams.f_attention_scale == 0.0f
+ ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
+ cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ cb(cur, "mtp_attn_out", il);
+ }
+
+ cur = ggml_add(ctx0, cur, inpSA);
+ cb(cur, "mtp_attn_residual", il);
+
+ // gather the output rows here so the MoE FFN below only runs on the positions we keep
+ if (crop_before_ffn) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ }
+
+ // MoE FFN sub-layer (mtp.layers.1)
+ ggml_tensor * ffn_residual = cur;
+ cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "mtp_attn_post_norm", il);
+
+ {
+ ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur);
+ cb(router_logits, "mtp_ffn_moe_logits", il);
+
+ ggml_tensor * moe_out =
+ build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ nullptr, // no gate
+ layer.ffn_down_exps,
+ layer.ffn_exp_probs_b,
+ n_expert, n_expert_used,
+ LLM_FFN_RELU_SQR, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,
+ il,
+ router_logits, nullptr,
+ layer.ffn_up_exps_s,
+ nullptr, // no gate
+ layer.ffn_down_exps_s);
+ cb(moe_out, "mtp_ffn_moe_out", il);
+
+ ggml_tensor * ffn_shexp = build_ffn(cur,
+ layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
+ NULL, NULL, NULL,
+ layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
+ NULL,
+ LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "mtp_ffn_out", il);
+ }
+
+ cur = ggml_add(ctx0, cur, ffn_residual);
+ cb(cur, "mtp_post_ffn", il);
+
+ // final head norm: the MTP head has its own LayerNorm
+ GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm");
+ cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1);
+
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ if (!crop_before_ffn && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ }
+
+ // LM head
+ ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+ ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+ GGML_ASSERT(head_w != nullptr && "NEMOTRON_H_MOE MTP requires an output projection");
+ cur = build_lora_mm(head_w, cur, head_s);
+ cb(cur, "result_output", -1);
+
+ res->t_logits = cur;
+ ggml_build_forward_expand(gf, cur);
+}
ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group);
+ // NextN/MTP: optional draft head appended as extra trailing block(s)
+ ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
+ GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
+
// A layer is recurrent IFF the n_head_kv value is set to 0 and
- // the n_ff value is set to 0
- for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
- hparams.is_recr_impl[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0);
+ // the n_ff value is set to 0. Appended MTP blocks are dense (non-recurrent)
+ for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+ hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0;
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); // MTP head final_layernorm
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false);
ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
}
}
-void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
+void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
+ const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
+ const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+ const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
+
// mamba2 Mixer SSM params
// NOTE: int64_t for tensor dimensions
const int64_t d_conv = hparams.ssm_d_conv;
auto & layer = layers[i];
// all blocks use the attn norm
- layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, trunk_flags);
if (hparams.is_recr(i)) {
// ssm layers
- layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0);
+ layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, trunk_flags);
- layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0);
+ layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, trunk_flags);
layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, TENSOR_NOT_REQUIRED);
- layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, 0);
+ layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, trunk_flags);
// no "weight" suffix for these
- layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, 0);
- layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, 0);
+ layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, trunk_flags);
+ layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, trunk_flags);
- layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0);
+ layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, trunk_flags);
// out_proj
- layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, 0);
+ layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, trunk_flags);
} else if (hparams.n_ff(i) == 0) {
// attention layers (with optional bias)
const int64_t n_head_i = hparams.n_head(i);
const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
- create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, 0);
- layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, 0);
+ create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, trunk_flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, trunk_flags);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
} else {
if (n_expert != 0) {
const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
const int64_t n_ff_shexp = hparams.n_ff_shexp;
- layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, 0);
- layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, 0);
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, trunk_flags);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, trunk_flags);
// MoE branch
layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, TENSOR_NOT_REQUIRED);
- layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, 0);
- layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, trunk_flags);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, trunk_flags);
// Shared expert branch
- layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);
- layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, trunk_flags);
} else {
// mlp layers
- layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, 0);
- layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, 0);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, trunk_flags);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, trunk_flags);
layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {hparams.n_ff(i)}, TENSOR_NOT_REQUIRED);
}
}
}
+
+ // NextN/MTP draft head: each predict layer folds an attention sub-layer and a MoE
+ // sub-layer into a single trailing block
+ for (int i = n_layer; i < n_layer_all; ++i) {
+ auto & layer = layers[i];
+
+ const int64_t n_head_i = hparams.n_head(i);
+ const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
+ const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
+ const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
+ const int64_t n_ff_shexp = hparams.n_ff_shexp;
+
+ // NextN input-fusion tensors
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags);
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags);
+
+ // attention sub-layer
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags);
+ create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, mtp_flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, mtp_flags);
+ layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED);
+
+ // MoE sub-layer
+ layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, mtp_flags);
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, mtp_flags);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, mtp_flags);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, mtp_flags);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, mtp_flags);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, mtp_flags);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, mtp_flags);
+ }
}
std::unique_ptr<llm_graph_context> llama_model_nemotron_h::build_arch_graph(const llm_graph_params & params) const {
cur = build_ffn_layer(cur, model, il);
}
- if (il == n_layer - 1 && inp_out_ids) {
+ if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+ // seed for the MTP/NextN draft head
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ }
+
cb(cur, "result_norm", -1);
res->t_embd = cur;
overview / pkg / ggml / sources / llama.cpp / commitdiff