Felipe97/llama-cpp-compiled
01.1k
1from __future__ import annotations2 3import re4from typing import Any, Callable, Iterable, TYPE_CHECKING5 6import torch7 8if TYPE_CHECKING:9 from torch import Tensor10 11from .base import MmprojModel, ModelBase, gguf, logger12 13from .llama import LlamaModel14from .mamba import Mamba2Model15 16 17@ModelBase.register("GraniteForCausalLM")18@ModelBase.example("ibm-granite/granite-3.3-2b-instruct")19class GraniteModel(LlamaModel):20 """Conversion for IBM's GraniteForCausalLM"""21 model_arch = gguf.MODEL_ARCH.GRANITE22 23 def set_gguf_parameters(self):24 """Granite uses standard llama parameters with the following differences:25 26 - No head_dim support27 - New multiplier params:28 - attention_scale29 - embedding_scale30 - residual_scale31 - logits_scaling32 """33 if head_dim := self.hparams.pop("head_dim", None):34 logger.warning("Ignoring head_dim (%s) from config for Granite", head_dim)35 super().set_gguf_parameters()36 # NOTE: Convert _multiplier params to _scale params for naming37 # consistency38 if attention_scale := self.hparams.get("attention_multiplier"):39 self.gguf_writer.add_attention_scale(attention_scale)40 logger.info("gguf: (granite) attention_scale = %s", attention_scale)41 if embedding_scale := self.hparams.get("embedding_multiplier"):42 self.gguf_writer.add_embedding_scale(embedding_scale)43 logger.info("gguf: (granite) embedding_scale = %s", embedding_scale)44 if residual_scale := self.hparams.get("residual_multiplier"):45 self.gguf_writer.add_residual_scale(residual_scale)46 logger.info("gguf: (granite) residual_scale = %s", residual_scale)47 if logits_scale := self.hparams.get("logits_scaling"):48 self.gguf_writer.add_logit_scale(logits_scale)49 logger.info("gguf: (granite) logits_scale = %s", logits_scale)50 51 # If being used as the base for Granite4 Vision, add deepstack_layer_arr52 if self.hparams.get("spatial_target_layers") or self.hparams.get("deepstack_layer_map"):53 normalized_projector_map = Granite4VisionMmprojModel.get_normalized_projector_map(self.hparams)54 deepstack_mapping_arr = [-1 for _ in range(self.block_count)] # Populate with -1 sentinels55 for proj_idx, (_, llm_layer, _, _) in enumerate(normalized_projector_map):56 # Skip the first projector which is handled as the base embedding57 # stream like normal58 if proj_idx == 0:59 continue60 deepstack_mapping_arr[llm_layer] = proj_idx61 self.gguf_writer.add_deepstack_mapping(deepstack_mapping_arr)62 63 @classmethod64 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:65 name, gen = item66 # Skip multimodal tensors67 if (68 name.startswith(("encoder."))69 or "image_" in name70 or "layerwise_projectors" in name71 or "spatial_projectors" in name72 ):73 return74 return super().filter_tensors(item)75 76 77@ModelBase.register("GraniteSWAForCausalLM")78class GraniteSWAModel(GraniteModel):79 """Conversion for IBM's GraniteSWAForCausalLM (interleaved sliding window attention)"""80 model_arch = gguf.MODEL_ARCH.GRANITE_SWA81 82 @classmethod83 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:84 name, gen = item85 86 if name.endswith("sinks"):87 name += ".weight"88 89 return super().filter_tensors((name, gen))90 91 def set_gguf_parameters(self):92 """GraniteSWA uses Granite parameters plus sliding window configuration."""93 super().set_gguf_parameters()94 95 # Add sliding_window from config96 sliding_window = self.hparams.get("sliding_window", 128)97 self.gguf_writer.add_sliding_window(sliding_window)98 logger.info("gguf: (granite_swa) sliding_window = %s", sliding_window)99 100 # Derive sliding_window_pattern from layer_types101 if layer_types := self.hparams.get("layer_types"):102 is_swa = [t == "sliding_attention" for t in layer_types]103 self.gguf_writer.add_sliding_window_pattern(is_swa)104 logger.info("gguf: (granite_swa) sliding_window_pattern = %d SWA layers / %d total",105 sum(is_swa), len(is_swa))106 else:107 # Fall back to period-based pattern: i % 4 != 0108 # This matches the transformers default pattern109 n_layers = self.block_count110 is_swa = [i % 4 != 0 for i in range(n_layers)]111 self.gguf_writer.add_sliding_window_pattern(is_swa)112 logger.info("gguf: (granite_swa) sliding_window_pattern (inferred) = %d SWA layers / %d total",113 sum(is_swa), n_layers)114 115 # Add rope_pattern from no_rope_layers116 if no_rope_layers := self.hparams.get("no_rope_layers"):117 # Convert 1/0 to bool (1 = use RoPE, 0 = NoPE)118 rope_pattern = [bool(x) for x in no_rope_layers]119 self.gguf_writer.add_rope_pattern(rope_pattern)120 logger.info("gguf: (granite_swa) rope_pattern = %d RoPE layers / %d total",121 sum(rope_pattern), len(rope_pattern))122 123 124@ModelBase.register("GraniteMoeSWAForCausalLM")125class GraniteMoeSWAModel(GraniteSWAModel):126 """Conversion for IBM's GraniteMoeSWAForCausalLM (unified dense + MoE with iSWA)"""127 model_arch = gguf.MODEL_ARCH.GRANITE_SWA128 129 def set_gguf_parameters(self):130 super().set_gguf_parameters()131 if shared_intermediate_size := self.hparams.get("shared_intermediate_size"):132 self.gguf_writer.add_expert_shared_feed_forward_length(shared_intermediate_size)133 logger.info("gguf: (granitemoewa) shared_intermediate_size = %s", shared_intermediate_size)134 135 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:136 """Split merged MoE tensors (gate+up) following standard MoE pattern."""137 138 # Handle expert FFN tensors (merged gate+up) - swash format: experts.gate_up_proj139 # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps140 # tensor for the routed experts.141 if name.endswith("block_sparse_moe.experts.gate_up_proj"):142 ffn_dim = self.hparams["intermediate_size"]143 assert data_torch.shape[-2] == 2 * ffn_dim, f"Merged FFN tensor size must be 2 * intermediate_size, got {data_torch.shape[-2]}"144 yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid)145 return146 147 # Handle expert FFN down projection - swash format: experts.down_proj148 if name.endswith("block_sparse_moe.experts.down_proj"):149 yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), bid)150 return151 152 # Handle expert FFN tensors (merged gate+up) - standard granite format: input_linear.weight153 # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps154 # tensor for the routed experts.155 if name.endswith("block_sparse_moe.input_linear.weight"):156 ffn_dim = self.hparams["intermediate_size"]157 assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * intermediate_size"158 yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid)159 return160 161 # Handle shared expert FFN tensors (if present) - kept fused since162 # inference (build_ffn) supports a single ffn_up_shexp tensor with163 # LLM_FFN_SWIGLU for the shared expert.164 if name.endswith("shared_mlp.input_linear.weight"):165 ffn_dim = self.hparams.get("shared_intermediate_size", self.hparams["intermediate_size"])166 assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * shared_intermediate_size"167 yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), bid)168 return169 170 # Handle shared expert output (if present)171 if name.endswith("shared_mlp.output_linear.weight"):172 yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, bid), bid)173 return174 175 # Pass through to parent for all other tensors (including sinks)176 yield from super().modify_tensors(data_torch, name, bid)177 178 179@ModelBase.register("GraniteMoeForCausalLM", "GraniteMoeSharedForCausalLM")180@ModelBase.example("ibm-granite/granite-3.1-3b-a800m-instruct")181class GraniteMoeModel(GraniteModel):182 """Conversion for IBM's GraniteMoeForCausalLM"""183 model_arch = gguf.MODEL_ARCH.GRANITE_MOE184 185 def set_gguf_parameters(self):186 """GraniteMoeShared uses GraniteMoe parameters plus the following:187 - shared_intermediate_size188 """189 super().set_gguf_parameters()190 if shared_feed_forward_length := self.hparams.get("shared_intermediate_size"):191 self.gguf_writer.add_expert_shared_feed_forward_length(shared_feed_forward_length)192 logger.info("gguf: (granitemoeshared) shared_feed_forward_length = %s", shared_feed_forward_length)193 194 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:195 """In modeling_granitemoe, the JetMoe implementation of parallel experts196 is used. This essentially merges w1 and w3 into a single tensor with 2x197 the hidden size that is then split during forward. To keep compatibility198 with existing mixtral support, we pull them apart here.199 """200 201 if name.endswith("block_sparse_moe.input_linear.weight"):202 ffn_dim = self.hparams["intermediate_size"]203 assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * intermediate_size"204 gate, up = data_torch.split(ffn_dim, dim=-2)205 yield from ModelBase.modify_tensors(self, gate, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), bid)206 yield from ModelBase.modify_tensors(self, up, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), bid)207 return208 209 has_experts = bool(self.hparams.get('num_local_experts'))210 211 if name.endswith("shared_mlp.input_linear.weight"):212 ffn_dim = self.hparams["shared_intermediate_size"]213 assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * shared_intermediate_size"214 gate, up = data_torch.split(ffn_dim, dim=-2)215 if has_experts:216 yield from ModelBase.modify_tensors(self, gate,self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_SHEXP, bid), bid)217 yield from ModelBase.modify_tensors(self, up, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), bid)218 return219 yield from ModelBase.modify_tensors(self, gate, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), bid)220 yield from ModelBase.modify_tensors(self, up, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), bid)221 return222 223 if not has_experts and name.endswith("shared_mlp.output_linear.weight"):224 yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN, bid), bid)225 return226 227 yield from super().modify_tensors(data_torch, name, bid)228 229 230@ModelBase.register("GraniteSwitchForCausalLM")231@ModelBase.example("ibm-granite/granite-switch-4.1-3b-preview")232class GraniteSwitchModel(GraniteMoeModel):233 """Dense, all-attention Granite with N per-token embedded LoRA adapters, stacked234 over the adapter dim with a zero adapter at slot 0 (N = num_adapters + 1)."""235 model_arch = gguf.MODEL_ARCH.GRANITE_SWITCH236 237 # permute q/k per-slice below (NORM-rope layout), not via the parent's auto-permute238 undo_permute = False239 240 def __init__(self, *args, **kwargs):241 super().__init__(*args, **kwargs)242 # the weightless switch reserves one cache slot: one fewer block than num_hidden_layers243 self.block_count = self.block_count - 1244 self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)245 246 self._n_adapters = int(self.hparams["num_adapters"])247 self._max_lora_rank = int(self.hparams["max_lora_rank"])248 self._n_slots = self._n_adapters + 1 # +1 for the zero slot at index 0249 250 n_head = int(self.hparams["num_attention_heads"])251 n_kv_head = int(self.hparams["num_key_value_heads"])252 head_dim = (253 self.hparams.get("projection_head_dim")254 or self.hparams.get("head_dim")255 or (self.hparams["hidden_size"] // n_head)256 )257 self._n_head = n_head258 self._n_kv_head = n_kv_head259 self._head_dim = int(head_dim)260 self._q_size = n_head * self._head_dim261 self._kv_size = n_kv_head * self._head_dim262 263 def set_gguf_parameters(self):264 super().set_gguf_parameters()265 266 # dense: pin expert_used_count to 0 (config carries a leftover num_experts_per_tok)267 if not self.hparams.get("num_local_experts"):268 self.gguf_writer.add_expert_used_count(0)269 270 self.gguf_writer.add_adapter_count(self._n_adapters)271 self.gguf_writer.add_adapter_lora_rank(self._max_lora_rank)272 self.gguf_writer.add_adapter_token_ids_activate(self.hparams["adapter_token_ids"])273 self.gguf_writer.add_adapter_token_ids_substitute(self.hparams["adapter_substitute_token_ids"])274 router_gain = float(self.hparams.get("control_token_gain", 15.0))275 self.gguf_writer.add_adapter_router_gain(router_gain)276 logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain)277 278 def _lora_a(self, data: Tensor) -> Tensor:279 # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in]280 a = data.squeeze(1)281 zero = torch.zeros_like(a[:1])282 return torch.cat([zero, a], dim=0).contiguous()283 284 def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor:285 # on-disk B: [n_adapters, 1, out, max_rank] -> [n_adapters+1, out, max_rank]286 b = data.squeeze(1)287 if permute_n_head is not None:288 # permute each adapter's B output rows to match the permuted q/k base289 b = torch.stack([self.permute(b[i], permute_n_head, permute_n_head) for i in range(b.shape[0])], dim=0)290 zero = torch.zeros_like(b[:1])291 return torch.cat([zero, b], dim=0).contiguous()292 293 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:294 T = gguf.MODEL_TENSOR295 296 # skip the weightless switch + control-token buffers (rebuilt at load time)297 bare = name.split(".")[-1]298 if (299 name.startswith("model.switch.") or name.startswith("switch.")300 or bare in ("adapter_token_ids", "control_to_substitute_lut")301 ):302 return303 304 if "self_attn.qkv_proj" in name:305 if name.endswith("base_layer.weight"):306 # fused [q|k|v] rows: permute q/k row-blocks for ggml's NORM-rope layout307 q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0)308 q = self.permute(q, self._n_head, self._n_head)309 k = self.permute(k, self._n_kv_head, self._n_kv_head)310 fused = torch.cat([q, k, v], dim=0)311 yield (self.format_tensor_name(T.ATTN_QKV, bid), fused)312 return313 if "lora_A_slices." in name:314 slot = int(name.rsplit(".", 1)[1])315 key = {0: T.ATTN_Q, 1: T.ATTN_K, 2: T.ATTN_V}[slot]316 yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))317 return318 if "lora_B_slices." in name:319 slot = int(name.rsplit(".", 1)[1])320 key, ph = {321 0: (T.ATTN_Q, self._n_head),322 1: (T.ATTN_K, self._n_kv_head),323 2: (T.ATTN_V, None),324 }[slot]325 yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch, ph))326 return327 raise ValueError(f"Unexpected qkv_proj tensor: {name}")328 329 if "self_attn.o_proj" in name:330 if name.endswith("base_layer.weight"):331 yield (self.format_tensor_name(T.ATTN_OUT, bid), data_torch)332 return333 if name.endswith("lora_A"):334 yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_a"), self._lora_a(data_torch))335 return336 if name.endswith("lora_B"):337 yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_b"), self._lora_b(data_torch))338 return339 raise ValueError(f"Unexpected o_proj tensor: {name}")340 341 if "shared_mlp.input_linear" in name:342 ffn = self.hparams["shared_intermediate_size"]343 if name.endswith("base_layer.weight"):344 gate, up = data_torch.split([ffn, ffn], dim=0)345 yield (self.format_tensor_name(T.FFN_GATE, bid), gate)346 yield (self.format_tensor_name(T.FFN_UP, bid), up)347 return348 if "lora_A_slices." in name:349 slot = int(name.rsplit(".", 1)[1])350 key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]351 yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))352 return353 if "lora_B_slices." in name:354 slot = int(name.rsplit(".", 1)[1])355 key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]356 yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch))357 return358 raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}")359 360 if "shared_mlp.output_linear" in name:361 if name.endswith("base_layer.weight"):362 yield (self.format_tensor_name(T.FFN_DOWN, bid), data_torch)363 return364 if name.endswith("lora_A"):365 yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_a"), self._lora_a(data_torch))366 return367 if name.endswith("lora_B"):368 yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_b"), self._lora_b(data_torch))369 return370 raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}")371 372 if bid is not None and ".layers." in name and (373 "input_layernorm" in name or "post_attention_layernorm" in name374 ):375 key = T.ATTN_NORM if "input_layernorm" in name else T.FFN_NORM376 yield (self.format_tensor_name(key, bid), data_torch)377 return378 379 if name in ("model.embed_tokens.weight", "embed_tokens.weight"):380 yield (self.format_tensor_name(T.TOKEN_EMBD), data_torch)381 return382 if name in ("model.norm.weight", "norm.weight"):383 yield (self.format_tensor_name(T.OUTPUT_NORM), data_torch)384 return385 if name == "lm_head.weight":386 return # tied to token_embd387 388 raise ValueError(f"graniteswitch: unhandled tensor {name!r} (bid={bid})")389 390 391@ModelBase.register("GraniteMoeHybridForCausalLM", "BambaForCausalLM")392@ModelBase.example("ibm-granite/granite-4.0-h-tiny", "ibm-ai-platform/Bamba-9B-v2")393class GraniteHybridModel(Mamba2Model, GraniteMoeModel):394 """GraniteHybrid is a hybrid SSM + Attention model that uses Mamba2 SSM395 layers and optionally uses MoE w/ a shared expert"""396 model_arch = gguf.MODEL_ARCH.GRANITE_HYBRID397 undo_permute = True398 399 def __init__(self, *args, **kwargs):400 401 # Hybrid mamba models use a prefix for the mamba-specific params.402 # TODO: Extend this if the prefix(es) need to be configurable403 self.hparam_prefixes = ["mamba"]404 405 super().__init__(*args, **kwargs)406 407 # Lists of which layers use ssm vs attention408 self._attn_layers = self.get_attn_layers()409 self._ssm_layers = [410 i for i in range(self.block_count)411 if i not in self._attn_layers412 ]413 414 # There are some models in this family that are non-hybrid, but keep the415 # same parent class by setting all layers to "attention." If this is the416 # case, the model architecture needs to be updated to a standard417 # "granite" or "granitemoe" model418 if not self._ssm_layers:419 has_experts = self.find_hparam(["num_experts_per_tok", "num_experts_per_token"], optional=True)420 new_arch = (421 gguf.MODEL_ARCH.GRANITE_MOE422 if has_experts else423 gguf.MODEL_ARCH.GRANITE424 )425 self.model_arch = new_arch426 self.gguf_writer.arch = gguf.MODEL_ARCH_NAMES[new_arch]427 self.gguf_writer.add_architecture()428 429 # n_group and d_inner are used during reshape_tensors for mamba2430 # NOTE: Explicitly include hparam prefix prefix for d_model to431 # disambiguate with top-level head_dim432 # NOTE 2: If needed for future models, this can be isolated in a method433 # to separate the prefix setting and the keys used434 self.d_model = self.find_hparam([f"{self.hparam_prefixes[0]}_head_dim", "hidden_size", "d_model"])435 self.n_group = self.find_hparam(["n_groups", "num_groups"])436 self.d_inner = self.find_hparam(["expand", "num_heads"]) * self.d_model437 438 def get_attn_layers(self):439 # Explicit list of layer type names440 if layer_types := self.hparams.get("layer_types"):441 return [442 i for i, typ in enumerate(layer_types)443 if typ == "attention"444 ]445 446 # Layer types indicated by index or period447 attn_layers = self.hparams.get("attn_layer_indices", [])448 if not attn_layers:449 attn_period = self.hparams.get("attn_layer_period")450 assert attn_period, "Didn't find attn_layer_indices or attn_layer_period"451 attn_offset = self.hparams.get("attn_layer_offset")452 assert attn_offset is not None, "No attention layer offset set with attn_layer_period"453 attn_layers = [454 i for i in range(self.block_count)455 if i % attn_period == attn_offset456 ]457 return attn_layers458 459 def find_hparam(self, keys: Iterable[str], *args, **kwargs) -> Any:460 prefixed = []461 for pfx in self.hparam_prefixes:462 prefixed.extend(463 "_".join([pfx, k])464 for k in keys465 )466 keys = list(keys) + prefixed467 return Mamba2Model.find_hparam(self, keys, *args, **kwargs)468 469 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:470 if (471 name.endswith("block_sparse_moe.input_linear.weight")472 or "shared_mlp" in name473 ):474 yield from GraniteMoeModel.modify_tensors(self, data_torch, name, bid)475 return476 477 # Determine whether this is a mamba layer or an attention layer478 if bid in self._ssm_layers:479 yield from Mamba2Model.modify_tensors(self, data_torch, name, bid)480 return481 elif bid in self._attn_layers:482 yield from GraniteMoeModel.modify_tensors(self, data_torch, name, bid)483 return484 yield from ModelBase.modify_tensors(self, data_torch, name, bid)485 486 def set_gguf_parameters(self):487 """This method merges params from both parents and some that are488 specific to this model. The result is some duplication of how the params489 get set. The following warnings are expected during conversion:490 491 WARNING:Duplicated key name 'granitehybrid.attention.head_count_kv'492 WARNING:Duplicated key name 'granitehybrid.context_length'493 """494 GraniteMoeModel.set_gguf_parameters(self)495 496 ## Mamba mixer params ##497 self.gguf_writer.add_ssm_conv_kernel(self.find_hparam(["conv_kernel", "d_conv"]))498 self.gguf_writer.add_ssm_state_size(self.find_hparam(["state_size", "d_state", "state_dim", "ssm_state_size"]))499 self.gguf_writer.add_ssm_group_count(self.n_group)500 self.gguf_writer.add_ssm_inner_size(self.d_inner)501 # NOTE: The mamba_dt_rank is _not_ the right field for how this is used502 # in llama.cpp503 self.gguf_writer.add_ssm_time_step_rank(self.find_hparam(["n_heads", "num_heads"]))504 505 ## Attention params ##506 head_count_kv = self.find_hparam(["num_key_value_heads", "n_head_kv"])507 head_count_kv_vec = [508 head_count_kv if i in self._attn_layers else 0 for i in range(self.block_count)509 ]510 if rope_dim := self.hparams.get("attn_rotary_emb"):511 self.gguf_writer.add_rope_dimension_count(rope_dim)512 self.gguf_writer.add_head_count_kv(head_count_kv_vec)513 514 ## If Bamba or non-hybrid, use rope, otherwise don't515 use_rope = (516 "BambaForCausalLM" in self.hparams["architectures"]517 or not self._ssm_layers518 )519 self.gguf_writer.add_rope_scaling_finetuned(use_rope)520 if not use_rope:521 self.gguf_writer.add_context_length(2**20)522 523 ## Validation ##524 d_head = self.find_hparam(["d_head"], optional=True) or 64525 assert self.hparams.get("hidden_act") in [None, "silu"], "Only SILU activation supported"526 assert self.d_inner % d_head == 0, f"SSM inner size {self.d_inner} not a multiple of head dim {d_head}"527 528 def set_vocab(self):529 # For models with no ssm layers, don't pad for mamba2530 self.hparams["pad_vocab_size_multiple"] = 8 if self._ssm_layers else 1531 Mamba2Model.set_vocab(self)532 533 534@ModelBase.register("GraniteSpeechForConditionalGeneration")535@ModelBase.example("ibm-granite/granite-speech-3.3-2b", "ibm-granite/granite-4.0-1b-speech")536class GraniteSpeechMmprojModel(MmprojModel):537 has_vision_encoder = False538 has_audio_encoder = True539 540 _batch_norm_tensors: list[dict[str, Tensor]] | None = None541 542 def get_audio_config(self) -> dict[str, Any] | None:543 return self.global_config.get("encoder_config")544 545 def set_gguf_parameters(self):546 assert self.hparams_audio is not None547 a = self.hparams_audio548 a["hidden_size"] = a["hidden_dim"]549 a["intermediate_size"] = a["hidden_dim"] * a["feedforward_mult"]550 a["num_attention_heads"] = a["num_heads"]551 a["num_hidden_layers"] = a["num_layers"]552 553 super().set_gguf_parameters()554 555 self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GRANITE_SPEECH)556 self.gguf_writer.add_audio_num_mel_bins(a["input_dim"])557 self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)558 self.gguf_writer.add_audio_chunk_size(a["context_size"])559 self.gguf_writer.add_audio_conv_kernel_size(a["conv_kernel_size"])560 self.gguf_writer.add_audio_max_pos_emb(a["max_pos_emb"])561 562 p = self.global_config563 self.gguf_writer.add_audio_projector_window_size(p["window_size"])564 self.gguf_writer.add_audio_projector_downsample_rate(p["downsample_rate"])565 self.gguf_writer.add_audio_projector_head_count(p["projector_config"]["num_attention_heads"])566 567 def tensor_force_quant(self, name, new_name, bid, n_dims):568 if "encoder" in name or "projector" in name:569 if ".conv" in name and ".weight" in name:570 return gguf.GGMLQuantizationType.F32571 return super().tensor_force_quant(name, new_name, bid, n_dims)572 573 @classmethod574 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:575 name, gen = item576 if "attention_dists" in name or "num_batches_tracked" in name:577 return None578 return super().filter_tensors(item)579 580 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:581 # fold running_mean, running_var and eps into weight and bias for batch_norm582 if "batch_norm" in name and "encoder.layers." in name:583 if self._batch_norm_tensors is None:584 self._batch_norm_tensors = [{} for _ in range(self.block_count)]585 assert bid is not None586 self._batch_norm_tensors[bid][name] = data_torch587 if len(self._batch_norm_tensors[bid]) < 4:588 return589 prefix = f"encoder.layers.{bid}.conv.batch_norm"590 weight = self._batch_norm_tensors[bid][f"{prefix}.weight"]591 bias = self._batch_norm_tensors[bid][f"{prefix}.bias"]592 running_mean = self._batch_norm_tensors[bid][f"{prefix}.running_mean"]593 running_var = self._batch_norm_tensors[bid][f"{prefix}.running_var"]594 eps = 1e-5595 a = weight / torch.sqrt(running_var + eps)596 b = bias - running_mean * a597 yield from super().modify_tensors(a, f"encoder.layers.{bid}.conv.batch_norm.weight", bid)598 yield from super().modify_tensors(b, f"encoder.layers.{bid}.conv.batch_norm.bias", bid)599 return600 601 if ".attn.to_kv.weight" in name:602 k_weight, v_weight = data_torch.chunk(2, dim=0)603 yield from super().modify_tensors(k_weight, name.replace("to_kv", "to_k"), bid)604 yield from super().modify_tensors(v_weight, name.replace("to_kv", "to_v"), bid)605 return606 607 if ("up_conv" in name or "down_conv" in name) and name.endswith(".weight"):608 if data_torch.ndim == 3 and data_torch.shape[2] == 1:609 data_torch = data_torch.squeeze(2)610 611 if "depth_conv" in name and name.endswith(".weight"):612 if data_torch.ndim == 3 and data_torch.shape[1] == 1:613 data_torch = data_torch.squeeze(1)614 615 yield from super().modify_tensors(data_torch, name, bid)616 617 618@ModelBase.register("GraniteSpeechPlusForConditionalGeneration")619@ModelBase.example("ibm-granite/granite-speech-4.1-2b-plus")620class GraniteSpeechPlusMmprojModel(GraniteSpeechMmprojModel):621 """Conversion for GraniteSpeechPlus - extends GraniteSpeech with feature layer concatenation"""622 has_vision_encoder = False623 has_audio_encoder = True624 625 def set_gguf_parameters(self):626 assert self.hparams_audio is not None627 super().set_gguf_parameters()628 629 # Add feature_layer if present in encoder config630 if feature_layers := self.hparams_audio.get("cat_hidden_layers"):631 self.gguf_writer.add_audio_feature_layers(feature_layers)632 logger.info(f"gguf: audio feature_layers = {feature_layers}")633 634 # Validate projector dimension matches concatenated encoder output635 hidden_dim = self.hparams_audio["hidden_dim"]636 expected_dim = hidden_dim * (len(feature_layers) + 1)637 projector_dim = self.global_config["projector_config"]["encoder_hidden_size"]638 639 if projector_dim != expected_dim:640 raise ValueError(641 f"Projector encoder_hidden_size ({projector_dim}) does not match "642 f"expected concatenated dimension ({expected_dim}). "643 f"Expected: hidden_dim ({hidden_dim}) * (len(feature_layers) + 1) = {expected_dim}"644 )645 646 647@ModelBase.register("Granite4VisionForConditionalGeneration")648@ModelBase.example("ibm-granite/granite-4.0-3b-vision")649class Granite4VisionMmprojModel(MmprojModel):650 has_vision_encoder = True651 has_audio_encoder = False652 653 @staticmethod654 def get_normalized_projector_map(global_config: dict) -> list[tuple[int, int, str, int]]:655 """Normalize both deepstack and spatial projector maps to the form:656 (vision_layer, llm_layer, <type>, type_index)657 658 This is then used to populate the following mappings:659 - vision_feature_layers (mmproj hparam): ordered list of all660 vision_layer values where order corresponds with the order of the661 stacked projector tensors662 NOTE: Values may appear multiple times for spatial projectors663 - tensor_prefix_map (mmproj tensors): mapping from tensor prefixes to664 the index of the corresponding projector in the stacked tensors665 - deepstack_layer_arr (llm hparam): per-text-layer array indicating666 which input vision feature should be injected at that layer667 (-1 if none)668 669 Output: (vision_layer, llm_layer, <type>, type_index)670 """671 deepstack_map = global_config.get("deepstack_layer_map", []) # [[vis_layer, llm_layer], ...]672 spatial_layers = global_config.get("spatial_target_layers", []) # [llm_layer, ...]673 n_text_layers = global_config["text_config"]["num_hidden_layers"]674 n_vision_layers = global_config["vision_config"]["num_hidden_layers"]675 normalized_projector_map = []676 if deepstack_map:677 for deepstack_idx, (vision_layer, llm_layer) in enumerate(sorted(deepstack_map)):678 if vision_layer < 0:679 vision_layer = n_vision_layers + vision_layer680 if llm_layer < 0:681 llm_layer = n_text_layers + llm_layer682 normalized_projector_map.append((vision_layer, llm_layer, "layerwise", deepstack_idx))683 if spatial_layers:684 spatial_vision_layer = global_config.get("spatial_vision_layer", -1)685 if spatial_vision_layer < 0:686 spatial_vision_layer = n_vision_layers + spatial_vision_layer687 for spatial_idx, llm_layer in enumerate(spatial_layers):688 normalized_projector_map.append((spatial_vision_layer, llm_layer, "spatial", spatial_idx))689 return list(sorted(normalized_projector_map, key=(lambda entry: entry[1])))690 691 def __init__(self, *args, **kwargs):692 super().__init__(*args, **kwargs)693 normalized_projector_map = self.get_normalized_projector_map(self.global_config)694 self._n_proj = len(normalized_projector_map)695 696 self._tensor_prefix_map = {697 f"model.{proj_type}_projectors.{type_idx}": proj_idx698 for proj_idx, (_, _, proj_type, type_idx) in enumerate(normalized_projector_map)699 }700 self._vision_feature_layers = [vision_layer for vision_layer, _, _, _ in normalized_projector_map]701 self._spatial_offsets = [702 type_idx if proj_type == "spatial" else -1703 for _, _, proj_type, type_idx in normalized_projector_map704 ]705 706 def set_gguf_parameters(self):707 assert self.hparams_vision is not None708 super().set_gguf_parameters()709 710 self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GRANITE4_VISION)711 712 # SigLIP encoder hparams713 self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams.get("layer_norm_eps", 1e-6))714 self.gguf_writer.add_vision_use_gelu(True)715 716 # Preprocessor717 self.gguf_writer.add_vision_preproc_image_size(self.hparams.get("image_size", 384))718 719 # QFormer projector config720 ds_rate = self.global_config["downsample_rate"]721 ds_parts = ds_rate.split("/")722 assert len(ds_parts) == 2, f"Invalid 'downsample_rate' value: {ds_rate}"723 query_side, window_side = [int(p) for p in ds_parts]724 self.gguf_writer.add_vision_projector_query_side(query_side)725 self.gguf_writer.add_vision_projector_window_side(window_side)726 727 # Set vision feature layers728 self.gguf_writer.add_vision_feature_layers(self._vision_feature_layers)729 730 # Set the spatial offests per projector731 self.gguf_writer.add_vision_spatial_offsets(self._spatial_offsets)732 733 # Add flattened image grind pinpoints (resolution candidates internally)734 if pinpoints := self.global_config.get("image_grid_pinpoints"):735 # Flatten with h, w -> w, h inversion736 pinpoints = [val for h, w in pinpoints for val in (w, h)]737 self.gguf_writer.add_vision_image_grid_pinpoints(pinpoints)738 739 @classmethod740 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:741 name, _ = item742 if ("vision_model.head" in name or name.startswith("lm_head")):743 return None744 return super().filter_tensors(item)745 746 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:747 748 # Detect projector tensors and bin them749 projector_idx = None750 for prefix, proj_idx in self._tensor_prefix_map.items():751 if name.startswith(prefix):752 projector_idx = proj_idx753 break754 if projector_idx is not None:755 # If this projector tensor has a block id within the projector,756 # alias the bid to projector_idx757 #758 # TODO: currently, none of the Granite 4 Vision models have759 # projectors with multiple QFormer layers, so the `layer.{}` index760 # is always 0. This allows us to simply map to a single `bid` that761 # matches the projector index. If this changes, we'll need a762 # convention that merges the two IDs.763 id_matches = list(re.finditer(r"\.([0-9]+)\.", name))764 all_ids = [int(m.group(1)) for m in id_matches]765 assert len(all_ids) >= 1 and len(all_ids) <= 2, "Must have at least 1 and at most 2 ids in tensor names"766 # If not layer id, just use the projector index767 new_bid = projector_idx768 if len(all_ids) == 1:769 new_name = name[:id_matches[0].span(1)[0]] + str(new_bid) + name[id_matches[0].span(1)[1]:]770 else: # len(all_ids) == 2771 new_bid = projector_idx # + all_ids[1]772 new_name = name[:id_matches[0].span(0)[0]] + name[id_matches[0].span(1)[1]:id_matches[1].span(1)[0]] + str(new_bid) + name[id_matches[1].span(1)[1]:]773 yield from super().modify_tensors(data_torch, new_name, new_bid)774 return775 yield from super().modify_tensors(data_torch, name, bid)776