Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3from typing import Any, Callable, Iterable, TYPE_CHECKING4 5import torch6 7if TYPE_CHECKING:8    from torch import Tensor9 10from .base import MmprojModel, ModelBase, TextModel, gguf, logger11 12from .granite import GraniteHybridModel13 14 15@ModelBase.register(16    "NemotronH_Nano_VL_V2",17    "RADIOModel",18)19class NemotronNanoV2VLModel(MmprojModel):20    # ViT-Huge architecture parameters for RADIO v2.5-h21    _vit_hidden_size = 128022    _vit_intermediate_size = 512023    _vit_num_layers = 3224    _vit_num_heads = 1625 26    def get_vision_config(self) -> dict[str, Any] | None:27        # RADIO config doesn't have standard ViT parameters, so they need to be constructed manually28        vision_config = self.global_config.get("vision_config")29        if vision_config is None:30            return None31        # Add ViT-H parameters32        vision_config = {33            **vision_config,34            "hidden_size": self._vit_hidden_size,35            "intermediate_size": self._vit_intermediate_size,36            "num_hidden_layers": self._vit_num_layers,37            "num_attention_heads": self._vit_num_heads,38            "image_size": self.global_config.get("force_image_size", 512),39        }40        return vision_config41 42    def get_audio_config(self) -> dict[str, Any] | None:43        return self.global_config.get("sound_config")44 45    def set_gguf_parameters(self):46        if "image_mean" not in self.preprocessor_config:47            self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406]48        if "image_std" not in self.preprocessor_config:49            self.preprocessor_config["image_std"] = [0.229, 0.224, 0.225]50 51        if self.hparams_audio is not None:52            self.has_vision_encoder = True53            self.has_audio_encoder = True54            self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"])55            self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)56            self.gguf_writer.add_audio_subsampling_factor(self.hparams_audio["subsampling_factor"])57            self.gguf_writer.add_audio_conv_kernel_size(self.hparams_audio["conv_kernel_size"])58            self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.PARAKEET)59            self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)60        else:61            self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)62 63        super().set_gguf_parameters()64        hparams = self.global_config65        self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)66        self.gguf_writer.add_vision_use_gelu(True)67        downsample_ratio = hparams.get("downsample_ratio", 0.5)68        self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio))69 70    def tensor_force_quant(self, name, new_name, bid, n_dims):71        if "sound_encoder" in name or new_name.startswith("mm.a."):72            if "bias" in new_name or "norm" in new_name:73                return gguf.GGMLQuantizationType.F3274            if "conv" in new_name and "weight" in new_name:75                return gguf.GGMLQuantizationType.F3276 77        return super().tensor_force_quant(name, new_name, bid, n_dims)78 79    @classmethod80    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:81        if (titem := super().filter_tensors(item)) is None:82            return None83        name, gen = titem84 85        if "input_conditioner" in name:86            return None87 88        # mtmd does not support video yet so skip tensors related to video.89        if "radio_model.model.patch_generator.video_embedder" in name:90            return None91 92        if not name.startswith(("vision_model.radio_model.model.", "mlp1.", "sound_encoder.", "sound_projection.")):93            return None94 95        if "patch_generator.pos_embed" in name:96            if not name.endswith(".weight"):97                name += ".weight"98 99        # num_batches is only used for training not inference.100        if "conv.norm" in name and "num_batches" in name:101            return None102 103        return name, gen104 105    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:106        # RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it107        if "patch_generator.pos_embed" in name:108            # Downsample position embeddings for fixed 512x512 image size109            import torch.nn.functional as F110            n_embd = self.hparams["hidden_size"]111            image_size = self.global_config.get("force_image_size", 512)112            patch_size = self.hparams["patch_size"]113            target_patches_per_side = image_size // patch_size  # 32114            max_patches_per_side = int((data_torch.shape[1]) ** 0.5)  # 128115            if target_patches_per_side != max_patches_per_side:116                # Reshape to grid, interpolate, flatten back117                data_torch = data_torch.reshape(1, max_patches_per_side, max_patches_per_side, n_embd)118                data_torch = data_torch.permute(0, 3, 1, 2).float()  # [1, n_embd, 128, 128]119                data_torch = F.interpolate(data_torch, size=(target_patches_per_side, target_patches_per_side),120                                           mode='bilinear', align_corners=True)121                data_torch = data_torch.permute(0, 2, 3, 1)  # [1, 32, 32, n_embd]122                data_torch = data_torch.reshape(1, target_patches_per_side * target_patches_per_side, n_embd)123 124        # Reshape linear patch embedding to conv2d format for ggml_conv_2d125        # From [n_embd, patch_size*patch_size*3] to [n_embd, 3, patch_size, patch_size]126        if "patch_generator.embedder" in name:127            patch_size = self.hparams["patch_size"]128            n_embd = self.hparams["hidden_size"]129            data_torch = data_torch.reshape(n_embd, 3, patch_size, patch_size)130 131        if "depthwise_conv.weight" in name:132            data_torch = data_torch.unsqueeze(-1)133            data_torch = data_torch.permute(3, 1, 0, 2).contiguous()134 135        if "pointwise_conv" in name and name.endswith(".weight"):136            if len(data_torch.shape) == 3 and data_torch.shape[2] == 1:137                data_torch = data_torch.reshape(data_torch.shape[0], data_torch.shape[1])138 139        if "subsampling.layers" in name and name.endswith(".bias"):140            if len(data_torch.shape) == 1:141                data_torch = data_torch.reshape(1, -1, 1, 1)142 143        if "pointwise_conv" in name and name.endswith(".bias"):144            if len(data_torch.shape) == 1:145                data_torch = data_torch.reshape(1, -1, 1, 1)146 147        for mapped_name, tensor in super().modify_tensors(data_torch, name, bid):148            if name.startswith("sound_projection.") and mapped_name.startswith("mm.model.mlp."):149                mapped_name = mapped_name.replace("mm.model.mlp.", "mm.a.mlp.")150            yield mapped_name, tensor151 152 153@ModelBase.register("NemotronForCausalLM")154class NemotronModel(TextModel):155    model_arch = gguf.MODEL_ARCH.NEMOTRON156 157    def set_vocab(self):158        self._set_vocab_sentencepiece()159        self.gguf_writer.add_pad_token_id(0)160        self.gguf_writer.add_unk_token_id(1)161 162    def set_gguf_parameters(self):163        super().set_gguf_parameters()164        hparams = self.hparams165        self.gguf_writer.add_vocab_size(hparams["vocab_size"])166 167        f_norm_eps = self.find_hparam(["layer_norm_eps", "layer_norm_epsilon", "norm_epsilon", "norm_eps"])168        self.gguf_writer.add_layer_norm_eps(f_norm_eps)169 170        # * Partial RoPE171        rot_pct = self.rope_parameters["partial_rotary_factor"]172        n_embd = self.find_hparam(["hidden_size", "n_embd"])173        n_head = self.find_hparam(["num_attention_heads", "n_head"])174        self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head)175 176        # * RopeScaling for Nemotron177        factor = self.hparams.get("factor") or self.rope_parameters.get("factor")178        if factor is None:179            self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)180        else:181            self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR)182            self.gguf_writer.add_rope_scaling_factor(factor)183 184    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:185        # * Adding +1 to LayerNorm's weights here to implement layernorm1p w/o changing anything on the GGML engine side186        #   model.layers.{l}.input_layernorm.weight187        #   model.layers.{l}.post_attention_layernorm.weight188        #   model.norm.weight189        if name.endswith("norm.weight"):190            data_torch = data_torch + 1191 192        yield from super().modify_tensors(data_torch, name, bid)193 194 195@ModelBase.register("NemotronHForCausalLM")196class NemotronHModel(GraniteHybridModel):197    """Hybrid mamba2/attention model from NVIDIA"""198    model_arch = gguf.MODEL_ARCH.NEMOTRON_H199    is_moe: bool = False200    supports_mtp_export = True201 202    def __init__(self, *args, **kwargs):203        # We have to determine the correct model architecture (MoE vs non-MoE) before204        # calling the parent __init__. This is because the parent constructor205        # uses self.model_arch to build the tensor name map, and all MoE-specific206        # mappings would be missed if it were called with the default non-MoE arch.207        hparams = ModelBase.load_hparams(args[0], self.is_mistral_format)208        has_moe_params = (209            "num_experts_per_tok" in hparams210            or (isinstance(hparams.get("llm_config"), dict) and "num_experts_per_tok" in hparams["llm_config"])211        )212        if has_moe_params:213            self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE214            self.is_moe = True215 216        super().__init__(*args, **kwargs)217 218        # Save the top-level head_dim for later219        self.head_dim = self.hparams.get("head_dim", self.hparams.get("attention_head_dim"))220        assert self.head_dim is not None, "Could not find the attention head dim in config"221 222        # Don't use expand to calculate d_inner223        self.d_inner = self.find_hparam(["num_heads"]) * self.d_model224 225        # Update the ssm / attn / mlp layers226        # M: Mamba2, *: Attention, -: MLP227        # MoE:228        # M: Mamba2, *: Attention, E: Expert229        pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")230        if pattern is None:231            self._ssm_layers = []232            self._mlp_layers = []233        elif isinstance(pattern, str):234            self._ssm_layers = [i for i, val in enumerate(pattern) if val == "M"]235            self._mlp_layers = [i for i, val in enumerate(pattern) if val == ("E" if self.is_moe else "-")]236        else:237            self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"]238            self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"]239 240        # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head241        self._mtp_bid: int | None = None242        if self.is_moe and not self.no_mtp:243            n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0244            if n_nextn > 0:245                assert n_nextn == 1, (246                    "NemotronH MTP conversion currently supports num_nextn_predict_layers == 1"247                )248                self._mtp_bid = self.block_count249                self.block_count += 1250                # The folded MTP block carries both an attention sub-layer and a251                # MoE sub-layer, so register it as both so the per-layer metadata arrays cover it252                self._attn_layers.append(self._mtp_bid)253                self._mlp_layers.append(self._mtp_bid)254                self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)255 256        if self.mtp_only and self._mtp_bid is None:257            raise ValueError("--mtp was requested, but this model does not contain a supported MTP head")258 259    def get_attn_layers(self):260        pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")261        if pattern is None:262            return []263        assert len(pattern) == self.block_count, f"Mismatch between pattern ({len(pattern)}) and block_count ({self.block_count})!"264        if isinstance(pattern, str):265            return [i for i, val in enumerate(pattern) if val == "*"]266 267        return [i for i, val in enumerate(pattern) if val == "attention"]268 269    @classmethod270    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:271        name, gen = item272        if name.startswith("mtp."):273            # --no-mtp: drop the MTP head entirely274            if cls.no_mtp:275                return None276        elif cls.mtp_only:277            # --mtp: export the MTP head plus the tensors it shares with the target model278            keep = name in (279                "backbone.embeddings.weight",280                "backbone.norm_f.weight",281                "lm_head.weight",282            )283            if not keep:284                return None285        return super().filter_tensors((name, gen))286 287    def prepare_metadata(self, vocab_only: bool):288        from_dir = self.fname_out.is_dir()289        super().prepare_metadata(vocab_only=vocab_only)290 291        if not self.mtp_only or not from_dir:292            return293        output_type: str = self.ftype.name.partition("_")[2]294        fname_default: str = gguf.naming_convention(295            self.metadata.name, self.metadata.basename, self.metadata.finetune,296            self.metadata.version, size_label=None, output_type=output_type, model_type=None)297        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"298 299    def set_gguf_parameters(self):300        super().set_gguf_parameters()301 302        head_dim = self.head_dim303        if head_dim is None:304            raise ValueError("Could not find the attention head dim in config")305        self.gguf_writer.add_key_length(head_dim)306        self.gguf_writer.add_value_length(head_dim)307 308        # Set feed_forward_length309        # NOTE: This will trigger an override warning. This is preferable to310        #   duplicating all the parent logic311        if not self.is_moe:312            n_ff = self.find_hparam(["intermediate_size", "n_inner", "hidden_dim"])313            self.gguf_writer.add_feed_forward_length([314                n_ff if i in self._mlp_layers else 0 for i in range(self.block_count)315            ])316        else:317            moe_intermediate_size = self.hparams["moe_intermediate_size"]318            self.gguf_writer.add_feed_forward_length([319                moe_intermediate_size if i in self._mlp_layers else 0 for i in range(self.block_count)320            ])321            self.gguf_writer.add_expert_used_count(self.hparams["num_experts_per_tok"])322            self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])323            self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"])324            self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])325            self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])326            self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])327            self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])328            self.gguf_writer.add_expert_group_count(self.hparams["n_group"])329 330            # number of experts used per token (top-k)331            if (n_experts_used := self.hparams.get("num_experts_per_tok")) is not None:332                self.gguf_writer.add_expert_used_count(n_experts_used)333 334            if (latent_size := self.hparams.get("moe_latent_size")) is not None:335                self.gguf_writer.add_moe_latent_size(latent_size)336 337        # MTP head: number of trailing NextN blocks338        if self._mtp_bid is not None:339            self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"])340 341    def set_vocab(self):342        # The NemotronH config uses pattern characters (e.g. '-') that may not343        # be supported by the installed transformers version. AutoTokenizer344        # internally calls AutoConfig which triggers this parsing failure.345        # Using trust_remote_code=True to load the model's own config class.346        tokens: list[str] = []347        toktypes: list[int] = []348 349        from transformers import AutoTokenizer350        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)351 352        # Pad vocab size (from Mamba2Model/GraniteHybridModel)353        self.hparams["pad_vocab_size_multiple"] = 8 # Setting this here since GraniteHybridModel.set_vocab() isn't being invoked now.354        # From Mamba2Model.set_vocab():355        vocab_size = self.hparams["vocab_size"]356        pad_vocab = self.hparams.get("pad_vocab_size_multiple", 16)357        # ref: https://stackoverflow.com/a/17511341/22827863358        vocab_size = -(vocab_size // -pad_vocab) * pad_vocab359        self.hparams["vocab_size"] = vocab_size360 361        assert max(tokenizer.vocab.values()) < vocab_size  # ty: ignore[unresolved-attribute]362 363        tokpre = self.get_vocab_base_pre(tokenizer)364 365        reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()}  # ty: ignore[unresolved-attribute]366        added_vocab = tokenizer.get_added_vocab()  # ty: ignore[unresolved-attribute]367 368        added_tokens_decoder = tokenizer.added_tokens_decoder  # ty: ignore[unresolved-attribute]369 370        for i in range(vocab_size):371            if i not in reverse_vocab:372                tokens.append(f"[PAD{i}]")373                toktypes.append(gguf.TokenType.UNUSED)374            else:375                token: str = reverse_vocab[i]376                if token in added_vocab:377                    if not added_tokens_decoder[i].normalized:378                        previous_token = token379                        token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False))  # ty: ignore[unresolved-attribute, invalid-assignment]380                        if previous_token != token:381                            logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")382 383                    if added_tokens_decoder[i].special or self.does_token_look_special(token):384                        toktypes.append(gguf.TokenType.CONTROL)385                    else:386                        token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ")  # pre-normalize user-defined spaces387                        toktypes.append(gguf.TokenType.USER_DEFINED)388                else:389                    toktypes.append(gguf.TokenType.NORMAL)390                tokens.append(token)391 392        # From TextModel.set_vocab_gpt2():393        self.gguf_writer.add_tokenizer_model("gpt2")394        self.gguf_writer.add_tokenizer_pre(tokpre)395        self.gguf_writer.add_token_list(tokens)396        self.gguf_writer.add_token_types(toktypes)397 398        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)399        special_vocab.add_to_gguf(self.gguf_writer)400 401        # The tokenizer _does_ add a BOS token (via post_processor type402        # TemplateProcessing) but does not set add_bos_token to true in the403        # config, so we need to explicitly override it here.404        if not self.is_moe:405            self.gguf_writer.add_add_bos_token(True)406 407    _MTP_SPECIAL_RENAMES = {408        "mtp.layers.0.enorm.weight":           "model.layers.{bid}.enorm.weight",409        "mtp.layers.0.hnorm.weight":           "model.layers.{bid}.hnorm.weight",410        "mtp.layers.0.eh_proj.weight":         "model.layers.{bid}.eh_proj.weight",411        "mtp.layers.1.norm.weight":            "model.layers.{bid}.post_attention_layernorm.weight",412        "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight",413    }414 415    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:416        #   mtp.layers.0: NextN input fusion + attention417        #   mtp.layers.1: MoE + final head norm418        if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")):419            suffix = name.split(".", 3)[3]420            bid = self._mtp_bid421            renamed = self._MTP_SPECIAL_RENAMES.get(name)422            name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}"423 424        if self.is_moe and bid is not None:425            if name.endswith("mixer.gate.e_score_correction.bias"):426                yield from ModelBase.modify_tensors(self, data_torch, name, bid)427                return428 429            if name.endswith("mixer.dt_bias"):430                new_name = name.replace("dt_bias", "dt.bias")431                yield from ModelBase.modify_tensors(self, data_torch, new_name, bid)432                return433 434            if name.endswith("mixer.conv1d.weight"):435                squeezed_data = data_torch.squeeze()436                yield from ModelBase.modify_tensors(self, squeezed_data, name, bid)437                return438 439            if name.endswith("mixer.A_log"):440                transformed_data = -torch.exp(data_torch)441                reshaped_data = transformed_data.squeeze().reshape(-1, 1)442                yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)443                return444 445            if name.endswith("mixer.D"):446                reshaped_data = data_torch.squeeze().reshape(-1, 1)447                yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)448                return449 450            if name.endswith("mixer.norm.weight"):451                reshaped_data = data_torch.reshape(self.n_group, -1)452                yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)453                return454 455            if name.find("mixer.experts") != -1:456                n_experts = self.hparams["n_routed_experts"]457                assert bid is not None458 459                if self._experts is None:460                    self._experts = [{} for _ in range(self.block_count)]461 462                self._experts[bid][name] = data_torch463 464                if len(self._experts[bid]) >= n_experts * 2:465                    # merge the experts into a single tensor466                    for w_name in ["down_proj", "up_proj"]:467                        datas: list[Tensor] = []468 469                        for xid in range(n_experts):470                            ename = f"backbone.layers.{bid}.mixer.experts.{xid}.{w_name}.weight"471                            datas.append(self._experts[bid][ename])472                            del self._experts[bid][ename]473 474                        data_torch = torch.stack(datas, dim=0)475                        merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"476 477                        yield from ModelBase.modify_tensors(self, data_torch, merged_name, bid)478                    return479                else:480                    return481 482        yield from super().modify_tensors(data_torch, name, bid)483 484    def prepare_tensors(self):485        super().prepare_tensors()486 487        if self._experts is not None:488            # flatten `list[dict[str, Tensor]]` into `list[str]`489            experts = [k for d in self._experts for k in d.keys()]490            if len(experts) > 0:491                raise ValueError(f"Unprocessed experts: {experts}")492