diff --git a/conversion/__init__.py b/conversion/__init__.py index 3232a1050..5ae6ad819 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -109,6 +109,8 @@ TEXT_MODEL_MAP: dict[str, str] = { "GraniteSwitchForCausalLM": "granite", "GraniteSpeechForConditionalGeneration": "granite", "GraniteSpeechPlusForConditionalGeneration": "granite", + "GraniteSWAForCausalLM": "granite", + "GraniteMoeSWAForCausalLM": "granite", "Grok1ForCausalLM": "grok", "GrokForCausalLM": "grok", "GroveMoeForCausalLM": "grovemoe", diff --git a/conversion/granite.py b/conversion/granite.py index 5f1e3e847..796d37cca 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -74,6 +74,108 @@ class GraniteModel(LlamaModel): return super().filter_tensors(item) +@ModelBase.register("GraniteSWAForCausalLM") +class GraniteSWAModel(GraniteModel): + """Conversion for IBM's GraniteSWAForCausalLM (interleaved sliding window attention)""" + model_arch = gguf.MODEL_ARCH.GRANITE_SWA + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + if name.endswith("sinks"): + name += ".weight" + + return super().filter_tensors((name, gen)) + + def set_gguf_parameters(self): + """GraniteSWA uses Granite parameters plus sliding window configuration.""" + super().set_gguf_parameters() + + # Add sliding_window from config + sliding_window = self.hparams.get("sliding_window", 128) + self.gguf_writer.add_sliding_window(sliding_window) + logger.info("gguf: (granite_swa) sliding_window = %s", sliding_window) + + # Derive sliding_window_pattern from layer_types + if layer_types := self.hparams.get("layer_types"): + is_swa = [t == "sliding_attention" for t in layer_types] + self.gguf_writer.add_sliding_window_pattern(is_swa) + logger.info("gguf: (granite_swa) sliding_window_pattern = %d SWA layers / %d total", + sum(is_swa), len(is_swa)) + else: + # Fall back to period-based pattern: i % 4 != 0 + # This matches the transformers default pattern + n_layers = self.block_count + is_swa = [i % 4 != 0 for i in range(n_layers)] + self.gguf_writer.add_sliding_window_pattern(is_swa) + logger.info("gguf: (granite_swa) sliding_window_pattern (inferred) = %d SWA layers / %d total", + sum(is_swa), n_layers) + + # Add rope_pattern from no_rope_layers + if no_rope_layers := self.hparams.get("no_rope_layers"): + # Convert 1/0 to bool (1 = use RoPE, 0 = NoPE) + rope_pattern = [bool(x) for x in no_rope_layers] + self.gguf_writer.add_rope_pattern(rope_pattern) + logger.info("gguf: (granite_swa) rope_pattern = %d RoPE layers / %d total", + sum(rope_pattern), len(rope_pattern)) + + +@ModelBase.register("GraniteMoeSWAForCausalLM") +class GraniteMoeSWAModel(GraniteSWAModel): + """Conversion for IBM's GraniteMoeSWAForCausalLM (unified dense + MoE with iSWA)""" + model_arch = gguf.MODEL_ARCH.GRANITE_SWA + + def set_gguf_parameters(self): + super().set_gguf_parameters() + if shared_intermediate_size := self.hparams.get("shared_intermediate_size"): + self.gguf_writer.add_expert_shared_feed_forward_length(shared_intermediate_size) + logger.info("gguf: (granitemoewa) shared_intermediate_size = %s", shared_intermediate_size) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + """Split merged MoE tensors (gate+up) following standard MoE pattern.""" + + # Handle expert FFN tensors (merged gate+up) - swash format: experts.gate_up_proj + # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps + # tensor for the routed experts. + if name.endswith("block_sparse_moe.experts.gate_up_proj"): + ffn_dim = self.hparams["intermediate_size"] + assert data_torch.shape[-2] == 2 * ffn_dim, f"Merged FFN tensor size must be 2 * intermediate_size, got {data_torch.shape[-2]}" + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid) + return + + # Handle expert FFN down projection - swash format: experts.down_proj + if name.endswith("block_sparse_moe.experts.down_proj"): + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), bid) + return + + # Handle expert FFN tensors (merged gate+up) - standard granite format: input_linear.weight + # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps + # tensor for the routed experts. + if name.endswith("block_sparse_moe.input_linear.weight"): + ffn_dim = self.hparams["intermediate_size"] + assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * intermediate_size" + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid) + return + + # Handle shared expert FFN tensors (if present) - kept fused since + # inference (build_ffn) supports a single ffn_up_shexp tensor with + # LLM_FFN_SWIGLU for the shared expert. + if name.endswith("shared_mlp.input_linear.weight"): + ffn_dim = self.hparams.get("shared_intermediate_size", self.hparams["intermediate_size"]) + assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * shared_intermediate_size" + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), bid) + return + + # Handle shared expert output (if present) + if name.endswith("shared_mlp.output_linear.weight"): + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, bid), bid) + return + + # Pass through to parent for all other tensors (including sinks) + yield from super().modify_tensors(data_torch, name, bid) + + @ModelBase.register("GraniteMoeForCausalLM", "GraniteMoeSharedForCausalLM") @ModelBase.example("ibm-granite/granite-3.1-3b-a800m-instruct") class GraniteMoeModel(GraniteModel): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d043c9b6e..fad8d1fd8 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -208,6 +208,7 @@ class Keys: SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers" SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" + ROPE_PATTERN = "{arch}.attention.rope_pattern" class Indexer: HEAD_COUNT = "{arch}.attention.indexer.head_count" @@ -549,6 +550,7 @@ class MODEL_ARCH(IntEnum): GRANITE_MOE = auto() GRANITE_HYBRID = auto() GRANITE_SWITCH = auto() + GRANITE_SWA = auto() CHAMELEON = auto() WAVTOKENIZER_DEC = auto() PLM = auto() @@ -1265,6 +1267,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.GRANITE_MOE: "granitemoe", MODEL_ARCH.GRANITE_HYBRID: "granitehybrid", MODEL_ARCH.GRANITE_SWITCH: "graniteswitch", + MODEL_ARCH.GRANITE_SWA: "granite_swa", MODEL_ARCH.CHAMELEON: "chameleon", MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec", MODEL_ARCH.PLM: "plm", @@ -4152,6 +4155,31 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.GRANITE_SWA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + # MoE (GraniteMoeSWA) + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + # Shared expert - gate+up kept fused in FFN_UP_SHEXP (LLM_FFN_SWIGLU) + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], MODEL_ARCH.CHAMELEON: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 9e0914fd8..16ae9f999 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -824,6 +824,9 @@ class GGUFWriter: else: self.add_array(key, value) + def add_rope_pattern(self, value: Sequence[bool]) -> None: + self.add_array(Keys.Attention.ROPE_PATTERN.format(arch=self.arch), value) + def add_dense_features_dims(self, dense:str, in_f:int, out_f:int) -> None: self.add_uint32(Keys.LLM.DENSE_FEAT_IN_SIZE.format(arch=self.arch, dense=dense), in_f) self.add_uint32(Keys.LLM.DENSE_FEAT_OUT_SIZE.format(arch=self.arch, dense=dense), out_f) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 3292942b4..a0571ccd3 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -458,6 +458,7 @@ class TensorNameMap: "transformer.decoder_layer.{bid}.router", # Grok "transformer.blocks.{bid}.ffn.router.layer", # dbrx "model.layers.{bid}.block_sparse_moe.router.layer", # granitemoe + "model.layers.{bid}.block_sparse_moe.router", # granite_swa "model.layers.{bid}.feed_forward.router", # llama4 jamba "encoder.layers.{bid}.mlp.router.layer", # nomic-bert-moe "model.layers.{bid}.mlp.router", # openai-moe diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5b88bde14..955c2d796 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -102,6 +102,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_GRANITE_HYBRID, "granitehybrid" }, { LLM_ARCH_GRANITE_SWITCH, "graniteswitch" }, + { LLM_ARCH_GRANITE_SWA, "granite_swa" }, { LLM_ARCH_CHAMELEON, "chameleon" }, { LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" }, { LLM_ARCH_PLM, "plm" }, @@ -261,6 +262,8 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, "%s.attention.relative_buckets_count" }, { LLM_KV_ATTENTION_SLIDING_WINDOW, "%s.attention.sliding_window" }, { LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, "%s.attention.sliding_window_pattern" }, + { LLM_KV_ATTENTION_ROPE_PATTERN, "%s.attention.rope_pattern" }, + { LLM_KV_ATTENTION_SCALE, "%s.attention.scale" }, { LLM_KV_ATTENTION_OUTPUT_SCALE, "%s.attention.output_scale" }, { LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 8042120a2..48fe051a9 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -107,6 +107,7 @@ enum llm_arch { LLM_ARCH_GRANITE_MOE, LLM_ARCH_GRANITE_HYBRID, LLM_ARCH_GRANITE_SWITCH, + LLM_ARCH_GRANITE_SWA, LLM_ARCH_CHAMELEON, LLM_ARCH_WAVTOKENIZER_DEC, LLM_ARCH_PLM, @@ -267,6 +268,8 @@ enum llm_kv { LLM_KV_ATTENTION_SLIDING_WINDOW, LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, LLM_KV_ATTENTION_SCALE, + LLM_KV_ATTENTION_ROPE_PATTERN, + LLM_KV_ATTENTION_OUTPUT_SCALE, LLM_KV_ATTENTION_VALUE_SCALE, LLM_KV_ATTENTION_TEMPERATURE_LENGTH, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index e3f0cf0ed..cbe31134f 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -291,7 +291,11 @@ bool llama_hparams::has_rope(uint32_t il) const { return false; } - return true; + if (il < n_layer_all) { + return rope_pattern[il] != 0; + } + + GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all); } uint32_t llama_hparams::n_layer() const { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index e91ce1cc3..f6af36436 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -144,6 +144,10 @@ struct llama_hparams { std::array rope_sections; + // Per-layer RoPE enable flags (1 = use RoPE, 0 = NoPE) + // by default, all layers use RoPE (controlled by rope_finetuned) + std::array rope_pattern; + // Sliding Window Attention (SWA) llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE; // the size of the sliding window (0 - no SWA) diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index be9524d40..b9e0a6009 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -30,6 +30,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_MELLUM: case LLM_ARCH_LAGUNA: + case LLM_ARCH_GRANITE_SWA: return false; default: return true; @@ -272,6 +273,7 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_VALUE_RESIDUAL_MIX_LORA_RANK, hparams.n_lora_value_res_mix); add_kv(LLM_KV_ATTENTION_GATE_LORA_RANK, hparams.n_lora_gate); add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, hparams.n_rel_attn_bkts); + add_kv(LLM_KV_ATTENTION_ROPE_PATTERN, hparams.rope_pattern, true); add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); // add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, ???); add_kv(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0d74a2135..3759c8625 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -246,6 +246,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_minicpm(params); case LLM_ARCH_GRANITE_HYBRID: return new llama_model_granite_hybrid(params); + case LLM_ARCH_GRANITE_SWA: + return new llama_model_granite_swa(params); case LLM_ARCH_CHAMELEON: return new llama_model_chameleon(params); case LLM_ARCH_WAVTOKENIZER_DEC: @@ -1157,6 +1159,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_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); std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0); std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), llm_arch_is_recurrent(ml.get_arch()) ? 1 : 0); std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0); @@ -2639,6 +2642,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GRANITE_MOE: case LLM_ARCH_GRANITE_HYBRID: case LLM_ARCH_GRANITE_SWITCH: + case LLM_ARCH_GRANITE_SWA: case LLM_ARCH_CHAMELEON: case LLM_ARCH_BAILINGMOE: case LLM_ARCH_BAILINGMOE3: diff --git a/src/models/granite-hybrid.cpp b/src/models/granite-hybrid.cpp index eb23095ae..8a8f7e19f 100644 --- a/src/models/granite-hybrid.cpp +++ b/src/models/granite-hybrid.cpp @@ -16,7 +16,8 @@ void llama_model_granite_hybrid::load_arch_hparams(llama_model_loader & ml) { // Granite uses rope_finetuned as a switch for rope, so default to true bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; + hparams.rope_finetuned = rope_finetuned; // needed for round trip save + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned); // A layer is recurrent IFF the n_head_kv value is set to 0 for (uint32_t i = 0; i < hparams.n_layer(); ++i) { @@ -147,7 +148,7 @@ llama_model_granite_hybrid::graph::graph(const llama_model & model, const llm_gr // Positional embeddings populated if rope enabled ggml_tensor * inp_pos = nullptr; - if (hparams.rope_finetuned) { + if (hparams.has_rope(0)) { inp_pos = build_inp_pos(); } @@ -206,8 +207,7 @@ ggml_tensor * llama_model_granite_hybrid::graph::build_attention_layer(ggml_tens const int il) { auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); - const bool use_rope = hparams.rope_finetuned; - if (use_rope) { + if (hparams.has_rope(il)) { ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); diff --git a/src/models/granite-moe.cpp b/src/models/granite-moe.cpp index 115263c41..09be49393 100644 --- a/src/models/granite-moe.cpp +++ b/src/models/granite-moe.cpp @@ -7,11 +7,6 @@ void llama_model_granite_moe::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); - // Granite uses rope_finetuned as a switch for rope, so default to true - bool rope_finetuned = true; - ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; - switch (hparams.n_layer()) { case 32: type = LLM_TYPE_3B; break; case 40: type = LLM_TYPE_3B; break; diff --git a/src/models/granite-swa.cpp b/src/models/granite-swa.cpp new file mode 100644 index 000000000..3aa2b63b2 --- /dev/null +++ b/src/models/granite-swa.cpp @@ -0,0 +1,319 @@ +#include "models.h" + +#include + +void llama_model_granite_swa::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); + ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false); + ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); + ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); + + // MoE expert configuration + ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert, false); + ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used, false); + + // iSWA configuration + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + + // Granite4 Vision uses array deepstack_mapping + ml.get_arr(LLM_KV_DEEPSTACK_MAPPING, hparams.deepstack_mapping_arr, false); + + // Count the unique deepstack input indices + std::unordered_set unique_deepstack_idxs; + for (const auto val : hparams.deepstack_mapping_arr) { + if (val >= 0) { + unique_deepstack_idxs.insert(val); + } + } + hparams.n_deepstack_layers = unique_deepstack_idxs.size(); + + // Ensure all values are valid (avoid overflow attacks) + for (const auto val : unique_deepstack_idxs) { + if (val > hparams.n_deepstack_layers) { + std::stringstream ss; + ss << "Invalid deepstack index: " << val << " > " << hparams.n_deepstack_layers; + throw std::runtime_error(ss.str()); + } + } + + // Per-layer RoPE pattern (optional) + ml.get_arr(LLM_KV_ATTENTION_ROPE_PATTERN, hparams.rope_pattern, false); + + switch (hparams.n_layer()) { + case 32: type = LLM_TYPE_3B; break; + case 40: type = LLM_TYPE_3B; break; + // Add additional layer/vocab/etc checks here for other model sizes + default: type = LLM_TYPE_UNKNOWN; + } + + // For Granite MoE Shared + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); +} + +void llama_model_granite_swa::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // output + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + + // if output is NULL, init from the input tok embed + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + // optional bias tensors + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); + + // Per-layer attention sinks for iSWA + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) { + layer.rope_long = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + } + else { + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + } + + if (n_expert == 0) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + + // optional MLP bias + layer.ffn_gate_b = create_tensor(tn(LLM_TENSOR_FFN_GATE, "bias", i), {n_ff}, TENSOR_NOT_REQUIRED); + 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), {n_ff}, TENSOR_NOT_REQUIRED); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff, n_embd, n_expert}, 0); + create_tensor_gate_up_exps(layer, i, n_embd, n_ff, n_expert, 0); + + // For Granite MoE Shared - gate+up kept fused in ffn_up_shexp (see LLM_FFN_SWIGLU below) + if (hparams.n_ff_shexp > 0) { + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, 2*hparams.n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, 0); + } + } + } +} + +std::unique_ptr llama_model_granite_swa::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_granite_swa::graph::graph( + const llama_model & model, + const llm_graph_params & params) + : llm_graph_context(params) { + + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + // inp_pos - built only if rope enabled + ggml_tensor * inp_pos = build_inp_pos(); + auto * inp_attn = build_attn_inp_kv_iswa(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + + // Granite Vision 4.1 deepstack: inject the projector stream that + // targets decoder layer `il` before the decoder runs. + // NOTE: skip the first deepstack layer since that's inpL + const auto & deepstack_emb_idx = hparams.deepstack_mapping_arr[il]; + if (il > 0 && deepstack_emb_idx >= 0) { + ggml_tensor * ds = ggml_view_2d(ctx0, + res->t_inp_embd, n_embd, n_tokens, + res->t_inp_embd->nb[1], + deepstack_emb_idx * n_embd * sizeof(float)); + inpL = ggml_add(ctx0, inpL, ds); + cb(inpL, "deepstack_in", il); + } + + ggml_tensor * inpSA = inpL; + + // norm + cur = build_norm(inpL, + model.layers[il].attn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self-attention + cur = build_attention_layer( + cur, inp_pos, inp_attn, + model, n_embd_head, il); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + // ffn + cur = build_layer_ffn(cur, inpSA, model, il); + + // input for next layer + inpL = cur; + } + cur = inpL; + + cur = build_norm(cur, + model.output_norm, NULL, + LLM_NORM_RMS, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head + cur = build_lora_mm(model.output, cur, model.output_s); + + // For Granite architectures - scale logits + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +ggml_tensor * llama_model_granite_swa::graph::build_attention_layer( + ggml_tensor * cur, + ggml_tensor * inp_pos, + llm_graph_input_attn_kv_iswa * inp_attn, + const llama_model & model, + const int64_t n_embd_head, + const int il) { + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); + + const bool use_rope = hparams.has_rope(il); + if (use_rope) { + ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + } + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + // Pass layer.attn_sinks to build_attn for sink-based attention modulation + cur = build_attn(inp_attn, + model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, model.layers[il].attn_sinks, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + return cur; +} + +ggml_tensor * llama_model_granite_swa::graph::build_layer_ffn( + ggml_tensor * cur, + ggml_tensor * inpSA, + const llama_model & model, + const int il) { + + // For Granite architectures - scale residual + if (hparams.f_residual_scale) { + cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); + } + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + // feed-forward network (non-MoE) + if (model.layers[il].ffn_gate_inp == nullptr) { + + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL, + model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL, + model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + } else { + // MoE branch + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + ggml_tensor * moe_out = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, + il, + nullptr, model.layers[il].ffn_gate_up_exps); + cb(moe_out, "ffn_moe_out", il); + + // For Granite MoE Shared - gate+up kept fused in ffn_up_shexp + if (hparams.n_ff_shexp > 0) { + ggml_tensor * ffn_shexp = build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, NULL, + NULL, NULL, NULL, + model.layers[il].ffn_down_shexp, NULL, NULL, + NULL, + LLM_FFN_SWIGLU, LLM_FFN_SEQ, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } else { + cur = moe_out; + } + } + + // For Granite architectures - scale residual + if (hparams.f_residual_scale) { + cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); + } + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "ffn_out", il); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + return cur; +} diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 80f6b86ed..7c9a901c8 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -11,7 +11,8 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; + hparams.rope_finetuned = rope_finetuned; // needed for round trip save + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned); switch (hparams.n_layer()) { case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break; @@ -254,7 +255,7 @@ llama_model_granite_switch::graph::graph( cb(inpL, "inp_embd", -1); ggml_tensor * inp_pos = nullptr; - if (hparams.rope_finetuned) { + if (hparams.has_rope(0)) { inp_pos = build_inp_pos(); } auto * inp_attn = build_attn_inp_kv(); @@ -361,7 +362,7 @@ ggml_tensor * llama_model_granite_switch::graph::build_attention_layer( 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); - if (hparams.rope_finetuned) { + if (hparams.has_rope(il)) { ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, diff --git a/src/models/granite.cpp b/src/models/granite.cpp index 4a75c5ff3..9e9f97e94 100644 --- a/src/models/granite.cpp +++ b/src/models/granite.cpp @@ -33,7 +33,8 @@ void llama_model_granite::load_arch_hparams(llama_model_loader & ml) { // Granite uses rope_finetuned as a switch for rope, so default to true bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; + hparams.rope_finetuned = rope_finetuned; // needed for round trip save + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned); switch (hparams.n_layer()) { case 32: type = LLM_TYPE_3B; break; @@ -127,7 +128,7 @@ llama_model_granite::graph::graph( // inp_pos - built only if rope enabled ggml_tensor * inp_pos = nullptr; - if (hparams.rope_finetuned) { + if (hparams.has_rope(0)) { inp_pos = build_inp_pos(); } auto * inp_attn = build_attn_inp_kv(); @@ -203,8 +204,7 @@ ggml_tensor * llama_model_granite::graph::build_attention_layer( auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); - const bool use_rope = hparams.rope_finetuned; - if (use_rope) { + if (hparams.has_rope(il)) { ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); Qcur = ggml_rope_ext( ctx0, Qcur, inp_pos, rope_factors, diff --git a/src/models/models.h b/src/models/models.h index 180b30a46..1dd30dfd1 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1719,6 +1719,34 @@ struct llama_model_granite_hybrid : public llama_model_base { }; +struct llama_model_granite_swa : public llama_model_base { + llama_model_granite_swa(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + private: + ggml_tensor * build_attention_layer( + ggml_tensor * cur, + ggml_tensor * inp_pos, + llm_graph_input_attn_kv_iswa * inp_attn, + const llama_model & model, + const int64_t n_embd_head, + const int il); + + ggml_tensor * build_layer_ffn( + ggml_tensor * cur, + ggml_tensor * inpSA, + const llama_model & model, + const int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_chameleon : public llama_model_base { llama_model_chameleon(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 032ac3537..4eb3763cb 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -197,7 +197,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA, 10000.0f); // SWA pattern: every 5th layer is full attention (matches E2B layer_types) ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5)); - } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER) { + } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_GRANITE_SWA) { std::vector pattern; pattern.reserve(n_layer); for (uint32_t il = 0; il < n_layer; il++) {