Felipe97/llama-cpp-compiled
01.1k
1from __future__ import annotations2 3import json4import math5 6from typing import Callable, Iterable, TYPE_CHECKING7 8import numpy as np9import torch10 11if TYPE_CHECKING:12 from torch import Tensor13 14from .base import ModelBase, TextModel, gguf, logger15 16 17@ModelBase.register(18 "LLaMAForCausalLM",19 "LlamaForCausalLM",20 "MistralForCausalLM",21 "MixtralForCausalLM",22 "VLlama3ForCausalLM",23 "LlavaForConditionalGeneration",24 "VoxtralForConditionalGeneration",25 "LlamaForCausalLMEagle3",26 "Eagle3LlamaForCausalLM",27 "Eagle3Speculator",28 "Eagle3DraftModel",29 "IQuestCoderForCausalLM",30 "LlamaModel")31# [TAG_HF_EXAMPLE_GATED] meta-llama/Llama-3.2-1B-Instruct is gated32@ModelBase.example("unsloth/Llama-3.2-1B-Instruct", "mistralai/Mistral-7B-Instruct-v0.3", "mistralai/Mixtral-8x7B-Instruct-v0.1")33class LlamaModel(TextModel):34 model_arch = gguf.MODEL_ARCH.LLAMA35 undo_permute = True36 37 def __init__(self, *args, **kwargs):38 super().__init__(*args, **kwargs)39 # fix for SmolVLM2, missing `num_attention_heads` in config.json40 if self.hf_arch == "VLlama3ForCausalLM":41 self.hparams["num_attention_heads"] = self.hparams.get("num_attention_heads", 32)42 # Mistral consolidated format has no config.json; origin_hf_arch is HF-only.43 if self.is_mistral_format:44 self.origin_hf_arch = None45 else:46 hparams = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)47 self.origin_hf_arch = hparams.get('architectures', [None])[0]48 49 # Detect eagle3 draft checkpoint by hparams (some models don't use a distinct HF arch name)50 if "draft_vocab_size" in self.hparams and self.hparams["num_hidden_layers"] == 1:51 self.is_eagle3 = True52 self.model_arch = gguf.MODEL_ARCH.EAGLE353 logger.info("Detected EAGLE-3 draft model, switching to EAGLE3 architecture")54 # Re-initialize tensor_map with eagle3 architecture55 self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)56 # Update gguf_writer architecture57 self.gguf_writer.arch = gguf.MODEL_ARCH_NAMES[self.model_arch]58 self.gguf_writer.add_architecture()59 if self.target_model_dir is None:60 raise ValueError(61 "EAGLE-3 model requires --target-model-dir to be specified. "62 "Please provide the path to the target model directory to read config.json"63 )64 # Read both eagle3 raw config and target model config65 with open(self.dir_model / "config.json", 'r', encoding='utf-8') as f:66 eagle3_raw_config = json.load(f)67 with open(self.target_model_dir / "config.json", 'r', encoding='utf-8') as f:68 target_config = json.load(f)69 70 if "text_config" in target_config:71 target_config = {**target_config, **target_config["text_config"]}72 self.target_vocab_size = target_config["vocab_size"]73 74 # target_layers: use the eagle3 config's explicit aux hidden-state layer ids75 # if present, else derive from the target layer count.76 target_num_layers = target_config["num_hidden_layers"]77 aux_layer_ids = eagle3_raw_config.get("eagle_aux_hidden_state_layer_ids")78 if aux_layer_ids:79 target_layers = aux_layer_ids80 else:81 target_layers = [2, target_num_layers // 2, target_num_layers - 3]82 logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)")83 self.gguf_writer.add_target_layers(target_layers)84 85 # target_hidden_size: prefer eagle3 config, fallback to target config86 if eagle3_raw_config.get("target_hidden_size") is not None:87 target_hidden_size = eagle3_raw_config["target_hidden_size"]88 src = "EAGLE-3 config"89 else:90 target_hidden_size = target_config["hidden_size"]91 src = "target model config"92 logger.info(f"EAGLE-3: target_hidden_size = {target_hidden_size} (from {src})")93 self.gguf_writer.add_target_hidden_size(target_hidden_size)94 95 # norm_before_residual (RedHat-style eagle3 specific)96 norm_before_residual = eagle3_raw_config.get("norm_before_residual", False)97 logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}")98 self.gguf_writer.add_norm_before_residual(norm_before_residual)99 100 # norm_before_fc: RMSNorm applied to the fused target features before the101 # fc projection (e.g. nvidia/gpt-oss-120b-Eagle3-v3)102 norm_before_fc = eagle3_raw_config.get("norm_before_fc", False)103 logger.info(f"EAGLE-3: norm_before_fc = {norm_before_fc}")104 self.gguf_writer.add_norm_before_fc(norm_before_fc)105 106 def set_vocab(self):107 # eagle3: use tokenizer from target model if provided108 original_dir_model = None109 if getattr(self, 'is_eagle3', False):110 assert self.target_model_dir is not None111 logger.info(f"EAGLE-3: Using tokenizer from target model: {self.target_model_dir}")112 original_dir_model = self.dir_model113 self.dir_model = self.target_model_dir114 115 if self.origin_hf_arch == "GlmasrModel":116 return self._set_vocab_glmedge()117 118 if self.is_mistral_format:119 return self._set_vocab_mistral()120 121 path_tekken_json = self.dir_model / "tekken.json"122 path_tokenizer_json = self.dir_model / "tokenizer.json"123 if path_tekken_json.is_file() and not path_tokenizer_json.is_file():124 return self._set_vocab_mistral()125 126 tokenizer_config_file = self.dir_model / 'tokenizer_config.json'127 if tokenizer_config_file.is_file():128 with open(tokenizer_config_file, "r", encoding="utf-8") as f:129 tokenizer_config_json = json.load(f)130 if (add_prefix_space := tokenizer_config_json.get("add_prefix_space")) is not None:131 self.gguf_writer.add_add_space_prefix(add_prefix_space)132 if tokenizer_config_json.get("tokenizer_class") == "HybridDNATokenizer":133 return self._set_vocab_hybriddna()134 135 try:136 self._set_vocab_sentencepiece()137 except FileNotFoundError:138 try:139 self._set_vocab_llama_hf()140 except (FileNotFoundError, TypeError):141 # Llama 3142 self._set_vocab_gpt2()143 144 # Apply to CodeLlama only (and ignore for Llama 3 with a vocab size of 128256)145 if self.hparams.get("vocab_size", 32000) == 32016:146 special_vocab = gguf.SpecialVocab(147 self.dir_model, load_merges=False,148 special_token_types = ['prefix', 'suffix', 'middle', 'eot']149 )150 special_vocab._set_special_token("prefix", 32007)151 special_vocab._set_special_token("suffix", 32008)152 special_vocab._set_special_token("middle", 32009)153 special_vocab._set_special_token("eot", 32010)154 special_vocab.add_to_gguf(self.gguf_writer)155 156 # Apply to granite small models only157 if self.hparams.get("vocab_size", 32000) == 49152:158 self.gguf_writer.add_add_bos_token(False)159 160 # eagle3: Restore original dir_model161 if original_dir_model is not None:162 self.dir_model = original_dir_model163 164 def set_gguf_parameters(self):165 super().set_gguf_parameters()166 hparams = self.hparams167 168 if not self.is_mistral_format:169 self.gguf_writer.add_vocab_size(hparams["vocab_size"])170 171 if (rope_dim := hparams.get("head_dim")) is None:172 rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]173 self.gguf_writer.add_rope_dimension_count(rope_dim)174 175 @staticmethod176 def permute(weights: Tensor, n_head: int, n_head_kv: int | None):177 if n_head_kv is not None and n_head != n_head_kv:178 n_head = n_head_kv179 return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])180 .swapaxes(1, 2)181 .reshape(weights.shape))182 183 def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor, input_scale: Tensor):184 # Mirror the BF16 Q/K RoPE permutation site in modify_tensors; the NVFP4 path bypasses it.185 if self.undo_permute:186 n_head = self.find_hparam(["n_heads", "num_attention_heads"], optional=True)187 n_kv_head = self.find_hparam(["n_kv_heads", "num_key_value_heads"], optional=True)188 if n_head is not None:189 if name.endswith("q_proj.weight"):190 weight = LlamaModel.permute(weight, n_head, n_head)191 scale = LlamaModel.permute(scale, n_head, n_head)192 elif name.endswith("k_proj.weight"):193 weight = LlamaModel.permute(weight, n_head, n_kv_head)194 scale = LlamaModel.permute(scale, n_head, n_kv_head)195 super()._repack_nvfp4(name, weight, scale, scale2, input_scale)196 197 _experts: list[dict[str, Tensor]] | None = None198 199 @classmethod200 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:201 name, gen = item202 203 if "text_model." in name:204 name = name.replace("text_model.", "") # for SmolVLM205 206 return super().filter_tensors((name, gen))207 208 def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:209 tensors = super().index_tensors(remote_hf_model_id)210 211 # Handle Eagle3Speculator nested config212 if "transformer_layer_config" in self.hparams:213 self.hparams = {**self.hparams, **self.hparams["transformer_layer_config"]}214 215 # eagle3 detection216 if "draft_vocab_size" in self.hparams and self.hparams["num_hidden_layers"] == 1:217 logger.info("EAGLE-3: renaming midlayer.* / layers.0.* to model.layers.0.*")218 new_tensors = {}219 for name, gen in tensors.items():220 if name.startswith("midlayer."):221 new_name = "model.layers.0." + name[len("midlayer."):]222 new_tensors[new_name] = gen223 elif name.startswith("layers.0."): # Eagle3Speculator format224 new_name = "model." + name225 new_tensors[new_name] = gen226 else:227 new_tensors[name] = gen228 return new_tensors229 230 return tensors231 232 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:233 # eagle3: special tensors that bypass standard llama mapping234 if getattr(self, 'is_eagle3', False):235 if name == "fc.weight":236 yield (name, data_torch)237 return238 if name == "input_norm.weight":239 yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch)240 return241 if name == "d2t":242 # store for manual int64 handling in prepare_tensors (avoid F32 conversion)243 if not hasattr(self, '_eagle3_int_tensors'):244 self._eagle3_int_tensors = {}245 self._eagle3_int_tensors[name] = data_torch246 return247 if name == "t2d":248 # not used at runtime, skip249 return250 if name.endswith(".hidden_norm.weight"):251 yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_NORM_2, bid), data_torch)252 return253 254 n_head = self.find_hparam(["n_heads", "num_attention_heads"])255 n_kv_head = self.find_hparam(["n_kv_heads", "num_key_value_heads"])256 257 if self.hf_arch == "LlamaModel":258 name = "model." + name259 260 if self.undo_permute:261 if name.endswith(("q_proj.weight", "q_proj.bias")):262 data_torch = LlamaModel.permute(data_torch, n_head, n_head)263 if name.endswith(("k_proj.weight", "k_proj.bias")):264 data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)265 266 # process the experts separately267 if name.find("block_sparse_moe.experts") != -1:268 n_experts = self.hparams["num_local_experts"]269 270 assert bid is not None271 272 if self._experts is None:273 self._experts = [{} for _ in range(self.block_count)]274 275 self._experts[bid][name] = data_torch276 277 if len(self._experts[bid]) >= n_experts * 3:278 # merge the experts into a single 3d tensor279 for wid in ["w1", "w2", "w3"]:280 datas: list[Tensor] = []281 282 for xid in range(n_experts):283 ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"284 datas.append(self._experts[bid][ename])285 del self._experts[bid][ename]286 287 data_torch = torch.stack(datas, dim=0)288 289 merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"290 291 yield from super().modify_tensors(data_torch, merged_name, bid)292 return293 else:294 return295 296 yield from super().modify_tensors(data_torch, name, bid)297 298 def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:299 if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):300 if rope_params.get("rope_type", '').lower() == "llama3":301 base = rope_params.get("rope_theta", 10000.0)302 if (dim := self.hparams.get("head_dim")) is None:303 dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]304 freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))305 306 factor = rope_params.get("factor", 8.0)307 low_freq_factor = rope_params.get("low_freq_factor", 1.0)308 high_freq_factor = rope_params.get("high_freq_factor", 4.0)309 old_context_len = rope_params.get("original_max_position_embeddings", 8192)310 311 low_freq_wavelen = old_context_len / low_freq_factor312 high_freq_wavelen = old_context_len / high_freq_factor313 # assert low_freq_wavelen != high_freq_wavelen # Errors for Llama4314 315 rope_factors = []316 for freq in freqs:317 wavelen = 2 * math.pi / freq318 if wavelen < high_freq_wavelen:319 rope_factors.append(1)320 elif wavelen > low_freq_wavelen:321 rope_factors.append(factor)322 else:323 smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)324 rope_factors.append(1 / ((1 - smooth) / factor + smooth))325 326 yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))327 328 def prepare_tensors(self):329 # eagle3: collect d2t original dtype before parent converts tensors to F32330 eagle3_original_dtypes = {}331 if getattr(self, 'is_eagle3', False):332 for name, data_torch in self.get_tensors():333 if name == "d2t":334 eagle3_original_dtypes[name] = data_torch.dtype335 336 super().prepare_tensors()337 338 # eagle3: write d2t as absolute target token ids339 if getattr(self, 'is_eagle3', False) and hasattr(self, '_eagle3_int_tensors'):340 for name, data_torch in self._eagle3_int_tensors.items():341 old_dtype = eagle3_original_dtypes.get(name, data_torch.dtype)342 data = data_torch.to(torch.int64).cpu().numpy()343 if name == "d2t":344 data = data.reshape(-1)345 data = data + np.arange(data.size, dtype=np.int64)346 if np.any((data < 0) | (data >= self.target_vocab_size)):347 raise ValueError(f"EAGLE-3 d2t target ids out of range for target vocab size {self.target_vocab_size}")348 if np.unique(data).size != data.size:349 raise ValueError("EAGLE-3 d2t contains duplicate target ids")350 data_qtype = gguf.GGMLQuantizationType.I64351 352 shape_str = f"{{{', '.join(str(n) for n in reversed(data.shape))}}}"353 logger.info(f"{name + ',':<30} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")354 self.gguf_writer.add_tensor(name, data, raw_dtype=data_qtype)355 356 if self._experts is not None:357 # flatten `list[dict[str, Tensor]]` into `list[str]`358 experts = [k for d in self._experts for k in d.keys()]359 if len(experts) > 0:360 raise ValueError(f"Unprocessed experts: {experts}")361 362 363@ModelBase.register("ArceeForCausalLM")364@ModelBase.example("arcee-ai/AFM-4.5B")365class ArceeModel(LlamaModel):366 model_arch = gguf.MODEL_ARCH.ARCEE367 368 def set_gguf_parameters(self):369 super().set_gguf_parameters()370 self._try_set_pooling_type()371 372 373@ModelBase.register(374 "Llama4ForConditionalGeneration",375 "Llama4ForCausalLM",376)377# [TAG_HF_EXAMPLE_GATED] meta-llama/Llama-4-Scout-17B-16E-Instruct is gated378@ModelBase.example("unsloth/Llama-4-Scout-17B-16E-Instruct")379class Llama4Model(LlamaModel):380 model_arch = gguf.MODEL_ARCH.LLAMA4381 undo_permute = False382 383 def __init__(self, *args, **kwargs):384 super().__init__(*args, **kwargs)385 # IMPORTANT: the normal "intermediate_size" is renamed to "intermediate_size_mlp", we need to undo this386 self.hparams["intermediate_size_moe"] = self.hparams["intermediate_size"]387 self.hparams["intermediate_size"] = self.hparams["intermediate_size_mlp"]388 389 def set_vocab(self):390 self._set_vocab_gpt2()391 392 def set_gguf_parameters(self):393 super().set_gguf_parameters()394 self.gguf_writer.add_interleave_moe_layer_step(self.hparams["interleave_moe_layer_step"])395 self.gguf_writer.add_expert_feed_forward_length(self.hparams["intermediate_size_moe"])396 if "layer_types" in self.hparams:397 if all(lt == "full_attention" for lt in self.hparams["layer_types"]):398 # all layers are full attention (for MobileLLM), disable swa399 self.gguf_writer.add_sliding_window(0)400 401 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):402 # split the gate_up into gate and up403 if "gate_up_proj" in name:404 name_up = name.replace("gate_up_proj", "up_proj.weight")405 name_gate = name.replace("gate_up_proj", "gate_proj.weight")406 dim_half = data_torch.shape[-1] // 2407 gate_proj_weight, up_proj_weight = data_torch.transpose(-1, -2).split(dim_half, dim=-2)408 yield from super().modify_tensors(gate_proj_weight, name_gate, bid)409 yield from super().modify_tensors(up_proj_weight, name_up, bid)410 return411 412 if name.endswith("down_proj"):413 name += ".weight"414 data_torch = data_torch.transpose(-1, -2)415 416 yield from super().modify_tensors(data_torch, name, bid)417 418 419@ModelBase.register("LlamaBidirectionalModel")420@ModelBase.example("nvidia/llama-embed-nemotron-8b")421class LlamaEmbedNemotronModel(LlamaModel):422 model_arch = gguf.MODEL_ARCH.LLAMA_EMBED423 424 425@ModelBase.register("SmolLM3ForCausalLM")426@ModelBase.example("HuggingFaceTB/SmolLM3-3B")427class SmolLM3Model(LlamaModel):428 model_arch = gguf.MODEL_ARCH.SMOLLM3429 430 431@ModelBase.register("ApertusForCausalLM")432@ModelBase.example("swiss-ai/Apertus-8B-Instruct-2509")433class ApertusModel(LlamaModel):434 model_arch = gguf.MODEL_ARCH.APERTUS435 undo_permute = False436 437 _alpha_n = {}438 _alpha_p = {}439 _beta = {}440 _eps = {}441 442 def modify_tensors(self, data_torch, name, bid):443 # Handle xIELU activation parameters444 n_layers = self.hparams["num_hidden_layers"]445 if name.endswith(".act_fn.alpha_n"):446 self._alpha_n[bid] = data_torch.to("cpu").float().item()447 if (len(self._alpha_n) == n_layers):448 self.gguf_writer.add_xielu_alpha_n([self._alpha_n[k] for k in sorted(self._alpha_n)])449 return450 if name.endswith(".act_fn.alpha_p"):451 self._alpha_p[bid] = data_torch.to("cpu").float().item()452 if (len(self._alpha_p) == n_layers):453 self.gguf_writer.add_xielu_alpha_p([self._alpha_p[k] for k in sorted(self._alpha_p)])454 return455 if name.endswith(".act_fn.beta"):456 self._beta[bid] = data_torch.to("cpu").float().item()457 if (len(self._beta) == n_layers):458 self.gguf_writer.add_xielu_beta([self._beta[k] for k in sorted(self._beta)])459 return460 if name.endswith(".act_fn.eps"):461 self._eps[bid] = data_torch.to("cpu").float().item()462 if (len(self._eps) == n_layers):463 self.gguf_writer.add_xielu_eps([self._eps[k] for k in sorted(self._eps)])464 return465 466 yield from super().modify_tensors(data_torch, name, bid)467