char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Kartik Guliaconvert : expand Nemotron H conversion fix (#28689)8c322d5bc

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
28.4 KiB602 linesraw
1from __future__ import annotations
2
3from typing import Any, Callable, Iterable, TYPE_CHECKING
4
5import torch
6
7if TYPE_CHECKING:
8    from pathlib import Path
9    from torch import Tensor
10
11from .base import MmprojModel, ModelBase, TextModel, gguf, logger
12
13from .granite import GraniteHybridModel
14
15
16@ModelBase.register(
17    "NemotronH_Nano_VL_V2",
18    "RADIOModel",
19)
20@ModelBase.example("nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16")
21class NemotronNanoV2VLModel(MmprojModel):
22    # ViT-Huge architecture parameters for RADIO v2.5-h
23    _vit_hidden_size = 1280
24    _vit_intermediate_size = 5120
25    _vit_num_layers = 32
26    _vit_num_heads = 16
27
28    def get_vision_config(self) -> dict[str, Any] | None:
29        # RADIO config doesn't have standard ViT parameters, so they need to be constructed manually
30        vision_config = self.global_config.get("vision_config")
31        if vision_config is None:
32            return None
33        # Add ViT-H parameters
34        vision_config = {
35            **vision_config,
36            "hidden_size": self._vit_hidden_size,
37            "intermediate_size": self._vit_intermediate_size,
38            "num_hidden_layers": self._vit_num_layers,
39            "num_attention_heads": self._vit_num_heads,
40            "image_size": self.global_config.get("force_image_size", 512),
41        }
42        return vision_config
43
44    def get_audio_config(self) -> dict[str, Any] | None:
45        return self.global_config.get("sound_config")
46
47    def set_gguf_parameters(self):
48        if "image_mean" not in self.preprocessor_config:
49            self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406]
50        if "image_std" not in self.preprocessor_config:
51            self.preprocessor_config["image_std"] = [0.229, 0.224, 0.225]
52
53        if self.hparams_audio is not None:
54            self.has_vision_encoder = True
55            self.has_audio_encoder = True
56            self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"])
57            self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
58            self.gguf_writer.add_audio_subsampling_factor(self.hparams_audio["subsampling_factor"])
59            self.gguf_writer.add_audio_conv_kernel_size(self.hparams_audio["conv_kernel_size"])
60            self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.PARAKEET)
61            self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
62        else:
63            self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)
64
65        super().set_gguf_parameters()
66        hparams = self.global_config
67        self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
68        self.gguf_writer.add_vision_use_gelu(True)
69        downsample_ratio = hparams.get("downsample_ratio", 0.5)
70        self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio))
71
72    def tensor_force_quant(self, name, new_name, bid, n_dims):
73        if "sound_encoder" in name or new_name.startswith("mm.a."):
74            if "bias" in new_name or "norm" in new_name:
75                return gguf.GGMLQuantizationType.F32
76            if "conv" in new_name and "weight" in new_name:
77                return gguf.GGMLQuantizationType.F32
78
79        return super().tensor_force_quant(name, new_name, bid, n_dims)
80
81    @classmethod
82    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
83        if (titem := super().filter_tensors(item)) is None:
84            return None
85        name, gen = titem
86
87        if "input_conditioner" in name:
88            return None
89
90        # mtmd does not support video yet so skip tensors related to video.
91        if "radio_model.model.patch_generator.video_embedder" in name:
92            return None
93
94        if not name.startswith(("vision_model.radio_model.model.", "mlp1.", "sound_encoder.", "sound_projection.")):
95            return None
96
97        if "patch_generator.pos_embed" in name:
98            if not name.endswith(".weight"):
99                name += ".weight"
100
101        # num_batches is only used for training not inference.
102        if "conv.norm" in name and "num_batches" in name:
103            return None
104
105        return name, gen
106
107    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
108        # RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it
109        if "patch_generator.pos_embed" in name:
110            # Downsample position embeddings for fixed 512x512 image size
111            import torch.nn.functional as F
112            n_embd = self.hparams["hidden_size"]
113            image_size = self.global_config.get("force_image_size", 512)
114            patch_size = self.hparams["patch_size"]
115            target_patches_per_side = image_size // patch_size  # 32
116            max_patches_per_side = int((data_torch.shape[1]) ** 0.5)  # 128
117            if target_patches_per_side != max_patches_per_side:
118                # Reshape to grid, interpolate, flatten back
119                data_torch = data_torch.reshape(1, max_patches_per_side, max_patches_per_side, n_embd)
120                data_torch = data_torch.permute(0, 3, 1, 2).float()  # [1, n_embd, 128, 128]
121                data_torch = F.interpolate(data_torch, size=(target_patches_per_side, target_patches_per_side),
122                                           mode='bilinear', align_corners=True)
123                data_torch = data_torch.permute(0, 2, 3, 1)  # [1, 32, 32, n_embd]
124                data_torch = data_torch.reshape(1, target_patches_per_side * target_patches_per_side, n_embd)
125
126        # Reshape linear patch embedding to conv2d format for ggml_conv_2d
127        # From [n_embd, patch_size*patch_size*3] to [n_embd, 3, patch_size, patch_size]
128        if "patch_generator.embedder" in name:
129            patch_size = self.hparams["patch_size"]
130            n_embd = self.hparams["hidden_size"]
131            data_torch = data_torch.reshape(n_embd, 3, patch_size, patch_size)
132
133        if "depthwise_conv.weight" in name:
134            data_torch = data_torch.unsqueeze(-1)
135            data_torch = data_torch.permute(3, 1, 0, 2).contiguous()
136
137        if "pointwise_conv" in name and name.endswith(".weight"):
138            if len(data_torch.shape) == 3 and data_torch.shape[2] == 1:
139                data_torch = data_torch.reshape(data_torch.shape[0], data_torch.shape[1])
140
141        if "subsampling.layers" in name and name.endswith(".bias"):
142            if len(data_torch.shape) == 1:
143                data_torch = data_torch.reshape(1, -1, 1, 1)
144
145        if "pointwise_conv" in name and name.endswith(".bias"):
146            if len(data_torch.shape) == 1:
147                data_torch = data_torch.reshape(1, -1, 1, 1)
148
149        for mapped_name, tensor in super().modify_tensors(data_torch, name, bid):
150            if name.startswith("sound_projection.") and mapped_name.startswith("mm.model.mlp."):
151                mapped_name = mapped_name.replace("mm.model.mlp.", "mm.a.mlp.")
152            yield mapped_name, tensor
153
154
155@ModelBase.register("NemotronForCausalLM")
156@ModelBase.example("nvidia/Minitron-4B-Base")
157class NemotronModel(TextModel):
158    model_arch = gguf.MODEL_ARCH.NEMOTRON
159
160    def set_vocab(self):
161        self._set_vocab_sentencepiece()
162        self.gguf_writer.add_pad_token_id(0)
163        self.gguf_writer.add_unk_token_id(1)
164
165    def set_gguf_parameters(self):
166        super().set_gguf_parameters()
167        hparams = self.hparams
168        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
169
170        f_norm_eps = self.find_hparam(["layer_norm_eps", "layer_norm_epsilon", "norm_epsilon", "norm_eps"])
171        self.gguf_writer.add_layer_norm_eps(f_norm_eps)
172
173        # * Partial RoPE
174        rot_pct = self.rope_parameters["partial_rotary_factor"]
175        n_embd = self.find_hparam(["hidden_size", "n_embd"])
176        n_head = self.find_hparam(["num_attention_heads", "n_head"])
177        self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head)
178
179        # * RopeScaling for Nemotron
180        factor = self.hparams.get("factor") or self.rope_parameters.get("factor")
181        if factor is None:
182            self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)
183        else:
184            self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR)
185            self.gguf_writer.add_rope_scaling_factor(factor)
186
187    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
188        # * Adding +1 to LayerNorm's weights here to implement layernorm1p w/o changing anything on the GGML engine side
189        #   model.layers.{l}.input_layernorm.weight
190        #   model.layers.{l}.post_attention_layernorm.weight
191        #   model.norm.weight
192        if name.endswith("norm.weight"):
193            data_torch = data_torch + 1
194
195        yield from super().modify_tensors(data_torch, name, bid)
196
197
198@ModelBase.register("NemotronHForCausalLM")
199@ModelBase.example("nvidia/Nemotron-H-8B-Base-8K")
200class NemotronHModel(GraniteHybridModel):
201    """Hybrid mamba2/attention model from NVIDIA"""
202    model_arch = gguf.MODEL_ARCH.NEMOTRON_H
203    is_moe: bool = False
204    supports_mtp_export = True
205    _experts: list[dict[str, Tensor]] | None = None
206
207    _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
208    _ATTN_LAYER_TYPES = {"attention", "full_attention"}
209    _MLP_LAYER_TYPES = {"moe"}
210
211    def __init__(self, *args, **kwargs):
212        # We have to determine the correct model architecture (MoE vs non-MoE) before
213        # calling the parent __init__. This is because the parent constructor
214        # uses self.model_arch to build the tensor name map, and all MoE-specific
215        # mappings would be missed if it were called with the default non-MoE arch.
216        hparams = kwargs.pop("hparams", None)
217        if hparams is None:
218            hparams = ModelBase.load_hparams(args[0], self.is_mistral_format)
219        llm_config = {**hparams, **(hparams.get("llm_config") or {})}
220
221        has_moe_params = "num_experts_per_tok" in llm_config
222        layers_block_type = llm_config.get("layers_block_type")
223
224        if has_moe_params:
225            self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE
226            self.is_moe = True
227        if layers_block_type is not None:
228            hparams["num_hidden_layers"] = len(layers_block_type)
229
230        super().__init__(*args, hparams=hparams, **kwargs)
231
232        # Save the top-level head_dim for later
233        self.head_dim = self.hparams.get("head_dim", self.hparams.get("attention_head_dim"))
234        assert self.head_dim is not None, "Could not find the attention head dim in config"
235
236        # Don't use expand to calculate d_inner
237        self.d_inner = self.find_hparam(["num_heads"]) * self.d_model
238
239        # Update the ssm / attn / mlp layers
240        # M: Mamba2, *: Attention, -: MLP
241        # MoE:
242        # M: Mamba2, *: Attention, E: Expert
243        pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")
244        if pattern is None:
245            self._ssm_layers = []
246            self._mlp_layers = []
247        elif isinstance(pattern, str):
248            self._ssm_layers = [i for i, val in enumerate(pattern) if val == "M"]
249            self._mlp_layers = [i for i, val in enumerate(pattern) if val == ("E" if self.is_moe else "-")]
250        else:
251            self._ssm_layers = [i for i, val in enumerate(pattern) if val in self._SSM_LAYER_TYPES]
252            self._mlp_layers = [i for i, val in enumerate(pattern) if val in self._MLP_LAYER_TYPES]
253
254        # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head
255        self._mtp_bid: int | None = None
256        if self.is_moe and not self.no_mtp:
257            n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0
258            if n_nextn > 0:
259                assert n_nextn == 1, (
260                    "NemotronH MTP conversion currently supports num_nextn_predict_layers == 1"
261                )
262                self._mtp_bid = self.block_count
263                self.block_count += 1
264                # The folded MTP block carries both an attention sub-layer and a
265                # MoE sub-layer, so register it as both so the per-layer metadata arrays cover it
266                self._attn_layers.append(self._mtp_bid)
267                self._mlp_layers.append(self._mtp_bid)
268                self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
269
270        if self.mtp_only and self._mtp_bid is None:
271            raise ValueError("--mtp was requested, but this model does not contain a supported MTP head")
272
273    def get_attn_layers(self):
274        pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")
275        if pattern is None:
276            return []
277        assert len(pattern) == self.block_count, f"Mismatch between pattern ({len(pattern)}) and block_count ({self.block_count})!"
278        if isinstance(pattern, str):
279            return [i for i, val in enumerate(pattern) if val == "*"]
280
281        return [i for i, val in enumerate(pattern) if val in self._ATTN_LAYER_TYPES]
282
283    @classmethod
284    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
285        name, gen = item
286        if name.startswith("mtp."):
287            # --no-mtp: drop the MTP head entirely
288            if cls.no_mtp:
289                return None
290        elif cls.mtp_only:
291            # --mtp: export the MTP head plus the tensors it shares with the target model
292            # Include lm_head scale sidecars so NVFP4 packing sees them.
293            keep = name in (
294                "backbone.embeddings.weight",
295                "backbone.norm_f.weight",
296                "lm_head.weight",
297                "lm_head.weight_scale",
298                "lm_head.weight_scale_2",
299                "lm_head.weight_scale_inv",
300                "lm_head.input_scale",
301                "lm_head.input_global_scale",
302                "lm_head.weight_global_scale",
303                "lm_head.weight_packed",
304            )
305            if not keep:
306                return None
307        # PEFT names adapter tensors using model.layers.*, while Nemotron-H checkpoints
308        # and the GGUF tensor map use backbone.layers.*
309        if name.startswith("model.layers.") and ".mixer." in name:
310            name = name.replace("model.layers.", "backbone.layers.", 1)
311        return super().filter_tensors((name, gen))
312
313    def prepare_metadata(self, vocab_only: bool):
314        from_dir = self.fname_out.is_dir()
315        super().prepare_metadata(vocab_only=vocab_only)
316
317        if not self.mtp_only or not from_dir:
318            return
319        output_type: str = self.ftype.name.partition("_")[2]
320        fname_default: str = gguf.naming_convention(
321            self.metadata.name, self.metadata.basename, self.metadata.finetune,
322            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
323        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
324
325    def set_gguf_parameters(self):
326        super().set_gguf_parameters()
327
328        head_dim = self.head_dim
329        if head_dim is None:
330            raise ValueError("Could not find the attention head dim in config")
331        self.gguf_writer.add_key_length(head_dim)
332        self.gguf_writer.add_value_length(head_dim)
333
334        # Set feed_forward_length
335        # NOTE: This will trigger an override warning. This is preferable to
336        #   duplicating all the parent logic
337        if not self.is_moe:
338            n_ff = self.find_hparam(["intermediate_size", "n_inner", "hidden_dim"])
339            self.gguf_writer.add_feed_forward_length([
340                n_ff if i in self._mlp_layers else 0 for i in range(self.block_count)
341            ])
342        else:
343            moe_intermediate_size = self.hparams["moe_intermediate_size"]
344            self.gguf_writer.add_feed_forward_length([
345                moe_intermediate_size if i in self._mlp_layers else 0 for i in range(self.block_count)
346            ])
347            self.gguf_writer.add_expert_used_count(self.hparams["num_experts_per_tok"])
348            self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])
349            self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"])
350            self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])
351            self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])
352            self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
353            self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])
354            self.gguf_writer.add_expert_group_count(self.hparams["n_group"])
355
356            # number of experts used per token (top-k)
357            if (n_experts_used := self.hparams.get("num_experts_per_tok")) is not None:
358                self.gguf_writer.add_expert_used_count(n_experts_used)
359
360            if (latent_size := self.hparams.get("moe_latent_size")) is not None:
361                self.gguf_writer.add_moe_latent_size(latent_size)
362
363        # MTP head: number of trailing NextN blocks
364        if self._mtp_bid is not None:
365            self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"])
366
367    def set_vocab(self):
368        # The NemotronH config uses pattern characters (e.g. '-') that may not
369        # be supported by the installed transformers version. AutoTokenizer
370        # internally calls AutoConfig which triggers this parsing failure.
371        # Using trust_remote_code=True to load the model's own config class.
372        tokens: list[str] = []
373        toktypes: list[int] = []
374
375        from transformers import AutoTokenizer
376        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
377
378        # Pad vocab size (from Mamba2Model/GraniteHybridModel)
379        self.hparams["pad_vocab_size_multiple"] = 8 # Setting this here since GraniteHybridModel.set_vocab() isn't being invoked now.
380        # From Mamba2Model.set_vocab():
381        vocab_size = self.hparams["vocab_size"]
382        pad_vocab = self.hparams.get("pad_vocab_size_multiple", 16)
383        # ref: https://stackoverflow.com/a/17511341/22827863
384        vocab_size = -(vocab_size // -pad_vocab) * pad_vocab
385        self.hparams["vocab_size"] = vocab_size
386
387        assert max(tokenizer.vocab.values()) < vocab_size  # ty: ignore[unresolved-attribute]
388
389        tokpre = self.get_vocab_base_pre(tokenizer)
390
391        reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()}  # ty: ignore[unresolved-attribute]
392        added_vocab = tokenizer.get_added_vocab()  # ty: ignore[unresolved-attribute]
393
394        added_tokens_decoder = tokenizer.added_tokens_decoder  # ty: ignore[unresolved-attribute]
395
396        for i in range(vocab_size):
397            if i not in reverse_vocab:
398                tokens.append(f"[PAD{i}]")
399                toktypes.append(gguf.TokenType.UNUSED)
400            else:
401                token: str = reverse_vocab[i]
402                if token in added_vocab:
403                    if not added_tokens_decoder[i].normalized:
404                        previous_token = token
405                        token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False))  # ty: ignore[unresolved-attribute, invalid-assignment]
406                        if previous_token != token:
407                            logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")
408
409                    if added_tokens_decoder[i].special or self.does_token_look_special(token):
410                        toktypes.append(gguf.TokenType.CONTROL)
411                    else:
412                        token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ")  # pre-normalize user-defined spaces
413                        toktypes.append(gguf.TokenType.USER_DEFINED)
414                else:
415                    toktypes.append(gguf.TokenType.NORMAL)
416                tokens.append(token)
417
418        # From TextModel.set_vocab_gpt2():
419        self.gguf_writer.add_tokenizer_model("gpt2")
420        self.gguf_writer.add_tokenizer_pre(tokpre)
421        self.gguf_writer.add_token_list(tokens)
422        self.gguf_writer.add_token_types(toktypes)
423
424        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
425        special_vocab.add_to_gguf(self.gguf_writer)
426
427        # The tokenizer _does_ add a BOS token (via post_processor type
428        # TemplateProcessing) but does not set add_bos_token to true in the
429        # config, so we need to explicitly override it here.
430        if not self.is_moe:
431            self.gguf_writer.add_add_bos_token(True)
432
433    _MTP_SPECIAL_RENAMES = {
434        "mtp.layers.0.enorm.weight":           "model.layers.{bid}.enorm.weight",
435        "mtp.layers.0.hnorm.weight":           "model.layers.{bid}.hnorm.weight",
436        "mtp.layers.0.eh_proj.weight":         "model.layers.{bid}.eh_proj.weight",
437        "mtp.layers.1.norm.weight":            "model.layers.{bid}.post_attention_layernorm.weight",
438        "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight",
439    }
440
441    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
442        #   mtp.layers.0: NextN input fusion + attention
443        #   mtp.layers.1: MoE + final head norm
444        if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")):
445            suffix = name.split(".", 3)[3]
446            bid = self._mtp_bid
447            renamed = self._MTP_SPECIAL_RENAMES.get(name)
448            name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}"
449
450        if self.is_moe and bid is not None:
451            if name.endswith("mixer.gate.e_score_correction.bias"):
452                yield from ModelBase.modify_tensors(self, data_torch, name, bid)
453                return
454
455            if name.endswith("mixer.dt_bias"):
456                new_name = name.replace("dt_bias", "dt.bias")
457                yield from ModelBase.modify_tensors(self, data_torch, new_name, bid)
458                return
459
460            if name.endswith("mixer.conv1d.weight"):
461                squeezed_data = data_torch.squeeze()
462                yield from ModelBase.modify_tensors(self, squeezed_data, name, bid)
463                return
464
465            if name.endswith("mixer.A_log"):
466                transformed_data = -torch.exp(data_torch)
467                reshaped_data = transformed_data.squeeze().reshape(-1, 1)
468                yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)
469                return
470
471            if name.endswith("mixer.D"):
472                reshaped_data = data_torch.squeeze().reshape(-1, 1)
473                yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)
474                return
475
476            if name.endswith("mixer.norm.weight"):
477                reshaped_data = data_torch.reshape(self.n_group, -1)
478                yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)
479                return
480
481            if name.find("mixer.experts") != -1:
482                n_experts = self.hparams["n_routed_experts"]
483                assert bid is not None
484
485                if self._experts is None:
486                    self._experts = [{} for _ in range(self.block_count)]
487
488                self._experts[bid][name] = data_torch
489
490                if len(self._experts[bid]) >= n_experts * 2:
491                    # merge the experts into a single tensor
492                    for w_name in ["down_proj", "up_proj"]:
493                        datas: list[Tensor] = []
494
495                        for xid in range(n_experts):
496                            ename = f"backbone.layers.{bid}.mixer.experts.{xid}.{w_name}.weight"
497                            datas.append(self._experts[bid][ename])
498                            del self._experts[bid][ename]
499
500                        data_torch = torch.stack(datas, dim=0)
501                        merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
502
503                        yield from ModelBase.modify_tensors(self, data_torch, merged_name, bid)
504                    return
505                else:
506                    return
507
508        yield from super().modify_tensors(data_torch, name, bid)
509
510    def prepare_tensors(self):
511        super().prepare_tensors()
512
513        if self._experts is not None:
514            # flatten `list[dict[str, Tensor]]` into `list[str]`
515            experts = [k for d in self._experts for k in d.keys()]
516            if len(experts) > 0:
517                raise ValueError(f"Unprocessed experts: {experts}")
518
519
520@ModelBase.register("NemotronHPuzzleForCausalLM")
521@ModelBase.example("nvidia/NVIDIA-Nemotron-Labs-3-Puzzle-75B-A9B-BF16")
522class NemotronHPuzzleModel(NemotronHModel):
523    """NVIDIA Puzzle: NemotronH with a per-block MoE config (block_configs).
524
525    The checkpoint also ships an MTP draft head (mtp.safetensors). It is skipped
526    here: there is no Puzzle MTP inference path in tree, and the head is laid out
527    by mtp_block_configs rather than the mtp.layers.* form NemotronHModel maps."""
528
529    model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE
530    is_moe: bool = True
531    supports_mtp_export = False
532
533    def __init__(self, dir_model: "Path", *args, **kwargs):
534        hparams = dict(kwargs.pop("hparams", None) or ModelBase.load_hparams(dir_model, self.is_mistral_format))
535
536        self.block_configs: list[dict] = hparams["block_configs"]
537        self.n_layer_trunk = len(self.block_configs)
538
539        # block_configs carries the per-block MoE shape, and is the authority on the
540        # block pattern too: the layers_block_type the HF config wrapper computes is
541        # not sized to it.
542        hparams["num_hidden_layers"] = self.n_layer_trunk
543        hparams["layers_block_type"] = [bc["block_type"] for bc in self.block_configs]
544
545        self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE
546
547        # Bypass NemotronHModel.__init__: it assumes a flat num_experts_per_tok /
548        # moe_intermediate_size and a layers_block_type sized to block_count, neither
549        # of which hold for Puzzle's per-block config.
550        GraniteHybridModel.__init__(self, dir_model, *args, hparams=hparams, **kwargs)
551
552        self.head_dim = self.find_hparam(["head_dim", "attention_head_dim"])
553        self.d_inner = self.find_hparam(["num_heads"]) * self.d_model
554
555        # NemotronHModel.__init__ folds an MTP block into block_count when the
556        # config carries num_nextn_predict_layers; Puzzle's config does, but its
557        # head has a different layout and no inference path, so stay opted out.
558        self._mtp_bid = None
559
560    def set_gguf_parameters(self):
561        GraniteHybridModel.set_gguf_parameters(self)
562
563        head_dim = self.head_dim
564        if head_dim is None:
565            raise ValueError("Could not find the attention head dim in config")
566        self.gguf_writer.add_key_length(head_dim)
567        self.gguf_writer.add_value_length(head_dim)
568
569        ffn_lengths = [bc.get("moe_intermediate_size") or 0 for bc in self.block_configs]
570        experts_used = [bc.get("num_experts_per_tok") or 0 for bc in self.block_configs]
571
572        self.gguf_writer.add_feed_forward_length(ffn_lengths)
573        self.gguf_writer.add_expert_feed_forward_length(ffn_lengths)
574        self.gguf_writer.add_expert_used_count(experts_used)
575
576        self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"])
577        self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])
578        self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])
579        self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])
580        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])
581        self.gguf_writer.add_expert_group_count(self.hparams["n_group"])
582        self.gguf_writer.add_moe_latent_size(self.hparams["moe_latent_size"])
583
584    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
585        # The official BF16 checkpoint (NVIDIA-Nemotron-Labs-3-Puzzle-75B-A9B-BF16)
586        # names the trunk "model.*" (model.layers.*, model.embeddings, model.norm_f)
587        # where the original release used the NemotronH-style "backbone.*", and spells
588        # the router bias "e_score_correction_bias" instead of "e_score_correction.bias";
589        # normalize so both convert identically.
590        if name.startswith("model."):
591            name = "backbone." + name[len("model."):]
592        if name.endswith("mixer.gate.e_score_correction_bias"):
593            name = name[: -len("e_score_correction_bias")] + "e_score_correction.bias"
594
595        yield from super().modify_tensors(data_torch, name, bid)
596
597    @classmethod
598    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
599        # Drop the MTP head unconditionally; see the class docstring.
600        if item[0].startswith("mtp."):
601            return None
602        return super().filter_tensors(item)