39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from .base import ModelBase, gguf
|
|
from .qwen import Qwen2MoeModel
|
|
|
|
|
|
@ModelBase.register("FlintForCausalLM")
|
|
@ModelBase.example("RecursiveSI/flint")
|
|
class FlintModel(Qwen2MoeModel):
|
|
model_arch = gguf.MODEL_ARCH.FLINT
|
|
|
|
def set_gguf_parameters(self):
|
|
super().set_gguf_parameters()
|
|
dense = self.hparams["mlp_only_layers"]
|
|
if dense != list(range(len(dense))):
|
|
raise ValueError("Flint dense layers must precede the MoE layers")
|
|
self.gguf_writer.add_leading_dense_block_count(len(dense))
|
|
self.gguf_writer.add_expert_shared_count(
|
|
self.hparams["shared_expert_intermediate_size"] // self.hparams["moe_intermediate_size"]
|
|
)
|
|
self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
|
|
self.gguf_writer.add_expert_weights_scale(1.0)
|
|
self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])
|
|
self.gguf_writer.add_sliding_window_pattern([
|
|
kind == "sliding_attention" for kind in self.hparams["layer_types"]
|
|
])
|
|
self.gguf_writer.add_rope_pattern([
|
|
i not in self.hparams["no_rope_layer_indices"] for i in range(self.block_count)
|
|
])
|
|
|
|
def modify_tensors(self, data_torch, name, bid):
|
|
if name.endswith(".softmax_offset"):
|
|
name += ".weight"
|
|
if name.endswith(".gate.expert_bias"):
|
|
if not self.hparams["router_expert_bias"]:
|
|
return
|
|
name += ".bias"
|
|
yield from super().modify_tensors(data_torch, name, bid)
|