Felipe97/llama-cpp-compiled
01.1k
1from __future__ import annotations2 3from typing import Callable, Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8 from torch import Tensor9 10from .base import ModelBase, TextModel, gguf11 12 13@ModelBase.register("Rwkv6ForCausalLM")14@ModelBase.example("RWKV/v6-Finch-1B6-HF")15class Rwkv6Model(TextModel):16 model_arch = gguf.MODEL_ARCH.RWKV617 18 def set_vocab(self):19 self._set_vocab_rwkv_world()20 21 def set_gguf_parameters(self):22 head_size = self.hparams["head_size"]23 hidden_size = self.hparams["hidden_size"]24 layer_norm_eps = self.hparams["layer_norm_epsilon"]25 rescale_every_n_layers = self.hparams["rescale_every"]26 intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else int((hidden_size * 3.5) // 32 * 32)27 time_mix_extra_dim = 64 if hidden_size == 4096 else 3228 time_decay_extra_dim = 128 if hidden_size == 4096 else 6429 30 # RWKV isn't context limited31 self.gguf_writer.add_context_length(1048576)32 self.gguf_writer.add_embedding_length(hidden_size)33 self.gguf_writer.add_block_count(self.block_count)34 self.gguf_writer.add_layer_norm_eps(layer_norm_eps)35 self.gguf_writer.add_rescale_every_n_layers(rescale_every_n_layers)36 self.gguf_writer.add_wkv_head_size(head_size)37 self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim)38 self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim)39 self.gguf_writer.add_feed_forward_length(intermediate_size)40 self.gguf_writer.add_file_type(self.ftype)41 42 # required by llama.cpp, unused43 self.gguf_writer.add_head_count(0)44 45 lerp_weights: dict[int, dict[str, Tensor]] = {}46 47 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:48 new_name = self.map_tensor_name(name)49 50 if not (new_name.endswith(".weight") or new_name.endswith(".bias")):51 new_name += ".weight"52 53 if new_name.endswith("time_mix_w1.weight") or new_name.endswith("time_mix_decay_w1.weight") or new_name.endswith("time_mix_decay_w2.weight"):54 data_torch = data_torch.transpose(0, 1)55 56 if new_name.endswith("time_mix_w2.weight"):57 data_torch = data_torch.permute(0, 2, 1)58 59 if new_name.endswith("time_mix_decay.weight") or "lerp" in new_name:60 data_torch = data_torch.squeeze()61 62 try:63 rescale_every_n_layers = self.hparams["rescale_every"]64 if rescale_every_n_layers > 0:65 if new_name.endswith("time_mix_output.weight") or new_name.endswith("channel_mix_value.weight"):66 data_torch = data_torch.div_(2 ** int(bid // rescale_every_n_layers))67 except KeyError:68 pass69 70 # concat time_mix_lerp weights to reduce some cpu overhead71 # also reduces the number of tensors in the model72 if bid is not None and "time_mix_lerp" in new_name and "time_mix_lerp_x" not in new_name:73 try:74 self.lerp_weights[bid][new_name] = data_torch75 except KeyError:76 self.lerp_weights[bid] = {new_name: data_torch}77 if all(f"blk.{bid}.time_mix_lerp_{i}.weight" in self.lerp_weights[bid].keys() for i in ["w", "k", "v", "r", "g"]):78 new_name = f"blk.{bid}.time_mix_lerp_fused.weight"79 data = torch.stack([self.lerp_weights[bid][f"blk.{bid}.time_mix_lerp_{i}.weight"].unsqueeze(0) for i in ["w", "k", "v", "r", "g"]], dim=0).unsqueeze(1)80 yield (new_name, data)81 return82 83 yield (new_name, data_torch)84 85 86@ModelBase.register("RWKV6Qwen2ForCausalLM")87@ModelBase.example("recursal/QRWKV6-32B-Instruct-Preview-v0.1")88class RWKV6Qwen2Model(Rwkv6Model):89 model_arch = gguf.MODEL_ARCH.RWKV6QWEN290 91 def set_vocab(self):92 try:93 self._set_vocab_sentencepiece()94 except FileNotFoundError:95 self._set_vocab_gpt2()96 97 def set_gguf_parameters(self):98 num_attention_heads = self.hparams["num_attention_heads"]99 num_key_value_heads = self.hparams["num_key_value_heads"]100 hidden_size = self.hparams["hidden_size"]101 head_size = hidden_size // num_attention_heads102 rms_norm_eps = self.hparams["rms_norm_eps"]103 intermediate_size = self.hparams["intermediate_size"]104 time_mix_extra_dim = self.hparams.get("lora_rank_tokenshift", 64 if hidden_size >= 4096 else 32)105 time_decay_extra_dim = self.hparams.get("lora_rank_decay", 128 if hidden_size >= 4096 else 64)106 107 # RWKV isn't context limited108 self.gguf_writer.add_context_length(1048576)109 self.gguf_writer.add_embedding_length(hidden_size)110 self.gguf_writer.add_block_count(self.block_count)111 self.gguf_writer.add_wkv_head_size(head_size)112 self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim)113 self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim)114 self.gguf_writer.add_feed_forward_length(intermediate_size)115 self.gguf_writer.add_file_type(self.ftype)116 117 # special parameters for time_mixing in RWKV6QWEN2118 self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)119 self.gguf_writer.add_token_shift_count(1)120 # RWKV6QWEN2 use grouped key/value like GQA121 self.gguf_writer.add_head_count_kv(num_key_value_heads)122 123 # required by llama.cpp, unused124 self.gguf_writer.add_head_count(0)125 126 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:127 for new_name, data in super().modify_tensors(data_torch, name, bid):128 if "time_mix_w1" in new_name or "time_mix_w2" in new_name:129 data = data.view(5, -1, data.shape[-1])130 # rwkv6qwen2 has a different order of rkvwg instead of the original wkvrg131 # permute them here to avoid code changes132 data = torch.stack([data[3], data[1], data[2], data[0], data[4]], dim=0).view(-1, data.shape[-1])133 if "w2" in new_name:134 data = data.view(5, -1, data.shape[-1])135 yield (new_name, data)136 continue137 yield (new_name, data)138 139 140@ModelBase.register("Rwkv7ForCausalLM", "RWKV7ForCausalLM")141@ModelBase.example("fla-hub/rwkv7-1.5B-world")142class Rwkv7Model(TextModel):143 model_arch = gguf.MODEL_ARCH.RWKV7144 145 def set_vocab(self):146 self._set_vocab_rwkv_world()147 148 def calc_lora_rank(self, hidden_size, exponent, multiplier):149 return max(1, round(hidden_size ** exponent * multiplier / 32)) * 32150 151 def set_gguf_parameters(self):152 try:153 head_size = self.hparams["head_size"]154 layer_norm_eps = self.hparams["layer_norm_epsilon"]155 except KeyError:156 head_size = self.hparams["head_dim"]157 layer_norm_eps = self.hparams["norm_eps"]158 hidden_size = self.hparams["hidden_size"]159 intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else (hidden_size * 4)160 161 # ICLR: In-Context-Learning-Rate162 try:163 lora_rank_decay = self.hparams["lora_rank_decay"] if self.hparams["lora_rank_decay"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)164 lora_rank_iclr = self.hparams["lora_rank_iclr"] if self.hparams["lora_rank_iclr"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)165 lora_rank_value_residual_mix = self.hparams["lora_rank_value_residual_mix"] if self.hparams["lora_rank_value_residual_mix"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3)166 lora_rank_gate = self.hparams["lora_rank_gate"] if self.hparams["lora_rank_gate"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6)167 except KeyError:168 lora_rank_decay = self.hparams["decay_low_rank_dim"] if self.hparams["decay_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)169 lora_rank_iclr = self.hparams["a_low_rank_dim"] if self.hparams["a_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8)170 lora_rank_value_residual_mix = self.hparams["v_low_rank_dim"] if self.hparams["v_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3)171 lora_rank_gate = self.hparams["gate_low_rank_dim"] if self.hparams["gate_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6)172 173 # RWKV isn't context limited174 self.gguf_writer.add_context_length(1048576)175 self.gguf_writer.add_embedding_length(hidden_size)176 self.gguf_writer.add_block_count(self.block_count)177 self.gguf_writer.add_layer_norm_eps(layer_norm_eps)178 self.gguf_writer.add_wkv_head_size(head_size)179 self.gguf_writer.add_decay_lora_rank(lora_rank_decay)180 self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr)181 self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix)182 self.gguf_writer.add_gate_lora_rank(lora_rank_gate)183 self.gguf_writer.add_feed_forward_length(intermediate_size)184 self.gguf_writer.add_file_type(self.ftype)185 186 # required by llama.cpp, unused187 self.gguf_writer.add_head_count(0)188 189 lerp_weights: dict[int, dict[str, Tensor]] = {}190 lora_needs_transpose: bool = True191 192 @classmethod193 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:194 name, gen = item195 196 # unify tensor names here to make life easier197 name = name.replace("blocks", "layers").replace("ffn", "feed_forward")198 name = name.replace("self_attn", "attention").replace("attn", "attention")199 name = name.replace("time_mixer.", "")200 201 name = name.replace("feed_forward_norm", "ln2")202 name = name.replace("g_norm", "ln_x")203 204 return super().filter_tensors((name, gen))205 206 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:207 # lora layer names in fla-hub's impl208 if "_lora.lora" in name:209 self.lora_needs_transpose = False210 name = name.replace("_lora.lora.0.weight", "1.weight")211 name = name.replace("_lora.lora.2.weight", "2.weight")212 name = name.replace("_lora.lora.2.bias", "0.weight")213 214 if "attention.v" in name and "value" not in self.map_tensor_name(name) and bid == 0:215 # some models have dummy v0/v1/v2 on first layer while others don't216 # ignore them all since they are not used217 return218 219 wkv_has_gate = self.hparams.get("wkv_has_gate", True)220 lerp_list = ["r", "w", "k", "v", "a", "g"] if wkv_has_gate else ["r", "w", "k", "v", "a"]221 222 if bid is not None and "attention.x_" in name:223 if "attention.x_x" in name:224 # already concatenated225 new_name = f"blk.{bid}.time_mix_lerp_fused.weight"226 data = data_torch.reshape(len(lerp_list), 1, 1, -1)227 yield (new_name, data)228 else:229 try:230 self.lerp_weights[bid][name] = data_torch231 except KeyError:232 self.lerp_weights[bid] = {name: data_torch}233 if all(f"model.layers.{bid}.attention.x_{i}" in self.lerp_weights[bid].keys() for i in lerp_list):234 new_name = f"blk.{bid}.time_mix_lerp_fused.weight"235 data = torch.stack([self.lerp_weights[bid][f"model.layers.{bid}.attention.x_{i}"] for i in lerp_list], dim=0)236 yield (new_name, data)237 return238 else:239 data_torch = data_torch.squeeze()240 new_name = self.map_tensor_name(name)241 242 if not (new_name.endswith(".weight") or new_name.endswith(".bias")):243 new_name += ".weight"244 245 if self.lora_needs_transpose and any(246 new_name.endswith(t) for t in [247 "time_mix_w1.weight", "time_mix_w2.weight",248 "time_mix_a1.weight", "time_mix_a2.weight",249 "time_mix_v1.weight", "time_mix_v2.weight",250 "time_mix_g1.weight", "time_mix_g2.weight",251 ]252 ):253 data_torch = data_torch.transpose(0, 1)254 255 if 'r_k' in new_name:256 data_torch = data_torch.flatten()257 258 if bid == 0 and "time_mix_a" in new_name:259 # dummy v0/v1/v2 on first layer260 # easiest way to make llama happy261 yield (new_name.replace("time_mix_a", "time_mix_v"), data_torch)262 263 yield (new_name, data_torch)264 265 266@ModelBase.register("RwkvHybridForCausalLM")267@ModelBase.example("RWKV-Red-Team/ARWKV-7B-Preview-0.1")268class ARwkv7Model(Rwkv7Model):269 model_arch = gguf.MODEL_ARCH.ARWKV7270 271 def set_vocab(self):272 try:273 self._set_vocab_sentencepiece()274 except FileNotFoundError:275 self._set_vocab_gpt2()276 277 def set_gguf_parameters(self):278 hidden_size = self.hparams["hidden_size"]279 head_size = self.hparams["head_size"]280 rms_norm_eps = self.hparams["rms_norm_eps"]281 intermediate_size = self.hparams["intermediate_size"]282 wkv_has_gate = self.hparams["wkv_has_gate"]283 assert self.hparams["wkv_version"] == 7284 285 # ICLR: In-Context-Learning-Rate286 lora_rank_decay = 64287 lora_rank_iclr = 64288 lora_rank_value_residual_mix = 32289 lora_rank_gate = 128 if wkv_has_gate else 0290 291 # RWKV isn't context limited292 self.gguf_writer.add_context_length(1048576)293 self.gguf_writer.add_embedding_length(hidden_size)294 self.gguf_writer.add_block_count(self.block_count)295 self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps)296 self.gguf_writer.add_wkv_head_size(head_size)297 self.gguf_writer.add_decay_lora_rank(lora_rank_decay)298 self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr)299 self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix)300 self.gguf_writer.add_gate_lora_rank(lora_rank_gate)301 self.gguf_writer.add_feed_forward_length(intermediate_size)302 self.gguf_writer.add_file_type(self.ftype)303 self.gguf_writer.add_token_shift_count(1)304 305 # required by llama.cpp, unused306 self.gguf_writer.add_head_count(0)307