diff --git a/conversion/spark2_5.py b/conversion/spark2_5.py index 44a0bd262e6a..a85442921ce8 100644 --- a/conversion/spark2_5.py +++ b/conversion/spark2_5.py @@ -1,12 +1,13 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Callable, Iterable +from copy import deepcopy from typing import TYPE_CHECKING if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, gguf +from .base import ModelBase, TextModel, gguf, torch @ModelBase.register("Spark2_5ForCausalLM") @@ -14,6 +15,12 @@ class Spark2_5Model(TextModel): model_arch = gguf.MODEL_ARCH.SPARK2_5 + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + if item[0].startswith("draft_model."): + return None + return super().filter_tensors(item) + def set_gguf_parameters(self) -> None: super().set_gguf_parameters() @@ -63,3 +70,62 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter ) yield from super().modify_tensors(data_torch, name, bid) + + +class Spark2_5DSparkModel(Spark2_5Model): + model_arch = gguf.MODEL_ARCH.DFLASH + + def __init__(self, dir_model, *args, **kwargs): + hparams = deepcopy(kwargs.pop("hparams", None) or ModelBase.load_hparams(dir_model, False)) + n_layers = hparams["num_draft_layers"] + hparams["num_hidden_layers"] = n_layers + hparams["layer_types"] = hparams.get("draft_layer_types", hparams["layer_types"][:n_layers]) + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + if self.hparams.get("markov_head_type", "vanilla") != "vanilla": + raise ValueError("Spark DSpark requires a vanilla Markov head") + if self.hparams.get("draft_vocab_size", self.hparams["vocab_size"]) != self.hparams["vocab_size"]: + raise ValueError("Spark DSpark currently requires the full target vocabulary") + if not self.hparams.get("confidence_head_with_markov", True): + raise ValueError("Spark DSpark requires confidence_head_with_markov") + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if not name.startswith("draft_model."): + return None + return TextModel.filter_tensors(("model." + name.removeprefix("draft_model."), gen)) + + def set_vocab(self) -> None: + super().set_vocab() + self.gguf_writer.add_mask_token_id(self.hparams["mask_token_id"]) + + def set_gguf_parameters(self) -> None: + super().set_gguf_parameters() + self.gguf_writer.add_block_size(self.hparams["dspark_block_size"]) + # Target layer outputs are extracted as the next layer's inputs. + self.gguf_writer.add_target_layers([i + 1 for i in self.hparams["target_layer_ids"]]) + self.gguf_writer.add_sample_from_anchor(self.hparams.get("sample_from_anchor", True)) + self.gguf_writer.add_causal_attention(False) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("model.fc_norm."): + if name == "model.fc_norm.0.weight": + norms = [self.model_tensors[f"model.fc_norm.{i}.weight"]().float() + for i in range(len(self.hparams["target_layer_ids"]))] + yield self.format_tensor_name(gguf.MODEL_TENSOR.FC_NORM), torch.stack(norms) + return + + if name.endswith(".self_attn.q_k_v_proj.weight"): + head_dim = self.hparams["head_dim"] + n_q = head_dim * self.hparams["num_attention_heads"] + n_kv = head_dim * self.hparams["num_key_value_heads"] + for tensor, value in zip( + (gguf.MODEL_TENSOR.ATTN_Q, gguf.MODEL_TENSOR.ATTN_K, gguf.MODEL_TENSOR.ATTN_V), + data_torch.split((n_q, n_kv, n_kv), dim=0), + strict=True, + ): + yield self.format_tensor_name(tensor, bid), value + return + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 5e263fd16fa7..acd36b1bb9c8 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -127,7 +127,7 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--dspark", action="store_true", - help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.", + help="Export only the DeepSeek-V4 or Spark2.5 DSpark draft tensors as a separate GGUF.", ) parser.add_argument( "--mistral-format", action="store_true", @@ -267,15 +267,21 @@ def main() -> None: sys.exit(1) if args.dspark: - if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"): - logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM") + if is_mistral_format: + logger.error("--dspark is not supported for Mistral format models") sys.exit(1) if model_architecture == "DeepseekV41ForCausalLM": from conversion.deepseek41 import DeepseekV41DSparkModel model_class = DeepseekV41DSparkModel - else: + elif model_architecture == "DeepseekV4ForCausalLM": from conversion.deepseek import DeepseekV4DSparkModel model_class = DeepseekV4DSparkModel + elif model_architecture == "Spark2_5ForCausalLM": + from conversion.spark2_5 import Spark2_5DSparkModel + model_class = Spark2_5DSparkModel + else: + logger.error("--dspark is only supported for DeepseekV4ForCausalLM, DeepseekV41ForCausalLM, and Spark2_5ForCausalLM") + sys.exit(1) if args.mtp or args.no_mtp: if not model_class.supports_mtp_export: diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index babc8bba4c6d..3c9822a1e00f 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1328,6 +1328,7 @@ class MODEL_TENSOR(IntEnum): NEXTN_HC_HEAD_UP = auto() # eagle3 FC = auto() # feature fusion layer + FC_NORM = auto() D2T = auto() # draft to target vocabulary mapping # dspark DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed @@ -2184,6 +2185,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: "blk.{bid}.nextn.hc_head_down", MODEL_TENSOR.NEXTN_HC_HEAD_UP: "blk.{bid}.nextn.hc_head_up", MODEL_TENSOR.FC: "fc", + MODEL_TENSOR.FC_NORM: "fc_norm", MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", @@ -5769,6 +5771,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K_NORM, MODEL_TENSOR.ATTN_SINKS, @@ -5801,6 +5804,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FC, + MODEL_TENSOR.FC_NORM, MODEL_TENSOR.ENC_OUTPUT_NORM, MODEL_TENSOR.D2T, # optional DSpark heads diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index ea13e9c3fcea..c56d24cd4572 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -792,6 +792,7 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_MASKED_EMBD_CENTROIDS, "masked_embd_centroids" }, { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, { LLM_TENSOR_FC, "fc" }, + { LLM_TENSOR_FC_NORM, "fc_norm" }, { LLM_TENSOR_D2T, "d2t" }, { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, @@ -1135,6 +1136,7 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_MASKED_EMBD_ORDERING, {LLM_TENSOR_LAYER_INPUT, GGML_OP_NONE}}, // eagle3 {LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FC_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, {LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_SET_ROWS}}, // dspark {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 7996bdc309b7..d15e54ef3bd0 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -798,6 +798,7 @@ enum llm_tensor { LLM_TENSOR_MASKED_EMBD_ORDERING, LLM_TENSOR_HRM_Z_L_INIT, LLM_TENSOR_FC, + LLM_TENSOR_FC_NORM, LLM_TENSOR_D2T, LLM_TENSOR_DSPARK_MARKOV_W1, LLM_TENSOR_DSPARK_MARKOV_W2, diff --git a/src/llama-model.h b/src/llama-model.h index 01fbda31928e..a01082736d01 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -790,6 +790,7 @@ struct llama_model { struct ggml_tensor * fc = nullptr; struct ggml_tensor * fc_s = nullptr; struct ggml_tensor * fc_in_s = nullptr; + struct ggml_tensor * fc_norm = nullptr; struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping // dspark diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 8b6d6abdbab5..1dad00788d4d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -113,12 +113,12 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { } // optional interleaved sliding-window attention with per-layer pattern array. - // DFlash has a single rope, so the SWA rope == main rope. if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) { hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); } type = LLM_TYPE_UNKNOWN; @@ -245,6 +245,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED); fc_in_s = create_tensor(tn(LLM_TENSOR_FC, "input_scale"), { 1 }, TENSOR_NOT_REQUIRED); + fc_norm = create_tensor(tn(LLM_TENSOR_FC_NORM, "weight"), { n_embd, n_embd_inp / n_embd }, TENSOR_NOT_REQUIRED); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm @@ -305,6 +306,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { return; } + const bool has_attn_gate = ml->get_tensor_meta("blk.0.attn_gate.weight") != nullptr; for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -315,8 +317,12 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, TENSOR_NOT_REQUIRED); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); - layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); - layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + if (has_attn_gate) { + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, n_head }, 0); + } else { + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + } layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED); layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), { n_embd }, TENSOR_NOT_REQUIRED); @@ -364,6 +370,11 @@ template <> llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { ggml_tensor * cur = build_inp_embd_enc(); + if (model.fc_norm) { + cur = ggml_reshape_3d(ctx0, cur, n_embd, model.target_layer_ids.size(), n_tokens); + cur = build_norm(cur, model.fc_norm, nullptr, LLM_NORM_RMS, -1); + cur = ggml_reshape_2d(ctx0, cur, model.fc->ne[0], n_tokens); + } cur = build_lora_mm(model.fc, cur, model.fc_s); cb(cur, "fc_out", -1); @@ -839,13 +850,16 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra int sections[4]; std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); - auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) { + auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos, int il) { + const int64_t n_rot_i = hparams.n_rot(il); + const float freq_base_i = model.get_rope_freq_base(cparams, il); + const float freq_scale_i = model.get_rope_freq_scale(cparams, il); return rope_type == GGML_ROPE_TYPE_MROPE ? ggml_rope_multi(ctx0, cur, pos, nullptr, - n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + n_rot_i, sections, rope_type, n_ctx_orig, freq_base_i, freq_scale_i, ext_factor, attn_factor, beta_fast, beta_slow) - : ggml_rope_ext(ctx0, cur, pos, model.layers[0].rope_freqs, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + : ggml_rope_ext(ctx0, cur, pos, model.layers[il].rope_freqs, + n_rot_i, rope_type, n_ctx_orig, freq_base_i, freq_scale_i, ext_factor, attn_factor, beta_fast, beta_slow); }; @@ -863,6 +877,11 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->add_input(std::move(inp)); // fuse the target features through the encoder + if (model.fc_norm) { + inp_target = ggml_reshape_3d(ctx0, inp_target, n_embd, model.target_layer_ids.size(), n_tokens); + inp_target = build_norm(inp_target, model.fc_norm, nullptr, LLM_NORM_RMS, -1); + inp_target = ggml_reshape_2d(ctx0, inp_target, n_embd_inp, n_tokens); + } ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s); inp_g = build_norm(inp_g, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); cb(inp_g, "inp_g_embeddings", -1); @@ -877,11 +896,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); - Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + if (layer.attn_k_norm) { + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + } if (shared_kv) { Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); } - Kcur = build_rope(Kcur, inp_pos); + Kcur = build_rope(Kcur, inp_pos, il); cb(Kcur, "Kcur_injected", il); cb(Vcur, "Vcur_injected", il); @@ -982,22 +1003,34 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); - Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); - Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + if (layer.attn_q_norm) { + Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + } if (shared_kv) { Vcur = ggml_rms_norm(ctx0, Vcur, hparams.f_norm_rms_eps); } - Qcur = build_rope(Qcur, inp_pos); - Kcur = build_rope(Kcur, inp_pos); + Qcur = build_rope(Qcur, inp_pos, il); + Kcur = build_rope(Kcur, inp_pos, il); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); // cache-aware, non-causal attention + ggml_tensor * wo = layer.wqkv_gate ? nullptr : layer.wo; ggml_tensor * cur = use_iswa - ? build_attn(inp_attn_iswa, layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il) - : build_attn(inp_attn, layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il); + ? build_attn(inp_attn_iswa, wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il) + : build_attn(inp_attn, wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il); + + if (layer.wqkv_gate) { + ggml_tensor * gate = ggml_sigmoid(ctx0, build_lora_mm(layer.wqkv_gate, noise_norm)); + cur = ggml_reshape_3d(ctx0, cur, n_embd_head, n_head, n_tokens); + gate = ggml_reshape_3d(ctx0, gate, 1, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, gate); + cur = ggml_reshape_2d(ctx0, cur, n_embd_head * n_head, n_tokens); + cur = build_lora_mm(layer.wo, cur, layer.wo_s); + } if (hparams.f_attn_value_scale != 0.0f) { cur = ggml_scale(ctx0, cur, hparams.f_attn_value_scale); @@ -1032,7 +1065,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra layer.ffn_gate, NULL, layer.ffn_gate_s, layer.ffn_down, NULL, layer.ffn_down_s, NULL, - hparams.llm_ffn_op, LLM_FFN_PAR, il); + layer.wqkv_gate ? LLM_FFN_GELU : hparams.llm_ffn_op, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); if (ffn_dynamic) { diff --git a/src/models/spark2-5.cpp b/src/models/spark2-5.cpp index 970b32ba603c..9a83d7c817f1 100644 --- a/src/models/spark2-5.cpp +++ b/src/models/spark2-5.cpp @@ -63,10 +63,13 @@ llama_model_spark2_5::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp_pos = build_inp_pos(); auto * inp_attn = build_attn_inp_kv_iswa(); ggml_tensor * inp_out_ids = build_inp_out_ids(); + const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]; const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + ggml_tensor * inpSA = inpL; ggml_tensor * cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); @@ -108,7 +111,7 @@ llama_model_spark2_5::graph::graph(const llama_model & model, const llm_graph_pa cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); cb(cur, "attn_out_proj", il); - if (il == n_layer - 1 && inp_out_ids) { + if (il == n_layer - 1 && inp_out_ids && !extract_final_inp) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -134,6 +137,10 @@ llama_model_spark2_5::graph::graph(const llama_model & model, const llm_graph_pa inpL = cur; } + res->t_layer_inp[n_layer] = inpL; + if (extract_final_inp && inp_out_ids) { + inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); + } ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1); cb(cur, "result_norm", -1); res->t_embd = cur; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 00267716de3f..bdbf23ce0392 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -9,6 +9,7 @@ // TODO: replace with #include "llama-ext.h" in the future #include "../src/llama-arch.h" +#include "../src/llama-ext.h" #include "../src/llama-model-saver.h" #include @@ -798,6 +799,40 @@ static bool check_causal_attn_toggle( return ok; } +static void test_spark_final_layer_features( + llama_model * model, llama_context * lctx, const std::vector & tokens) { + const uint32_t n_layer = llama_model_n_layer(model); + const uint32_t n_embd = llama_model_n_embd(model); + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + std::vector features_ref; + std::vector logits_ref; + common_batch batch(lctx); + + llama_set_embeddings_layer_inp(lctx, n_layer, true); + for (bool all_logits : {true, false}) { + llama_memory_clear(llama_get_memory(lctx), true); + batch.clear(); + for (size_t i = 0; i < tokens.size(); ++i) { + batch.add(tokens[i], i, 0, all_logits || i == tokens.size() - 1); + } + GGML_ASSERT(llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) == 0); + + const float * features = llama_get_embeddings_layer_inp(lctx, n_layer); + const float * logits = llama_get_logits_ith(lctx, -1); + GGML_ASSERT(features != nullptr && logits != nullptr); + const std::vector features_cur(features, features + tokens.size() * n_embd); + const std::vector logits_cur(logits, logits + n_vocab); + if (all_logits) { + features_ref = features_cur; + logits_ref = logits_cur; + } else { + GGML_ASSERT(nmse(features_ref, features_cur) < 1e-6); + GGML_ASSERT(nmse(logits_ref, logits_cur) < 1e-6); + } + } + llama_set_embeddings_layer_inp(lctx, n_layer, false); +} + static bool moe_mandatory(const llm_arch arch) { switch (arch) { case LLM_ARCH_FLINT: @@ -1175,6 +1210,12 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con test_executed = true; model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size); logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); + if (arch == LLM_ARCH_SPARK2_5) { + for (size_t n_tokens : {size_t(6), tokens.size()}) { + test_spark_final_layer_features(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), + std::vector(tokens.begin(), tokens.begin() + n_tokens)); + } + } const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); status_nmse = "\033[1;32mOK\033[0m";