llama.cpp k2-fast2: MBZUAI model/K2Horizon + FAST2 sliding-window override

c59fd93 verified about 1 month ago

8.02 kB

from __future__ import annotations

import re
from pathlib import Path
from typing import Iterable

import torch
from torch import Tensor

from .base import ModelBase, TextModel, gguf

@ModelBase.register(
    "K2HorizonForCausalLM",
    "K2AuroraForCausalLM", # TODO: DELETE
)
@ModelBase.example(
    "IFM/K2-Horizon-0.9B",
    "IFM/K2-Horizon-36B",
)
class K2HorizonModel(TextModel):
    model_arch = gguf.MODEL_ARCH.K2HORIZON

    def set_vocab(self):
        super().set_vocab()

        template_path = (
            Path(__file__).parent.parent
            / "models"
            / "templates"
            / "k2-horizon.jinja"
        )
        template = template_path.read_text(encoding="utf-8")
        self.gguf_writer.remove_key(gguf.Keys.Tokenizer.CHAT_TEMPLATE)
        self.gguf_writer.add_chat_template(template)

    def set_gguf_parameters(self):
        super().set_gguf_parameters()
        # generic
        rope_head_dim = self.hparams.get("rope_head_dim")
        norm_groups = int(self.hparams.get("layernorm_num_groups", 1))

        self.gguf_writer.add_group_norm_groups(norm_groups)
        if rope_head_dim is not None:
            self.gguf_writer.add_rope_dimension_count(int(rope_head_dim))

        # moe
        num_experts = int(self.hparams.get("num_experts", 0))
        if num_experts > 0:
            moe_ff = int(self.hparams["moe_intermediate_size"])
            dense_layers = self.hparams.get("num_dense_layers")
            mlp_only_layers = {int(layer) for layer in self.hparams.get("mlp_only_layers", [])}
            sparse_step = int(self.hparams.get("decoder_sparse_step", 1))
            shared_experts = int(self.hparams.get("num_shared_experts", 0))
            router_scale = self.hparams.get("router_scaling_factor")
            normalize_topk = bool(self.hparams.get("norm_topk_prob", False))
            router_func = self.hparams.get("router_score_func")

            if dense_layers is None:
                dense_layers = 0
                while dense_layers in mlp_only_layers:
                    dense_layers += 1

            self.gguf_writer.add_expert_feed_forward_length(moe_ff)
            self.gguf_writer.add_leading_dense_block_count(dense_layers)            
            self.gguf_writer.add_moe_every_n_layers(sparse_step)
            self.gguf_writer.add_expert_shared_count(shared_experts)
            self.gguf_writer.add_expert_weights_norm(normalize_topk)
            if shared_experts > 0:
                self.gguf_writer.add_expert_shared_feed_forward_length(moe_ff * shared_experts)
            if router_scale is not None:
                self.gguf_writer.add_expert_weights_scale(float(router_scale))
            match router_func:
                case "sigmoid":
                    gating_func = gguf.ExpertGatingFuncType.SIGMOID    
                case "softmax":
                    gating_func = gguf.ExpertGatingFuncType.SOFTMAX
                case _:
                    raise ValueError(f"Unsupported router_score_func: {router_func!r}")
            self.gguf_writer.add_expert_gating_func(gating_func)

        # mova
        value_experts = int(self.hparams.get("mova_num_experts", 0))
        value_experts_used = int(self.hparams.get("mova_num_experts_per_tok", 0))

        if value_experts > 0 and value_experts_used > 0:
            assert value_experts_used <= value_experts
            self.gguf_writer.add_attention_value_expert_count(value_experts)
            self.gguf_writer.add_attention_value_expert_used_count(value_experts_used)

        # gate func, only making sure it exists and is softplus
        gate_func = self.hparams.get("attention_gate_func")
        if gate_func not in (None, "softplus"):
            raise ValueError(f"Unsupported attention_gate_func: {gate_func!r}")

    _experts: list[dict[str, Tensor]] | None = None
    _value_experts: list[dict[str, Tensor]] | None = None
    def modify_tensors(
        self,
        data_torch: Tensor,
        name: str,
        bid: int | None
    ) -> Iterable[tuple[str, Tensor]]:
        # MoE: router
        if name.endswith(".mlp.gate.bias"):
            assert bid is not None
            yield (
                self.format_tensor_name(
                    gguf.MODEL_TENSOR.FFN_EXP_PROBS_B,
                    bid,
                    ".bias"
                ),
                data_torch
            )
            return

        # MoE: actual up down or gate
        is_moe_tensor = re.fullmatch(r"model\.layers\.\d+\.mlp\.experts\.\d+\.(down_proj|gate_proj|up_proj)\.weight", name)
        if is_moe_tensor:
            assert bid is not None
            num_experts = int(self.hparams["num_experts"])

            # allocate on first layer that has experts
            if self._experts is None:
                self._experts = [{} for _ in range(self.block_count)]

            # atp, this_blocks_experts contains all experts
            this_blocks_experts = self._experts[bid]
            this_blocks_experts[name] = data_torch

            # filling up self._experts until up down gate are all inside, then continue
            if len(this_blocks_experts) < num_experts * 3:
                return

            for projection in ("down_proj", "gate_proj", "up_proj"):
                tensors = []
                for expert_id in range(num_experts):
                    expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{projection}.weight"
                    tensors.append(this_blocks_experts.pop(expert_name))
                merged = torch.stack(tensors, dim=0)
                merged_name = f"model.layers.{bid}.mlp.experts.{projection}.weight"
                yield from super().modify_tensors(
                    merged,
                    merged_name,
                    bid
                )
            return
        
        # MoVA
        is_mova_weights = re.fullmatch(r"model\.layers\.\d+\.self_attn\.v_experts\.\d+\.weight", name)
        if is_mova_weights:
            assert bid is not None
            num_value_experts = int(self.hparams["mova_num_experts"])
            if self._value_experts is None:
                self._value_experts = [{} for _ in range(self.block_count)]

            this_blocks_value_expert = self._value_experts[bid]
            this_blocks_value_expert[name] = data_torch

            # no need to * 3 because no up down gate like normal moe
            if len(this_blocks_value_expert) < num_value_experts:
                return

            tensors = []
            for value_exp_id in range(num_value_experts):
                value_exp_name = f"model.layers.{bid}.self_attn.v_experts.{value_exp_id}.weight"
                tensors.append(this_blocks_value_expert.pop(value_exp_name))

            merged = torch.stack(tensors, dim = 0)
            merged_name = f"model.layers.{bid}.self_attn.v_experts.weight"
            yield from super().modify_tensors(
                merged,
                merged_name,
                bid
            )
            return

        # fallback, the default way basically
        yield from super().modify_tensors(
            data_torch,
            name,
            bid
        )

    def prepare_tensors(self):
        super().prepare_tensors()

        # this is just checks basically
        if self._experts is not None:
            remaining_experts = [
                name
                for block in self._experts
                for name in block
            ]

            if remaining_experts:
                raise ValueError(
                    f"Unprocessed MoE experts: {remaining_experts}"
                )

        if self._value_experts is not None:
            remaining_value_experts = [
                name
                for block in self._value_experts
                for name in block
            ]

            if remaining_value_experts:
                raise ValueError(
                    "Unprocessed MoVA value experts: "
                    f"{remaining_value_experts}"
                )

	from __future__ import annotations

	import re
	from pathlib import Path
	from typing import Iterable

	import torch
	from torch import Tensor

	from .base import ModelBase, TextModel, gguf

	@ModelBase.register(
	"K2HorizonForCausalLM",
	"K2AuroraForCausalLM", # TODO: DELETE
	)
	@ModelBase.example(
	"IFM/K2-Horizon-0.9B",
	"IFM/K2-Horizon-36B",
	)
	class K2HorizonModel(TextModel):
	model_arch = gguf.MODEL_ARCH.K2HORIZON

	def set_vocab(self):
	super().set_vocab()

	template_path = (
	Path(__file__).parent.parent
	/ "models"
	/ "templates"
	/ "k2-horizon.jinja"
	)
	template = template_path.read_text(encoding="utf-8")
	self.gguf_writer.remove_key(gguf.Keys.Tokenizer.CHAT_TEMPLATE)
	self.gguf_writer.add_chat_template(template)

	def set_gguf_parameters(self):
	super().set_gguf_parameters()
	# generic
	rope_head_dim = self.hparams.get("rope_head_dim")
	norm_groups = int(self.hparams.get("layernorm_num_groups", 1))

	self.gguf_writer.add_group_norm_groups(norm_groups)
	if rope_head_dim is not None:
	self.gguf_writer.add_rope_dimension_count(int(rope_head_dim))

	# moe
	num_experts = int(self.hparams.get("num_experts", 0))
	if num_experts > 0:
	moe_ff = int(self.hparams["moe_intermediate_size"])
	dense_layers = self.hparams.get("num_dense_layers")
	mlp_only_layers = {int(layer) for layer in self.hparams.get("mlp_only_layers", [])}
	sparse_step = int(self.hparams.get("decoder_sparse_step", 1))
	shared_experts = int(self.hparams.get("num_shared_experts", 0))
	router_scale = self.hparams.get("router_scaling_factor")
	normalize_topk = bool(self.hparams.get("norm_topk_prob", False))
	router_func = self.hparams.get("router_score_func")

	if dense_layers is None:
	dense_layers = 0
	while dense_layers in mlp_only_layers:
	dense_layers += 1

	self.gguf_writer.add_expert_feed_forward_length(moe_ff)
	self.gguf_writer.add_leading_dense_block_count(dense_layers)
	self.gguf_writer.add_moe_every_n_layers(sparse_step)
	self.gguf_writer.add_expert_shared_count(shared_experts)
	self.gguf_writer.add_expert_weights_norm(normalize_topk)
	if shared_experts > 0:
	self.gguf_writer.add_expert_shared_feed_forward_length(moe_ff * shared_experts)
	if router_scale is not None:
	self.gguf_writer.add_expert_weights_scale(float(router_scale))
	match router_func:
	case "sigmoid":
	gating_func = gguf.ExpertGatingFuncType.SIGMOID
	case "softmax":
	gating_func = gguf.ExpertGatingFuncType.SOFTMAX
	case _:
	raise ValueError(f"Unsupported router_score_func: {router_func!r}")
	self.gguf_writer.add_expert_gating_func(gating_func)

	# mova
	value_experts = int(self.hparams.get("mova_num_experts", 0))
	value_experts_used = int(self.hparams.get("mova_num_experts_per_tok", 0))

	if value_experts > 0 and value_experts_used > 0:
	assert value_experts_used <= value_experts
	self.gguf_writer.add_attention_value_expert_count(value_experts)
	self.gguf_writer.add_attention_value_expert_used_count(value_experts_used)

	# gate func, only making sure it exists and is softplus
	gate_func = self.hparams.get("attention_gate_func")
	if gate_func not in (None, "softplus"):
	raise ValueError(f"Unsupported attention_gate_func: {gate_func!r}")

	_experts: list[dict[str, Tensor]] \| None = None
	_value_experts: list[dict[str, Tensor]] \| None = None
	def modify_tensors(
	self,
	data_torch: Tensor,
	name: str,
	bid: int \| None
	) -> Iterable[tuple[str, Tensor]]:
	# MoE: router
	if name.endswith(".mlp.gate.bias"):
	assert bid is not None
	yield (
	self.format_tensor_name(
	gguf.MODEL_TENSOR.FFN_EXP_PROBS_B,
	bid,
	".bias"
	),
	data_torch
	)
	return

	# MoE: actual up down or gate
	is_moe_tensor = re.fullmatch(r"model\.layers\.\d+\.mlp\.experts\.\d+\.(down_proj\|gate_proj\|up_proj)\.weight", name)
	if is_moe_tensor:
	assert bid is not None
	num_experts = int(self.hparams["num_experts"])

	# allocate on first layer that has experts
	if self._experts is None:
	self._experts = [{} for _ in range(self.block_count)]

	# atp, this_blocks_experts contains all experts
	this_blocks_experts = self._experts[bid]
	this_blocks_experts[name] = data_torch

	# filling up self._experts until up down gate are all inside, then continue
	if len(this_blocks_experts) < num_experts * 3:
	return

	for projection in ("down_proj", "gate_proj", "up_proj"):
	tensors = []
	for expert_id in range(num_experts):
	expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{projection}.weight"
	tensors.append(this_blocks_experts.pop(expert_name))
	merged = torch.stack(tensors, dim=0)
	merged_name = f"model.layers.{bid}.mlp.experts.{projection}.weight"
	yield from super().modify_tensors(
	merged,
	merged_name,
	bid
	)
	return

	# MoVA
	is_mova_weights = re.fullmatch(r"model\.layers\.\d+\.self_attn\.v_experts\.\d+\.weight", name)
	if is_mova_weights:
	assert bid is not None
	num_value_experts = int(self.hparams["mova_num_experts"])
	if self._value_experts is None:
	self._value_experts = [{} for _ in range(self.block_count)]

	this_blocks_value_expert = self._value_experts[bid]
	this_blocks_value_expert[name] = data_torch

	# no need to * 3 because no up down gate like normal moe
	if len(this_blocks_value_expert) < num_value_experts:
	return

	tensors = []
	for value_exp_id in range(num_value_experts):
	value_exp_name = f"model.layers.{bid}.self_attn.v_experts.{value_exp_id}.weight"
	tensors.append(this_blocks_value_expert.pop(value_exp_name))

	merged = torch.stack(tensors, dim = 0)
	merged_name = f"model.layers.{bid}.self_attn.v_experts.weight"
	yield from super().modify_tensors(
	merged,
	merged_name,
	bid
	)
	return

	# fallback, the default way basically
	yield from super().modify_tensors(
	data_torch,
	name,
	bid
	)

	def prepare_tensors(self):
	super().prepare_tensors()

	# this is just checks basically
	if self._experts is not None:
	remaining_experts = [
	name
	for block in self._experts
	for name in block
	]

	if remaining_experts:
	raise ValueError(
	f"Unprocessed MoE experts: {remaining_experts}"
	)

	if self._value_experts is not None:
	remaining_value_experts = [
	name
	for block in self._value_experts
	for name in block
	]

	if remaining_value_experts:
	raise ValueError(
	"Unprocessed MoVA value experts: "
	f"{remaining_value_experts}"
	)