model: add NVIDIA Nemotron-3-Puzzle-75B-A9B (NemotronHPuzzle) support (#25444)

* hparams: add per-layer n_ff_exp/n_expert_used arrays with scalar-or-array loading

G1/G2 infrastructure for variable-per-layer expert FFN size and top-k routing
(required for Puzzle-75B which has 5 distinct n_ff_exp values and 7 top-k values
across its 40 MoE layers).

Design: rename scalar members to _impl suffix (following existing convention),
add LLAMA_MAX_LAYERS arrays, add n_ff_exp(il)/n_expert_used(il) accessors with
scalar fallback. No new GGUF keys: reuses existing expert_feed_forward_length and
expert_used_count keys via get_key_or_arr (scalar -> broadcast, array -> per-layer).

- llama-hparams.h: n_ff_exp -> n_ff_exp_impl, n_expert_used -> n_expert_used_impl;
  add n_ff_exp_arr / n_expert_used_arr arrays; add per-layer accessor declarations.
- llama-hparams.cpp: implement n_ff_exp(il) and n_expert_used(il); out-of-range
  il returns impl safely (shared code, no abort).
- llama-model.cpp: central n_expert_used load changed to get_key_or_arr; derive
  impl as max-of-array for validations and backward compat; zero both new arrays;
  HunyuanVL override also zeroes n_expert_used_arr.
- llama-graph.cpp: aggregation loop in build_moe_ffn uses hparams.n_expert_used(il)
  so per-layer top-k bounds the ggml_view loop correctly.
- All other files: mechanical rename hparams.n_{ff_exp,expert_used} -> *_impl.
  Scalar arches are unaffected (broadcast fills all array slots with the single value).

(cherry picked from commit 269a81e03d66e1c353e1a203a0c03a03eb2b1a4e)

* nemotron-h: use per-layer n_ff_exp(il) and n_expert_used(il) at MoE call-sites

Load n_ff_exp via get_key_or_arr into hparams.n_ff_exp_arr in load_arch_hparams;
derive impl as max for existing uniform GGUFs.

In load_arch_tensors, compute n_ff_exp_i = hparams.n_ff_exp(i) with fallback to
n_ff(i)/n_expert_used(i) for GGUFs that omit expert_feed_forward_length.

In build_ffn_layer, pass hparams.n_expert_used(il) to build_moe_ffn so per-layer
top-k is used for expert routing selection.

All other nemotron-h behaviour (mamba2, attention, shared-exp, latent projection,
routed_scaling_factor, expert_weights_norm, sigmoid gating) is unchanged.

(cherry picked from commit b1878a101793cd4e59868ac72635a86ea694987c)

* arch/*.cpp + gguf-py: mechanical rename n_ff_exp->n_ff_exp_impl, n_expert_used->n_expert_used_impl

All non-nemotron arch files continue using the scalar impl member directly.
Behaviour is identical: the impl value is the broadcast value from the GGUF scalar.

gguf_writer: add_expert_feed_forward_length and add_expert_used_count now accept
int | Sequence[int], mirroring add_feed_forward_length, so converters can write
per-layer arrays with the same existing GGUF keys.

(cherry picked from commit 8f009f54bea5ef9a6a354123bd25e9d5ea2d5e03)

* convert: support NemotronHPuzzleForCausalLM (per-block MoE config)

Parse block_configs/mtp_block_configs into per-layer arrays (scalar-or-array
keys), append the MTP [attention, moe] sub-blocks as blk.88/blk.89 with
nextn tensors, accept the backbone.* prefix, and register the arch.
Also fix a pre-existing undeclared _experts attribute on NemotronHModel.

(cherry picked from commit d1a592f278336e78457454eb6c96bca917135f10)

* nemotron-h: distinguish Nemotron 3 Puzzle (75B.A9B) from Super (120B.A12B)

Both have 88 layers; the per-layer expert_used_count array (heterogeneous
for Puzzle, broadcast-uniform for Super) is the discriminator.

(cherry picked from commit f824e09dc812169589cf5662c92d149a4c18c30a)

* convert: accept the official Puzzle BF16 checkpoint's tensor naming

The officially distributed BF16 checkpoint (NVIDIA-Nemotron-Labs-3-Puzzle-
75B-A9B-BF16) names the trunk model.* (model.layers.*, model.embeddings,
model.norm_f) where the original release used the NemotronH-style
backbone.*, and spells the router bias e_score_correction_bias instead of
e_score_correction.bias. Normalize both at the top of
NemotronHPuzzleModel.modify_tensors so either checkpoint converts; every
tensor name in the official index (42683 keys, MTP head included) resolves
through the tensor map after normalization.

(cherry picked from commit 189b67fc2c9d50970416c94b3317a6e7baa49b03)

* laguna: use n_ff_exp_impl for the uniform-MoE FFN size

Laguna landed after this branch was cut and reads hparams.n_ff_exp as a
scalar. This series turns it into a per-layer array with an n_ff_exp(il)
accessor, so the three scalar reads no longer compile. Laguna is a
uniform MoE, so point them at the scalar fallback n_ff_exp_impl, same as
deepseek2/qwen3moe/gemma4 in this series. No behaviour change.

(cherry picked from commit dbedc9e19c50dca0acdfb402362e2707bee424ae)

* arch: extend the n_ff_exp/n_expert_used rename to archs added upstream

kimi-k3, dflash, bailingmoe3, deepseek4, granite-swa and the nemotron-h MTP
block still referenced the scalar fields by their old names. n_ff_exp and
n_expert_used are accessors now, so those reads no longer compile; point the
non-per-layer archs at the _impl scalars and use the indexed form where the
call site is per-layer.

* convert: keep Puzzle opted out of the NemotronH MTP export path

#26725 added MTP export to NemotronHModel, keyed on num_nextn_predict_layers.
Puzzle's config carries that key, but NemotronHPuzzleModel bypasses
NemotronHModel.__init__ (its per-block config needs a different setup), so
_mtp_bid was never assigned and modify_tensors raised AttributeError on any
mtp.* tensor. Puzzle's head is also laid out by mtp_block_configs, not the
mtp.layers.* form the base maps.

Set _mtp_bid to None, drop mtp.* in filter_tensors, and declare
supports_mtp_export = False so --mtp / --no-mtp fail at the CLI.

* llama: replace n_ff_exp/n_expert_used scalars with per-layer accessors

Follow-up to review feedback: the previous revision kept the scalar
hparams fields alongside the new per-layer arrays, which duplicated
state that get_key_or_arr already handles by broadcasting a scalar
value over every layer.

Drop both scalars and expose n_ff_exp(il) / n_expert_used(il) built
exactly like the existing n_head_kv(il) and n_ff(il) accessors: they
index the array and GGML_ABORT out of range, with il defaulting to 0
so genuinely uniform call sites stay a plain n_ff_exp().

Arch loaders now read both keys through get_key_or_arr over
n_layer_all, and the n_expert_used validation checks the maximum
across layers instead of a single field.

* llama: restore per-key required flags on the expert hparam reads

The scalar-to-array conversion passed required=false at every call site,
which silently made mandatory keys optional. Each read now carries the
same required flag it had before the conversion.
This commit is contained in:
Yaniss Amazouz
2026-09-03 08:53:08 +02:00
committed by GitHub
parent 67a17c17ca
commit c61b98b875
56 changed files with 298 additions and 154 deletions
+30 -21
View File
@@ -634,7 +634,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
// the FFN is the same for Qwen 3 Next and Qwen 3.5:
if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
const int64_t n_ff_exp = hparams.n_ff_exp;
const int64_t n_ff_exp = hparams.n_ff_exp(il);
GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp);
return {{n_ff_exp, 2}};
}
@@ -657,7 +657,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
return {{tensor->ne[axis], 1}};
}
if (std::regex_match(tensor_name, pattern_ffn_gate_up_weight)) {
const int64_t n_ff_exp = hparams.n_ff_exp;
const int64_t n_ff_exp = hparams.n_ff_exp(il);
GGML_ASSERT(tensor->ne[axis] == 2*n_ff_exp);
return {{n_ff_exp, 2}};
}
@@ -943,6 +943,7 @@ const char * llm_type_name(llm_type type) {
case LLM_TYPE_31B_A3_5B: return "31B.A3.5B";
case LLM_TYPE_35B_A3B: return "35B.A3B";
case LLM_TYPE_48B_A3B: return "48B.A3B";
case LLM_TYPE_75B_A9B: return "75B.A9B";
case LLM_TYPE_80B_A3B: return "80B.A3B";
case LLM_TYPE_A3B: return "A3B";
case LLM_TYPE_100B_A6B: return "100B.A6B";
@@ -1226,14 +1227,15 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
GGML_ASSERT(hparams.n_layer_nextn <= hparams.n_layer_all);
ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert, false);
ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used, false);
std::fill(hparams.n_expert_used_arr.begin(), hparams.n_expert_used_arr.end(), 0);
ml.get_key_or_arr(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used_arr, hparams.n_layer_all, false);
ml.get_key(LLM_KV_EXPERT_GROUP_COUNT, hparams.n_expert_groups, false);
ml.get_key(LLM_KV_EXPERT_GROUP_USED_COUNT, hparams.n_group_used, false);
if (arch == LLM_ARCH_HUNYUAN_VL || arch == LLM_ARCH_HUNYUAN_DENSE) {
if (hparams.n_expert <= 1) {
hparams.n_expert = 0;
hparams.n_expert_used = 0;
hparams.n_expert = 0;
std::fill(hparams.n_expert_used_arr.begin(), hparams.n_expert_used_arr.end(), 0);
}
}
@@ -1251,10 +1253,16 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
GGML_ASSERT(hparams.convnext.n_layer <= hparams.n_layer_all);
}
// models may route a different number of experts per layer, so validate the maximum
uint32_t n_expert_used_max = 0;
for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
n_expert_used_max = std::max(n_expert_used_max, hparams.n_expert_used(il));
}
GGML_ASSERT(hparams.n_expert <= LLAMA_MAX_EXPERTS);
GGML_ASSERT(hparams.n_expert_used <= hparams.n_expert);
GGML_ASSERT(n_expert_used_max <= hparams.n_expert);
if (hparams.n_expert > 0) {
GGML_ASSERT(hparams.n_expert_used > 0);
GGML_ASSERT(n_expert_used_max > 0);
GGML_ASSERT(hparams.n_expert_groups < hparams.n_expert);
if (hparams.n_expert_groups > 1) {
GGML_ASSERT(hparams.n_expert % hparams.n_expert_groups == 0);
@@ -1262,13 +1270,14 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
GGML_ASSERT(hparams.n_group_used < hparams.n_expert_groups);
}
} else {
GGML_ASSERT(hparams.n_expert_used == 0);
GGML_ASSERT(n_expert_used_max == 0);
GGML_ASSERT(hparams.n_expert_groups == 0);
}
std::fill(hparams.n_head_arr.begin(), hparams.n_head_arr.end(), 0);
std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0);
std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0);
std::fill(hparams.n_head_arr.begin(), hparams.n_head_arr.end(), 0);
std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0);
std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0);
std::fill(hparams.n_ff_exp_arr.begin(), hparams.n_ff_exp_arr.end(), 0);
std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0);
std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), 1);
@@ -1501,7 +1510,7 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
const auto tn = LLM_TN(arch);
const int64_t n_expert = hparams.n_expert;
const int64_t n_expert_used = hparams.n_expert_used;
const int64_t n_expert_used = hparams.n_expert_used();
if (n_expert > 0 && n_expert_used == 0) {
throw std::runtime_error("model has expert layers but no expert layers are used");
@@ -1957,7 +1966,7 @@ void llama_model::print_info() const {
LLAMA_LOG_INFO("%s: f_attn_value_scale = %.4f\n", __func__, hparams.f_attn_value_scale);
LLAMA_LOG_INFO("%s: n_ff = %s\n", __func__, print_f([&](uint32_t il) { return hparams.n_ff(il); }, hparams.n_layer_all).c_str());
LLAMA_LOG_INFO("%s: n_expert = %u\n", __func__, hparams.n_expert);
LLAMA_LOG_INFO("%s: n_expert_used = %u\n", __func__, hparams.n_expert_used);
LLAMA_LOG_INFO("%s: n_expert_used = %u\n", __func__, hparams.n_expert_used());
LLAMA_LOG_INFO("%s: n_expert_groups = %d\n", __func__, hparams.n_expert_groups);
LLAMA_LOG_INFO("%s: n_group_used = %d\n", __func__, hparams.n_group_used);
LLAMA_LOG_INFO("%s: causal attn = %d\n", __func__, hparams.causal_attn);
@@ -2032,7 +2041,7 @@ void llama_model::print_info() const {
if (arch == LLM_ARCH_DEEPSEEK) {
LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
LLAMA_LOG_INFO("%s: n_expert_shared = %d\n", __func__, hparams.n_expert_shared);
LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale);
}
@@ -2045,7 +2054,7 @@ void llama_model::print_info() const {
LLAMA_LOG_INFO("%s: n_lora_kv = %d\n", __func__, hparams.n_lora_kv);
LLAMA_LOG_INFO("%s: n_embd_head_k_mla = %d\n", __func__, hparams.n_embd_head_k_mla());
LLAMA_LOG_INFO("%s: n_embd_head_v_mla = %d\n", __func__, hparams.n_embd_head_v_mla());
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
LLAMA_LOG_INFO("%s: n_expert_shared = %d\n", __func__, hparams.n_expert_shared);
LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale);
LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm);
@@ -2053,7 +2062,7 @@ void llama_model::print_info() const {
}
if (arch == LLM_ARCH_QWEN2MOE) {
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp);
}
@@ -2063,7 +2072,7 @@ void llama_model::print_info() const {
arch == LLM_ARCH_OPENAI_MOE ||
arch == LLM_ARCH_QWEN3VLMOE ||
arch == LLM_ARCH_RND1) {
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
}
if (arch == LLM_ARCH_MINICPM ||
@@ -2080,7 +2089,7 @@ void llama_model::print_info() const {
if (arch == LLM_ARCH_BAILINGMOE) {
LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
LLAMA_LOG_INFO("%s: n_expert_shared = %d\n", __func__, hparams.n_expert_shared);
LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale);
LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm);
@@ -2088,7 +2097,7 @@ void llama_model::print_info() const {
if (arch == LLM_ARCH_BAILINGMOE2 || arch == LLM_ARCH_BAILINGMOE3) {
LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp);
LLAMA_LOG_INFO("%s: n_expert_shared = %d\n", __func__, hparams.n_expert_shared);
LLAMA_LOG_INFO("%s: expert_weights_scale = %.1f\n", __func__, hparams.expert_weights_scale);
@@ -2098,12 +2107,12 @@ void llama_model::print_info() const {
}
if (arch == LLM_ARCH_SMALLTHINKER || arch == LLM_ARCH_LFM2MOE) {
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
LLAMA_LOG_INFO("%s: expert_gating_func = %s\n", __func__, llama_expert_gating_func_name((llama_expert_gating_func_type) hparams.expert_gating_func));
}
if (arch == LLM_ARCH_GROVEMOE) {
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp);
LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp());
LLAMA_LOG_INFO("%s: n_ff_chexp = %d\n", __func__, hparams.n_ff_chexp);
LLAMA_LOG_INFO("%s: n_group_experts = %d\n", __func__, hparams.n_group_experts);
LLAMA_LOG_INFO("%s: expert_group_scale = %.2f\n", __func__, hparams.expert_group_scale);