Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import math4import re5 6from typing import Callable, Iterable, TYPE_CHECKING7 8import torch9 10if TYPE_CHECKING:11    from torch import Tensor12 13from .base import MmprojModel, ModelBase, TextModel, _MISTRAL_COMMON_DATASET_MEAN, _MISTRAL_COMMON_DATASET_STD, gguf14 15from .qwen import Qwen3Model16 17 18@ModelBase.register("StepVLForConditionalGeneration", "Step3p7ForConditionalGeneration")19@ModelBase.example("stepfun-ai/Step3-VL-10B", "stepfun-ai/Step-3.7-Flash")20class Step3VLVisionModel(MmprojModel):21    def __init__(self, *args, **kwargs):22        super().__init__(*args, **kwargs)23        assert self.hparams_vision is not None24 25        if not self.hparams_vision.get("intermediate_size"):26            hidden_size = self.hparams_vision.get("hidden_size") or self.hparams_vision.get("width") or 027            assert hidden_size > 028            mlp_ratio = float(self.hparams_vision.get("mlp_ratio", 8960 / 1536))29            self.hparams_vision["intermediate_size"] = int(round(hidden_size * mlp_ratio))30 31        self.preprocessor_config.setdefault("image_mean", list(_MISTRAL_COMMON_DATASET_MEAN))32        self.preprocessor_config.setdefault("image_std", list(_MISTRAL_COMMON_DATASET_STD))33 34    def set_gguf_parameters(self):35        super().set_gguf_parameters()36        assert self.hparams_vision is not None37 38        projector_stride = int(self.global_config.get("understand_projector_stride", -1))39        hidden_size = int(self.hparams_vision.get("hidden_size", self.hparams_vision.get("width", -1)))40        num_layers = int(self.hparams_vision.get("num_hidden_layers", self.hparams_vision.get("layers", -1)))41        assert (projector_stride, int(self.hparams_vision.get("image_size", -1)), hidden_size, num_layers) == (2, 728, 1536, 47), (42            "current Step3-VL conversion path is only validated for Step3-VL-10B"43        )44 45        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.STEP3VL)46        self.gguf_writer.add_vision_attention_layernorm_eps(float(self.hparams_vision.get("layer_norm_eps", 1e-5)))47        self.gguf_writer.add_vision_projector_scale_factor(projector_stride ** 2)48        # 3024 max resize comes from step3-vl-10b processing_step3.py.49        self.gguf_writer.add_vision_preproc_image_size(3024)50 51    def tensor_force_quant(self, name, new_name, bid, n_dims):52        if ".position_embd." in new_name:53            return gguf.GGMLQuantizationType.F3254        if ("mm.0." in new_name or "mm.1." in new_name) and new_name.endswith(".weight"):55            return gguf.GGMLQuantizationType.F16 if self.ftype == gguf.LlamaFileType.MOSTLY_F16 else gguf.GGMLQuantizationType.F3256        return super().tensor_force_quant(name, new_name, bid, n_dims)57 58    @classmethod59    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:60        name, gen = item61 62        if name.startswith(("model.", "lm_head.")):63            return None64 65        return super().filter_tensors(item)66 67    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:68        if name.startswith("vision_model.vit_downsampler"):69            match = re.match(r"vision_model\.vit_downsampler(\d+)\.(weight|bias)", name)70            if match is None:71                raise ValueError(f"Unexpected Step3-VL projector tensor {name!r}")72 73            proj_id = int(match.group(1)) - 174            suffix = f".{match.group(2)}"75            yield (self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ, proj_id, suffix=suffix), data_torch)76            return77 78        if name == "vit_large_projector.weight":79            yield (self.format_tensor_name(gguf.MODEL_TENSOR.V_MMPROJ_FC), data_torch)80            return81 82        if name.startswith("vision_model."):83            if name == "vision_model.positional_embedding":84                name += ".weight"85            elif name.endswith(".gamma") and ".ls_" in name:86                name = name.removesuffix(".gamma") + ".weight"87 88            name = name.replace("attn.in_proj_weight", "attn.in_proj.weight")89            name = name.replace("attn.in_proj_bias", "attn.in_proj.bias")90 91            yield from super().modify_tensors(data_torch, name, bid)92 93 94@ModelBase.register("StepVLForConditionalGeneration")95@ModelBase.example("stepfun-ai/Step3-VL-10B")96class Step3VLTextModel(Qwen3Model):97    model_arch = gguf.MODEL_ARCH.QWEN398 99 100@ModelBase.register("Step3p5ForCausalLM", "Step3p7ForConditionalGeneration")101@ModelBase.example("stepfun-ai/Step-3.7-Flash")102class Step35Model(TextModel):103    model_arch = gguf.MODEL_ARCH.STEP35104    supports_mtp_export = True105 106    # The --mtp / --no-mtp toggles are ModelBase.mtp_only / no_mtp (set in107    # convert_hf_to_gguf.py main()). Unlike Qwen3.5, which stores MTP under a108    # `mtp.*` namespace, Step3.5 appends MTP layers at109    # `model.layers.{num_hidden_layers + i}`, so we filter them by layer index.110    # The trunk layer count is captured before indexing so the classmethod111    # filter_tensors can tell the appended MTP block(s) apart from the trunk.112    _n_main_layers: int | None = None113 114    def __init__(self, *args, **kwargs):115        super().__init__(*args, **kwargs)116        # NextN/MTP layers are appended past num_hidden_layers; extend the117        # tensor map to cover them so the MTP block's tensors get correctly118        # indexed names. When --no-mtp drops the MTP blocks, fall back to the119        # base num_hidden_layers so we don't reserve unused slots.120        n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0))121        if n_nextn > 0 and not self.no_mtp:122            self.block_count += n_nextn123            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)124 125    def index_tensors(self, remote_hf_model_id: str | None = None):126        # filter_tensors is a classmethod and can't reach self.hparams; stash127        # the trunk layer count here (before indexing runs) so it can detect128        # the appended MTP layers by index.129        hparams = {**self.hparams, **self.hparams.get("text_config", {})}130        key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)131        type(self)._n_main_layers = hparams.get(key)132        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)133 134    def set_gguf_parameters(self):135        rope_theta = self.hparams.get("rope_theta")136        if isinstance(rope_theta, list):137            self.hparams["rope_theta"] = float(rope_theta[0])138            self.hparams["local_rope_theta"] = float(rope_theta[1])139            self.rope_parameters["rope_theta"] = self.hparams["rope_theta"]140            self.rope_parameters["sliding_attention"] = {"rope_theta": self.hparams["local_rope_theta"]}141 142        super().set_gguf_parameters()143 144        layer_types = self.hparams.get("layer_types") or []145        partial_rotary_factors = self.hparams.get("partial_rotary_factors") or []146        attn_other = self.hparams.get("attention_other_setting") or {}147 148        n_head_base = self.hparams["num_attention_heads"]149        n_kv_base = self.hparams["num_attention_groups"]150 151        n_head_swa = attn_other.get("num_attention_heads", n_head_base)152        n_kv_swa = attn_other.get("num_attention_groups", n_kv_base)153 154        n_nextn = int(self.hparams.get("num_nextn_predict_layers", 0))155 156        # The Step3p5 HF checkpoint stores layer_types/partial_rotary_factors157        # entries for the MTP blocks past num_hidden_layers; preserve them so158        # the MTP layer's attention shape, SWA flag, and partial RoPE dim are159        # set correctly. Pad with full-attention defaults if the checkpoint160        # truncated them.161        def _pad(arr, n, default):162            arr = list(arr)163            if len(arr) < n:164                arr = arr + [default] * (n - len(arr))165            return arr[:n]166 167        layer_types = _pad(layer_types, self.block_count, "full_attention")168        partial_rotary_factors = _pad(169            partial_rotary_factors,170            self.block_count,171            0.5,  # full_attention default for Step3p5172        )173        assert [1.0 if lt == "sliding_attention" else 0.5 for lt in layer_types] == partial_rotary_factors174        head_arr = [n_head_swa if lt == "sliding_attention" else n_head_base for lt in layer_types]175        kv_arr = [n_kv_swa if lt == "sliding_attention" else n_kv_base for lt in layer_types]176        swa_pat = [lt == "sliding_attention" for lt in layer_types]177 178        self.gguf_writer.add_head_count(head_arr)179        self.gguf_writer.add_head_count_kv(kv_arr)180 181        self.gguf_writer.add_sliding_window(self.hparams["sliding_window"])182        self.gguf_writer.add_sliding_window_pattern(swa_pat)183 184        self.gguf_writer.add_value_length(self.hparams["head_dim"])185 186        # MoE params187        self.gguf_writer.add_expert_count(self.hparams["moe_num_experts"])188        self.gguf_writer.add_expert_used_count(self.hparams["moe_top_k"])189        self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])190        self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["share_expert_dim"])191 192        if (moe_router_scaling_factor := self.hparams.get("moe_router_scaling_factor")) is not None:193            self.gguf_writer.add_expert_weights_scale(moe_router_scaling_factor)194        if (norm_expert_weight := self.hparams.get("norm_expert_weight")) is not None:195            self.gguf_writer.add_expert_weights_norm(norm_expert_weight)196 197        # leading dense blocks198        leading_dense = 0199        moe_layers_enum = self.hparams.get("moe_layers_enum")200        if isinstance(moe_layers_enum, str) and moe_layers_enum.strip():201            moe_layers = sorted(int(i) for i in moe_layers_enum.strip().split(","))202            if moe_layers:203                leading_dense = max(0, moe_layers[0])204        self.gguf_writer.add_leading_dense_block_count(leading_dense)205        self.gguf_writer.add_moe_every_n_layers(int(self.hparams.get("moe_every_n_layer", 1)))206 207        self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("rms_norm_eps", 1e-5))208 209        # Optional per-layer SwiGLU clamps. MTP layers default to no clamping (0.0).210        if (limits := self.hparams.get("swiglu_limits")) is not None:211            limits_f = _pad(212                [0.0 if v is None else float(v) for v in limits],213                self.block_count,214                0.0,215            )216            self.gguf_writer.add_swiglu_clamp_exp(limits_f)217        if (limits_shared := self.hparams.get("swiglu_limits_shared")) is not None:218            limits_shared_f = _pad(219                [0.0 if v is None else float(v) for v in limits_shared],220                self.block_count,221                0.0,222            )223            self.gguf_writer.add_swiglu_clamp_shexp(limits_shared_f)224 225        if n_nextn > 0 and not self.no_mtp:226            self.gguf_writer.add_nextn_predict_layers(n_nextn)227 228    @classmethod229    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:230        if (titem := super().filter_tensors(item)) is None:231            return None232        name, gen = titem233 234        # Map router bias (expert selection bias) to a GGUF bias tensor235        if name.endswith(".moe.router_bias"):236            name += ".bias"237 238        # Step3.5 appends the MTP block(s) past num_hidden_layers.239        assert cls._n_main_layers is not None240        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers241 242        # --no-mtp: drop the appended MTP block(s) entirely.243        if is_mtp and cls.no_mtp:244            return None245        # --mtp: keep ONLY MTP-block tensors plus the shared embeddings/norm/246        # lm_head (so the resulting GGUF carries just the draft head).247        if cls.mtp_only and not is_mtp and name not in (248            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",249        ):250            return None251 252        # The checkpoint nests the per-MTP-layer shared head under253        # `model.layers.{N+i}.transformer.shared_head.{norm,output}.weight`;254        # strip the `transformer.` infix and rename `output` → `head` so the255        # existing NEXTN_SHARED_HEAD_{NORM,HEAD} tensor mapping picks them up.256        # Mirrors vllm's `_rewrite_spec_layer_name` (step3p5_mtp.py).257        if is_mtp:258            name = name.replace(".transformer.", ".")259            name = name.replace("shared_head.output", "shared_head.head")260 261        return name, gen262 263    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):264        if name.endswith("norm.weight"):265            data_torch += 1.0266 267        if name.endswith((".self_attn.g_proj.weight", ".moe.gate.weight", ".moe.up_proj.weight", ".moe.gate_proj.weight", ".moe.down_proj.weight")):268            data_torch = data_torch.squeeze().contiguous()269 270        yield from super().modify_tensors(data_torch, name, bid)271 272    def prepare_metadata(self, vocab_only: bool):273        from_dir = self.fname_out.is_dir()274        super().prepare_metadata(vocab_only=vocab_only)275 276        # Mirror Qwen3.5's behavior: when emitting a draft-only file into a277        # directory, prefix with "mtp-" so it doesn't collide with the trunk.278        if not self.mtp_only or not from_dir:279            return280 281        output_type: str = self.ftype.name.partition("_")[2]282        fname_default: str = gguf.naming_convention(283            self.metadata.name, self.metadata.basename, self.metadata.finetune,284            self.metadata.version, size_label=None, output_type=output_type, model_type=None)285        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"286 287    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:288        # Step35 can optionally use Llama-3 style RoPE scaling (HF: rope_scaling.rope_type == "llama3").289        # llama.cpp represents this via a single extra tensor: "rope_freqs.weight" (aka MODEL_TENSOR.ROPE_FREQS).290        rope_params = self.rope_parameters.get("full_attention", self.rope_parameters)291        rope_type = rope_params.get("rope_type") or ""292        if rope_type.lower() != "llama3":293            return294 295        # Step35 configs can carry per-layer rope_theta as a list; for llama3 rope factors we use the base value.296        rope_theta = self.hparams.get("rope_theta", 10000.0)297        if isinstance(rope_theta, list):298            rope_theta = rope_theta[0]299        base = float(rope_theta)300 301        if (storage_dim := self.hparams.get("head_dim")) is None:302            storage_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]303        storage_dim = int(storage_dim)304 305        # Llama 3 factors apply only to the rotary dims used by full_attention layers306        # (partial_rotary_factor * head_dim). Remaining slots are padded with 1.0 so307        # sliding_attention layers remain unaffected. set_gguf_parameters already308        # guarantees at least one full_attention layer.309        layer_types = (self.hparams.get("layer_types") or [])[: self.block_count]310        partial_rotary_factors = (self.hparams.get("partial_rotary_factors") or [])[: self.block_count]311        full_attention_factor = next(312            float(f) for lt, f in zip(layer_types, partial_rotary_factors) if lt == "full_attention"313        )314        rotary_dim = int(storage_dim * full_attention_factor)315 316        freqs = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim))317 318        factor = float(rope_params.get("factor", 8.0))319        low_freq_factor = float(rope_params.get("low_freq_factor", 1.0))320        high_freq_factor = float(rope_params.get("high_freq_factor", 4.0))321        old_context_len = int(rope_params.get("original_max_position_embeddings", 8192))322 323        low_freq_wavelen = old_context_len / low_freq_factor324        high_freq_wavelen = old_context_len / high_freq_factor325 326        rope_factors: list[float] = []327        for freq in freqs:328            wavelen = 2 * math.pi / float(freq)329            if wavelen < high_freq_wavelen:330                rope_factors.append(1.0)331            elif wavelen > low_freq_wavelen:332                rope_factors.append(factor)333            else:334                smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)335                rope_factors.append(1.0 / ((1.0 - smooth) / factor + smooth))336 337        # Pad to head_dim/2 with 1.0 so non-scaled layers remain neutral.338        if len(rope_factors) < storage_dim // 2:339            rope_factors.extend([1.0] * (storage_dim // 2 - len(rope_factors)))340 341        yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))342