diff --git a/conversion/zamba2.py b/conversion/zamba2.py index 0343b6946444..94715db49024 100644 --- a/conversion/zamba2.py +++ b/conversion/zamba2.py @@ -64,6 +64,15 @@ def __init__(self, *args, **kwargs): self._adapter_cache: dict[str, Tensor] = {} self._use_attn_adapters = self.hparams.get("use_shared_attention_adapter", False) self._use_mlp_adapters = self.hparams.get("use_shared_mlp_adapter", True) + # The MLP adapter is stored instead of merged when the checkpoint actually carries one + self._store_mlp_adapter = self._use_mlp_adapters and any( + "gate_up_proj_adapter_list" in t for t in self.model_tensors) + self._adapter_rank = int(self.hparams.get("adapter_rank", 0) or 0) + if self._store_mlp_adapter and self._adapter_rank <= 0: + raise ValueError( + f"{self.dir_model}: the MLP adapters are stored separately, so the checkpoint must declare " + "adapter_rank; without it the low-rank tensors cannot be sized" + ) adapter_pattern = re.compile( r"model\.layers\.(\d+)\.shared_transformer\." @@ -85,11 +94,16 @@ def __init__(self, *args, **kwargs): # Tails of a shared block that this checkpoint stores no per-layer adapter for are # byte-identical in every layer using the block, so they are written once, at the # block's source layer, and the loader aliases the other layers to it - # (llama_model_zamba2). A tail with an adapter is per layer after merging and stays - # per layer -- that one cannot be shared without changing the graph. + # (llama_model_zamba2). Attention adapters, when a checkpoint has them, are still merged + # into the per-layer copies -- that is unchanged. The MLP adapter is NOT merged: the base + # gate_up is shared like the rest of the block and each layer keeps its low-rank pair, + # applied in the graph as base + B*(A*x) (engine #1036 cause 2). Baking it per layer cost + # 1.13 B params on the 7B, and its 11 duplicated copies of a 102 MB matrix bought nothing. adapted = {tail for tail, atype in self._ADAPTER_MAP.items() - if any(atype in t for t in self.model_tensors)} + if atype.startswith("linear_") and any(atype in t for t in self.model_tensors)} self._shared_tails = self._SHARED_TAILS - adapted + if self._store_mlp_adapter: + self._shared_tails = self._shared_tails | {"feed_forward.gate_up_proj.weight"} def _get_adapter_contribution(self, src_layer: int, adapter_type: str, visit_idx: int) -> Tensor | None: """Compute up @ down for a given adapter, returning the additive correction.""" @@ -160,6 +174,11 @@ def set_gguf_parameters(self): self.gguf_writer.add_ssm_time_step_rank(n_heads) self.gguf_writer.add_ssm_group_count(self.n_group) + if self._store_mlp_adapter: + # the one key that tells the loader which of the two forms this file is: > 0 means ffn_up is the + # shared base and every hybrid layer carries a rank-sized A/B pair + self.gguf_writer.add_zamba2_ffn_gate_up_adapter_rank(self._adapter_rank) + self.gguf_writer.add_file_type(self.ftype) # Block tails that carry no per-layer adapter in a Zamba2 checkpoint, i.e. the ones the @@ -183,7 +202,27 @@ def set_gguf_parameters(self): } def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - # Skip adapter weights — they are pre-merged into shared weights below + # The per-layer MLP adapter is written as its two low-rank factors, at the hybrid layer that owns it, + # instead of being baked into that layer's copy of the shared gate_up matrix. The index in the + # checkpoint's adapter_list is the hybrid visit index, which is exactly how the layers map to blocks. + if self._store_mlp_adapter and "feed_forward.gate_up_proj_adapter_list" in name: + m = re.match( + r"model\.layers\.(\d+)\.shared_transformer\.feed_forward\.gate_up_proj_adapter_list\.(\d+)\.([01])\.weight$", + name) + if m is None or bid is None or bid not in self._shared_block_src: + # entries for other blocks are nn.Identity and carry no weights; duplicates from ties are + # read once, at the block's source layer, like the shared tails above + return + visit_idx = int(m.group(2)) + seq = int(m.group(3)) # 0 = down (hidden -> rank), 1 = up (rank -> 2*ff) + if visit_idx >= len(self._hybrid_layers): + raise ValueError(f"{name}: adapter index {visit_idx} has no hybrid layer") + target_bid = self._hybrid_layers[visit_idx] + tensor = gguf.MODEL_TENSOR.FFN_GATE_UP_ADAPTER_A if seq == 0 else gguf.MODEL_TENSOR.FFN_GATE_UP_ADAPTER_B + yield (self.format_tensor_name(tensor, target_bid), data_torch) + return + + # Skip adapter weights that are pre-merged into shared weights below if "adapter" in name.lower(): return @@ -210,14 +249,13 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_POST_NORM, target_bid), data_torch) return - # Check if this weight has a corresponding adapter type + # Check if this weight has a corresponding adapter type. Only the attention adapters are merged + # here; the MLP one is stored, so its base is shared and untouched. adapter_type = self._ADAPTER_MAP.get(tail) use_adapter = False if adapter_type: if adapter_type.startswith("linear_") and self._use_attn_adapters: use_adapter = True - elif adapter_type == "gate_up_proj_adapter_list" and self._use_mlp_adapters: - use_adapter = True # Delegate to the parent's tensor-name mapping, once per target layer. for visit_local, target_bid in enumerate(target_layers): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 8329fb204213..469a4de43bda 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -273,6 +273,10 @@ class SSM: class KDA: HEAD_DIM = "{arch}.kda.head_dim" + class Zamba2: + # names the arch explicitly: 0 or absent = the per-layer MLP adapter is baked into ffn_up + FFN_GATE_UP_ADAPTER_RANK = "zamba2.ffn_gate_up_adapter_rank" + class WKV: HEAD_SIZE = "{arch}.wkv.head_size" @@ -742,6 +746,8 @@ class MODEL_TENSOR(IntEnum): ZAYA_VLORA_DOWN_EXPS_A = auto() # Zaya1-VL ZAYA_VLORA_DOWN_EXPS_B = auto() # Zaya1-VL SSM_MIX = auto() # zamba2 + FFN_GATE_UP_ADAPTER_A = auto() # zamba2: per-layer low-rank MLP adapter, hidden -> rank + FFN_GATE_UP_ADAPTER_B = auto() # zamba2: per-layer low-rank MLP adapter, rank -> 2*ff TIME_MIX_W0 = auto() TIME_MIX_W1 = auto() TIME_MIX_W2 = auto() @@ -1396,6 +1402,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SSM_NORM: "blk.{bid}.ssm_norm", MODEL_TENSOR.SSM_OUT: "blk.{bid}.ssm_out", MODEL_TENSOR.SSM_MIX: "blk.{bid}.ssm_mix", + MODEL_TENSOR.FFN_GATE_UP_ADAPTER_A: "blk.{bid}.ffn_gate_up_adapter_a", + MODEL_TENSOR.FFN_GATE_UP_ADAPTER_B: "blk.{bid}.ffn_gate_up_adapter_b", MODEL_TENSOR.SSM_ALPHA: "blk.{bid}.ssm_alpha", # qwen3.5 MODEL_TENSOR.SSM_BETA_ALPHA: "blk.{bid}.ssm_ba", MODEL_TENSOR.SSM_CONV1D_Q: "blk.{bid}.ssm_conv1d_q", # Kimi Linear @@ -3316,6 +3324,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_NORM, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_UP_ADAPTER_A, + MODEL_TENSOR.FFN_GATE_UP_ADAPTER_B, ], MODEL_ARCH.ZAMBA: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 4e91ee3a8212..8ea70edf474c 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1143,6 +1143,9 @@ def add_ssm_time_step_rank(self, value: int) -> None: def add_ssm_group_count(self, value: int) -> None: self.add_uint32(Keys.SSM.GROUP_COUNT.format(arch=self.arch), value) + def add_zamba2_ffn_gate_up_adapter_rank(self, value: int) -> None: + self.add_uint32(Keys.Zamba2.FFN_GATE_UP_ADAPTER_RANK, value) + def add_ssm_dt_b_c_rms(self, value: bool) -> None: self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index a94b87652bb2..a580edca2bfe 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -412,6 +412,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_XIELU_ALPHA_P, "xielu.alpha_p" }, { LLM_KV_XIELU_BETA, "xielu.beta" }, { LLM_KV_XIELU_EPS, "xielu.eps" }, + { LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK, "zamba2.ffn_gate_up_adapter_rank" }, // deprecated { LLM_KV_TOKENIZER_PREFIX_ID, "tokenizer.ggml.prefix_token_id" }, @@ -481,6 +482,8 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, { LLM_TENSOR_SSM_OUT, "blk.%d.ssm_out" }, { LLM_TENSOR_SSM_MIX, "blk.%d.ssm_mix" }, + { LLM_TENSOR_FFN_GATE_UP_ADAPTER_A, "blk.%d.ffn_gate_up_adapter_a" }, + { LLM_TENSOR_FFN_GATE_UP_ADAPTER_B, "blk.%d.ffn_gate_up_adapter_b" }, { LLM_TENSOR_ROPE_FACTORS_LONG, "rope_factors_long" }, { LLM_TENSOR_ROPE_FACTORS_SHORT, "rope_factors_short" }, { LLM_TENSOR_SSM_X, "blk.%d.ssm_x" }, @@ -828,6 +831,8 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_SSM_DT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_OUT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_MIX, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_GATE_UP_ADAPTER_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_GATE_UP_ADAPTER_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_ALPHA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_BETA_ALPHA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_TIME_MIX_W1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 87317f0728db..c989e80888a5 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -412,6 +412,10 @@ enum llm_kv { LLM_KV_XIELU_BETA, LLM_KV_XIELU_EPS, + // zamba2: 0 or absent = the per-layer MLP adapters are baked into ffn_up (files written before this key + // existed); > 0 = ffn_up is the shared base and each hybrid layer carries a low-rank adapter pair of this rank + LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK, + // deprecated: LLM_KV_TOKENIZER_PREFIX_ID, LLM_KV_TOKENIZER_SUFFIX_ID, @@ -513,6 +517,8 @@ enum llm_tensor { LLM_TENSOR_SSM_NORM, LLM_TENSOR_SSM_OUT, LLM_TENSOR_SSM_MIX, // zamba2: linear mixing between transformer and mamba + LLM_TENSOR_FFN_GATE_UP_ADAPTER_A, // zamba2: per-layer low-rank MLP adapter, hidden -> adapter_rank + LLM_TENSOR_FFN_GATE_UP_ADAPTER_B, // zamba2: per-layer low-rank MLP adapter, adapter_rank -> 2*n_ff LLM_TENSOR_SSM_BETA_ALPHA, // qwen3next LLM_TENSOR_SSM_ALPHA, // qwen3.5 // Kimi Linear KDA (using SSM_ prefix for consistency) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 338a75045481..5a53798efc68 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1602,7 +1602,8 @@ ggml_tensor * llm_graph_context::build_ffn( ggml_tensor * act_scales, llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, - int il) const { + int il, + ggml_tensor * up_delta) const { // NVFP4 support is currently restricted to // 1) LORA absence (*_s would be applied after LORA residual, which is incorrect) // 2) bias absense (*_s would be applied after bias addition, which is incorrect) @@ -1634,6 +1635,13 @@ ggml_tensor * llm_graph_context::build_ffn( cb(tmp, "ffn_up_b", il); } + // a low-rank adapter added to the fused gate/up output (zamba2: base + B*(A*x)); an addition to the matmul + // result, so it is placed before the split and the activation, exactly where the checkpoint applies it + if (up_delta) { + tmp = ggml_add(ctx0, tmp, up_delta); + cb(tmp, "ffn_up_delta", il); + } + if (up_s) { tmp = ggml_mul(ctx0, tmp, up_s); cb(tmp, "ffn_up_s", il); diff --git a/src/llama-graph.h b/src/llama-graph.h index 9702974e6e43..e69b4603cbce 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1005,7 +1005,8 @@ struct llm_graph_context { ggml_tensor * act_scales, llm_ffn_op_type type_op, llm_ffn_gate_type type_gate, - int il) const; + int il, + ggml_tensor * up_delta = nullptr) const; // build MoE FFN without bias tensors ggml_tensor * build_moe_ffn( diff --git a/src/llama-hparams.h b/src/llama-hparams.h index cf22d37ba554..381e0fd64def 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -167,6 +167,10 @@ struct llama_hparams { uint32_t ssm_dt_rank = 0; uint32_t ssm_n_group = 0; + // zamba2: 0 or absent = the per-layer MLP adapter is baked into ffn_up; > 0 = ffn_up is the shared base and + // each hybrid layer carries a low-rank adapter pair of this rank (engine #1036 cause 2) + uint32_t zamba2_ffn_gate_up_adapter_rank = 0; + // for Kimi Linear KDA uint32_t n_embd_head_kda = 0; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index ef29e44b99e0..413c8dc0af17 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -401,6 +401,11 @@ void llama_model_saver::add_kv_from_model() { } else { add_kv(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); } + if (hparams.zamba2_ffn_gate_up_adapter_rank > 0) { + // only written when the tensors it describes are present, so the saver cannot claim the adapter form + // for a file that has the adapter baked into ffn_up + add_kv(LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK, hparams.zamba2_ffn_gate_up_adapter_rank); + } add_kv(LLM_KV_SSM_DT_B_C_RMS, hparams.ssm_dt_b_c_rms); add_kv(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); diff --git a/src/llama-model.h b/src/llama-model.h index a1de0c821f86..1df27f2c49cd 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -303,6 +303,10 @@ struct llama_layer { struct ggml_tensor * ffn_gate_enc = nullptr; struct ggml_tensor * ffn_down_enc = nullptr; struct ggml_tensor * ffn_up_enc = nullptr; + // zamba2: the per-layer MLP adapter (hidden -> rank -> 2*n_ff), applied as base + B*(A*x); null when the + // adapter was baked into ffn_up at conversion time + struct ggml_tensor * ffn_gate_up_adapter_a = nullptr; + struct ggml_tensor * ffn_gate_up_adapter_b = nullptr; // ff MoE struct ggml_tensor * ffn_gate_inp = nullptr; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index d992b677e974..fe2184551375 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -680,6 +680,12 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type return new_type; } +// zamba2's per-layer MLP adapter pair (engine #1036 cause 2). Matched by name because the pair is an arch-level +// tensor a quantize rule needs to see before the type is chosen, and it has no category of its own. +static bool is_zamba2_mlp_adapter(const char * name) { + return std::strstr(name, "ffn_gate_up_adapter_") != nullptr; +} + // outer wrapper: determine the ggml_type that this tensor should be quantized to static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_model_quantize_params * params, const ggml_tensor * tensor, ggml_type default_type, const tensor_metadata & tm) { if (!tensor_allows_quantization(params, qs.model.arch, tensor)) { @@ -728,8 +734,19 @@ static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_mod } // if not manual - use the standard logic for choosing the quantization type based on the selected mixture - if (!manual && !params->pure) { - new_type = llama_tensor_get_type_impl(qs, new_type, tensor, params->ftype, tm.category); + if (!manual) { + if (!params->pure && is_zamba2_mlp_adapter(tensor->name)) { + // zamba2's per-layer MLP adapter is a rank-128 pair applied at run time (base + B*(A*x)). Quantizing + // it saves little and costs the whole pair its place on the device: the B contraction has + // input_size = rank = 128, which is below the block formats' `input_size >= 256` minimum (their + // own block layout) while the float path accepts `>= 4 && % 4 == 0`. Measured on a 2.7B Q5_K_M, + // quantizing the pair puts 9 of 3051 nodes on the CPU (one per hybrid layer) and takes prefill from + // 589.5 to 463.7 tok/s; keeping it at F16 costs 31.8 MiB of 1804.9. --pure still forces the + // requested type on every tensor, and an explicit --tensor-type wins over this default. + new_type = GGML_TYPE_F16; + } else if (!params->pure) { + new_type = llama_tensor_get_type_impl(qs, new_type, tensor, params->ftype, tm.category); + } } // incompatible tensor shapes are handled here - fallback to a compatible type diff --git a/src/models/zamba2.cpp b/src/models/zamba2.cpp index 3959ea6b6a5d..cb7e15ad0bab 100644 --- a/src/models/zamba2.cpp +++ b/src/models/zamba2.cpp @@ -22,6 +22,8 @@ void llama_model_zamba2::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); 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); + // absent in files written before the MLP adapter was stored separately (engine #1036 cause 2) + ml.get_key(LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK, hparams.zamba2_ffn_gate_up_adapter_rank, false); // All layers have Mamba-2 (recurrent state needed for all) std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), true); @@ -38,7 +40,7 @@ void llama_model_zamba2::load_arch_hparams(llama_model_loader & ml) { hparams.f_attention_scale = 1.0f / sqrtf(float(hparams.n_embd_head_v_full) / 2.0f); } -void llama_model_zamba2::load_arch_tensors(llama_model_loader &) { +void llama_model_zamba2::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t d_conv = hparams.ssm_d_conv; @@ -51,6 +53,15 @@ void llama_model_zamba2::load_arch_tensors(llama_model_loader &) { const auto TENSOR_NOT_REQUIRED = llama_model_loader::TENSOR_NOT_REQUIRED; + // Two forms of the file, told apart by one key. Before this key existed the per-layer MLP adapter was baked + // into ffn_up at conversion time, so every hybrid layer carries its own full [n_embd, 2*n_ff] copy of a + // matrix that is otherwise shared per block (engine #1036 cause 2: 1.13 B params on the 7B). With the key + // set, ffn_up is the shared base -- stored once per block, like the rest of the shared transformer block -- + // and each hybrid layer carries the low-rank pair that makes its FFN different, applied in the graph as + // base + B*(A*x). Old files keep loading: rank 0 means ffn_up is per layer and no adapter is read. + uint32_t gate_up_adapter_rank = hparams.zamba2_ffn_gate_up_adapter_rank; + const bool ffn_up_adapter = gate_up_adapter_rank > 0; + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); @@ -87,8 +98,8 @@ void llama_model_zamba2::load_arch_tensors(llama_model_loader &) { // Everything belonging to the shared block is optional per layer: a layer that carries // no copy is aliased to its block's source after this loop, and a file with no copy at - // all is refused there rather than once per tensor. ffn_up is not shared -- the - // per-layer MLP adapter is baked into it at conversion time. + // all is refused there rather than once per tensor. ffn_up is part of the shared block + // only in the adapter form, where the per-layer difference is the low-rank pair below. const int shared = TENSOR_NOT_REQUIRED; // Pre-attention norm (4096-dim, operates on concat) @@ -100,9 +111,18 @@ void llama_model_zamba2::load_arch_tensors(llama_model_loader &) { // FFN (fused gate+up, split internally by GEGLU) layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, shared); - layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, 2*n_ff}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, 2*n_ff}, + ffn_up_adapter ? shared : 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, shared); + if (ffn_up_adapter) { + // GGUF stores matmul weights as [in, out]: A is hidden -> rank, B is rank -> 2*n_ff + layer.ffn_gate_up_adapter_a = create_tensor(tn(LLM_TENSOR_FFN_GATE_UP_ADAPTER_A, "weight", i), + {n_embd, gate_up_adapter_rank}, 0); + layer.ffn_gate_up_adapter_b = create_tensor(tn(LLM_TENSOR_FFN_GATE_UP_ADAPTER_B, "weight", i), + {gate_up_adapter_rank, 2*n_ff}, 0); + } + // Linear mixing (transformer output -> mamba input adjustment) layer.ssm_mix = create_tensor(tn(LLM_TENSOR_SSM_MIX, "weight", i), {n_embd, n_embd}, 0); @@ -136,6 +156,11 @@ void llama_model_zamba2::load_arch_tensors(llama_model_loader &) { layer.wo = src.wo; layer.ffn_norm = src.ffn_norm; layer.ffn_down = src.ffn_down; + // only when the file does not carry one: the adapter form stores ffn_up once per block, while a + // loader that fabricates tensors (no weights on disk, e.g. the arch test) gives every layer its own + if (layer.ffn_up == nullptr) { + layer.ffn_up = src.ffn_up; + } } } @@ -216,11 +241,22 @@ llama_model_zamba2::graph::graph(const llama_model & model, const llm_graph_para cur = build_norm(cur, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); + // The per-layer MLP adapter, applied to the same normed input as the fused gate+up matmul and added to + // its output (checkpoint: gate_up_proj(x) + adapter(x), no LoRA alpha/rank scaling). Absent in files + // written before the adapter form existed, where the adapter is baked into ffn_up instead. + ggml_tensor * gate_up_delta = NULL; + if (model.layers[il].ffn_gate_up_adapter_a != NULL) { + ggml_tensor * adapter = build_lora_mm(model.layers[il].ffn_gate_up_adapter_a, cur); + cb(adapter, "ffn_gate_up_adapter_a", il); + gate_up_delta = build_lora_mm(model.layers[il].ffn_gate_up_adapter_b, adapter); + cb(gate_up_delta, "ffn_gate_up_adapter_b", il); + } + cur = build_ffn(cur, model.layers[il].ffn_up, NULL, NULL, NULL, NULL, NULL, model.layers[il].ffn_down, NULL, NULL, - NULL, LLM_FFN_GEGLU, LLM_FFN_SEQ, il); + NULL, LLM_FFN_GEGLU, LLM_FFN_SEQ, il, gate_up_delta); cb(cur, "ffn_out", il); // Step 3: Linear mixing (transformer output projection) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index dfacc8bbefbb..9181d12fc422 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -183,6 +183,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { // the shared attention block runs on concat(hidden, embeddings): 2 * n_embd wide ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, 2 * n_embd_head); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, 2 * n_embd_head); + if (arch == LLM_ARCH_ZAMBA2) { + // exercise the low-rank MLP adapter form (engine #1036 cause 2) rather than the merged one: + // it is the newer path, it adds tensors the saver has to round-trip, and the merged form is + // what every published Zamba2 GGUF already exercises + ms.add_kv(LLM_KV_ZAMBA2_FFN_GATE_UP_ADAPTER_RANK, uint32_t(8)); + } } ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV, 1.0f); ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, 1e-5f);