Felipe97/llama-cpp-compiled
01.1k
1from __future__ import annotations2 3import math4 5from pathlib import Path6from typing import Callable, Iterable, TYPE_CHECKING7 8import torch9 10if TYPE_CHECKING:11 from torch import Tensor12 13from .base import MmprojModel, ModelBase, TextModel, gguf14from .qwenvl import Qwen2VLVisionModel15 16 17@ModelBase.register("ExaoneForCausalLM")18@ModelBase.example("LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct")19class ExaoneModel(TextModel):20 model_arch = gguf.MODEL_ARCH.EXAONE21 22 def set_gguf_parameters(self):23 super().set_gguf_parameters()24 hparams = self.hparams25 26 assert (hparams["activation_function"] == "silu")27 28 rotary_factor = self.rope_parameters.get("partial_rotary_factor")29 rotary_factor = rotary_factor if rotary_factor is not None else 1.030 self.gguf_writer.add_rope_dimension_count(int(rotary_factor * (hparams["hidden_size"] // hparams["num_attention_heads"])))31 32 def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:33 if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):34 if rope_params.get("rope_type", '').lower() == "llama3":35 base = self.rope_parameters.get("rope_theta", 10000.0)36 if (dim := self.hparams.get("head_dim")) is None:37 dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]38 freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))39 40 factor = rope_params.get("factor", 8.0)41 low_freq_factor = rope_params.get("low_freq_factor", 1.0)42 high_freq_factor = rope_params.get("high_freq_factor", 4.0)43 old_context_len = rope_params.get("original_max_position_embeddings", 8192)44 45 low_freq_wavelen = old_context_len / low_freq_factor46 high_freq_wavelen = old_context_len / high_freq_factor47 assert low_freq_wavelen != high_freq_wavelen48 49 rope_factors = []50 for freq in freqs:51 wavelen = 2 * math.pi / freq52 if wavelen < high_freq_wavelen:53 rope_factors.append(1)54 elif wavelen > low_freq_wavelen:55 rope_factors.append(factor)56 else:57 smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)58 rope_factors.append(1 / ((1 - smooth) / factor + smooth))59 60 yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))61 62 63@ModelBase.register("Exaone4ForCausalLM")64@ModelBase.example("LGAI-EXAONE/EXAONE-4.0-32B")65class Exaone4Model(TextModel):66 model_arch = gguf.MODEL_ARCH.EXAONE467 68 def set_vocab(self):69 tokens, toktypes, tokpre = self.get_vocab_base()70 self.gguf_writer.add_tokenizer_model("gpt2")71 self.gguf_writer.add_tokenizer_pre(tokpre)72 self.gguf_writer.add_token_list(tokens)73 self.gguf_writer.add_token_types(toktypes)74 75 special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)76 special_vocab.add_to_gguf(self.gguf_writer)77 78 def set_gguf_parameters(self):79 super().set_gguf_parameters()80 hparams = self.hparams81 self.gguf_writer.add_vocab_size(hparams["vocab_size"])82 83 if hparams.get("sliding_window") is not None:84 self.gguf_writer.add_sliding_window(hparams["sliding_window"])85 if "layer_types" in hparams:86 self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hparams["layer_types"]])87 elif "sliding_window_pattern" in hparams:88 sliding_window_pattern = []89 if isinstance(hparams["sliding_window_pattern"], str): # e.g. LLLG90 for i in range(hparams["num_hidden_layers"]):91 sliding_window_pattern.append(hparams["sliding_window_pattern"][i % len(hparams["sliding_window_pattern"])] == "L")92 if isinstance(hparams["sliding_window_pattern"], int): # e.g. 493 for i in range(hparams["num_hidden_layers"]):94 sliding_window_pattern.append((i + 1) % hparams["sliding_window_pattern"] != 0)95 if len(sliding_window_pattern) == hparams["num_hidden_layers"]:96 self.gguf_writer.add_sliding_window_pattern(sliding_window_pattern)97 98 def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:99 if rope_params := self.rope_parameters.get("full_attention", self.rope_parameters):100 if rope_params.get("rope_type", '').lower() == "llama3":101 base = rope_params.get("rope_theta", 10_000.0)102 if (dim := self.hparams.get("head_dim")) is None:103 dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]104 freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))105 106 factor = rope_params.get("factor", 16.0)107 low_freq_factor = rope_params.get("low_freq_factor", 1.0)108 high_freq_factor = rope_params.get("high_freq_factor", 4.0)109 old_context_len = rope_params.get("original_max_position_embeddings", 8192)110 111 low_freq_wavelen = old_context_len / low_freq_factor112 high_freq_wavelen = old_context_len / high_freq_factor113 114 rope_factors = []115 for freq in freqs:116 wavelen = 2 * math.pi / freq117 if wavelen < high_freq_wavelen:118 rope_factors.append(1)119 elif wavelen > low_freq_wavelen:120 rope_factors.append(factor)121 else:122 smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)123 rope_factors.append(1 / ((1 - smooth) / factor + smooth))124 125 yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))126 127 128# note: transformers >= 5.1 renamed the class to "ExaoneMoeForCausalLM" (lowercase 'e'),129# so accept both spellings - LG AI have updated the configs of already-released models130@ModelBase.register("ExaoneMoEForCausalLM", "ExaoneMoeForCausalLM")131@ModelBase.example("LGAI-EXAONE/K-EXAONE-236B-A23B")132class ExaoneMoEModel(Exaone4Model):133 model_arch = gguf.MODEL_ARCH.EXAONE_MOE134 135 def __init__(self, *args, **kwargs):136 super().__init__(*args, **kwargs)137 self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0)138 self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)139 140 def set_gguf_parameters(self):141 super().set_gguf_parameters()142 moe_intermediate_size = self.hparams["moe_intermediate_size"]143 num_shared_experts = self.hparams["num_shared_experts"]144 self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)145 self.gguf_writer.add_expert_shared_count(num_shared_experts)146 self.gguf_writer.add_expert_shared_feed_forward_length(moe_intermediate_size * num_shared_experts)147 self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])148 self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])149 n_dense_layer = self.hparams.get("first_k_dense_replace", self.hparams.get("first_last_k_dense_replace", 0))150 self.gguf_writer.add_leading_dense_block_count(n_dense_layer)151 self.gguf_writer.add_nextn_predict_layers(self.hparams.get("num_nextn_predict_layers", 0))152 153 self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)154 155 _experts: list[dict[str, Tensor]] | None = None156 157 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:158 if name.startswith("mtp."):159 if name.find("layers.") != -1:160 # `mtp.layers.0.[module_name]` format161 name = name.replace(f"mtp.layers.{bid}", f"model.layers.{bid + self.hparams['num_hidden_layers']}")162 else:163 # mtp fc/norm weights164 remapper = {165 "mtp.fc": "model.layers.{bid}.eh_proj",166 "mtp.pre_fc_norm_embedding": "model.layers.{bid}.enorm",167 "mtp.pre_fc_norm_hidden": "model.layers.{bid}.hnorm",168 "mtp.norm": "model.layers.{bid}.shared_head.norm",169 }170 _n = Path(name)171 new_name = remapper[_n.stem] + _n.suffix172 173 # set shared weights for all NextN/MTP layers174 for bid in range(self.hparams['num_hidden_layers'], self.block_count):175 yield from super().modify_tensors(data_torch, new_name.format(bid=bid), bid)176 return177 178 if name.find("mlp.experts") != -1:179 n_experts = self.find_hparam(["num_local_experts", "num_experts"])180 assert bid is not None181 182 if self._experts is None:183 self._experts = [{} for _ in range(self.block_count)]184 185 self._experts[bid][name] = data_torch186 187 if len(self._experts[bid]) >= n_experts * 3:188 # merge the experts into a single 3d tensor189 for w_name in ["down_proj", "gate_proj", "up_proj"]:190 datas: list[Tensor] = []191 192 for xid in range(n_experts):193 ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"194 datas.append(self._experts[bid][ename])195 del self._experts[bid][ename]196 197 data_torch = torch.stack(datas, dim=0)198 199 merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"200 201 new_name = self.map_tensor_name(merged_name)202 203 yield from super().modify_tensors(data_torch, new_name, bid)204 return205 else:206 return207 208 yield from super().modify_tensors(data_torch, name, bid)209 210 def prepare_tensors(self):211 super().prepare_tensors()212 if self._experts is not None:213 # flatten `list[dict[str, Tensor]]` into `list[str]`214 experts = [k for d in self._experts for k in d.keys()]215 if len(experts) > 0:216 raise ValueError(f"Unprocessed experts: {experts}")217 218 219@ModelBase.register("Exaone4_5_ForConditionalGeneration")220@ModelBase.example("LGAI-EXAONE/EXAONE-4.5-33B")221class Exaone4_5_TextModel(Exaone4Model):222 """Text tower of EXAONE 4.5; Tensors match EXAONE4"""223 224 model_arch = gguf.MODEL_ARCH.EXAONE4225 226 def __init__(self, *args, **kwargs):227 super().__init__(*args, **kwargs)228 n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0)229 if n_nextn > 0:230 self.block_count = self.hparams["num_hidden_layers"] + n_nextn231 self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)232 233 def set_gguf_parameters(self):234 super().set_gguf_parameters()235 n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0)236 if n_nextn > 0:237 self.gguf_writer.add_nextn_predict_layers(n_nextn)238 239 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:240 if name.startswith("mtp."):241 n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0)242 if n_nextn <= 0:243 return244 nh = self.hparams["num_hidden_layers"]245 if ".layers." in name:246 share = self.hparams.get("mtp_share_layers", False)247 mtp_bid = bid if bid is not None else 0248 if share:249 for k in range(n_nextn):250 nn = name.replace(f"mtp.layers.{mtp_bid}", f"model.layers.{nh + k}")251 yield from super().modify_tensors(data_torch, nn, nh + k)252 return253 name = name.replace(f"mtp.layers.{mtp_bid}", f"model.layers.{mtp_bid + nh}")254 else:255 remapper = {256 "mtp.fc": gguf.MODEL_TENSOR.NEXTN_EH_PROJ,257 "mtp.pre_fc_norm_embedding": gguf.MODEL_TENSOR.NEXTN_ENORM,258 "mtp.pre_fc_norm_hidden": gguf.MODEL_TENSOR.NEXTN_HNORM,259 "mtp.norm": gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,260 }261 _n = Path(name)262 key = _n.stem263 if key not in remapper:264 return265 for bid_mtp in range(nh, self.block_count):266 mapped_name = self.format_tensor_name(remapper[key], bid_mtp, suffix=_n.suffix)267 yield from ModelBase.modify_tensors(self, data_torch, mapped_name, bid_mtp)268 return269 270 yield from super().modify_tensors(data_torch, name, bid)271 272 273@ModelBase.register("Exaone4_5_ForConditionalGeneration")274@ModelBase.example("LGAI-EXAONE/EXAONE-4.5-33B")275class Exaone4_5VisionModel(Qwen2VLVisionModel):276 """Vision tower for EXAONE 4.5; Qwen2-VL-style ViT (GQA) + patch merger"""277 278 @classmethod279 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:280 name, gen = item281 name = name.replace("model.visual.", "visual.", 1)282 return super().filter_tensors((name, gen))283 284 def set_gguf_parameters(self):285 MmprojModel.set_gguf_parameters(self)286 assert self.hparams_vision is not None287 hparams = self.hparams_vision288 self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.EXAONE4_5)289 self.gguf_writer.add_vision_use_silu(True)290 self.gguf_writer.add_vision_min_pixels(self.preprocessor_config["min_pixels"])291 self.gguf_writer.add_vision_max_pixels(self.preprocessor_config["max_pixels"])292 num_kv_head = self.find_vparam(["num_key_value_heads"], optional=True)293 if num_kv_head is not None:294 self.gguf_writer.add_vision_head_count_kv(num_kv_head)295 eps = hparams.get("rms_norm_eps", self.global_config.get("rms_norm_eps", 1e-6))296 self.gguf_writer.add_vision_attention_layernorm_eps(eps)297 if (window_size := hparams.get("window_size")) is not None:298 self.gguf_writer.add_vision_window_size(window_size)299 fullatt_block_indexes = hparams.get("fullatt_block_indexes")300 if fullatt_block_indexes:301 n_wa_pattern = fullatt_block_indexes[0] + 1302 for i in range(1, len(fullatt_block_indexes)):303 if fullatt_block_indexes[i] - fullatt_block_indexes[i - 1] != n_wa_pattern:304 raise ValueError(f"Invalid EXAONE4.5 fullatt_block_indexes: {fullatt_block_indexes}")305 self.gguf_writer.add_vision_n_wa_pattern(n_wa_pattern)306 307 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:308 if ".qkv." in name:309 yield from ModelBase.modify_tensors(self, data_torch, name, bid)310 return311 312 yield from Qwen2VLVisionModel.modify_tensors(self, data_torch, name, bid)313