# ViT-Huge architecture parameters for RADIO v2.5-h

_vit_hidden_size = 1280

_vit_intermediate_size = 5120

_vit_num_layers = 32

_vit_num_heads = 16

def get_vision_config(self) -> dict[str, Any] | None:

# RADIO config doesn't have standard ViT parameters, so they need to be constructed manually

vision_config = self.global_config.get("vision_config")

if vision_config is None:

return None

# Add ViT-H parameters

vision_config = {

**vision_config,

"hidden_size": self._vit_hidden_size,

"intermediate_size": self._vit_intermediate_size,

"num_hidden_layers": self._vit_num_layers,

"num_attention_heads": self._vit_num_heads,

"image_size": self.global_config.get("force_image_size", 512),

}

return vision_config

def get_audio_config(self) -> dict[str, Any] | None:

return self.global_config.get("sound_config")

def set_gguf_parameters(self):

if "image_mean" not in self.preprocessor_config:

self.preprocessor_config["image_mean"] = [0.485, 0.456, 0.406]

if "image_std" not in self.preprocessor_config:

self.preprocessor_config["image_std"] = [0.229, 0.224, 0.225]

if self.hparams_audio is not None:

self.has_vision_encoder = True

self.has_audio_encoder = True

self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"])

self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)

self.gguf_writer.add_audio_subsampling_factor(self.hparams_audio["subsampling_factor"])

self.gguf_writer.add_audio_conv_kernel_size(self.hparams_audio["conv_kernel_size"])

self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.PARAKEET)

self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)

else:

self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_V2_VL)

super().set_gguf_parameters()

hparams = self.global_config

self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)

self.gguf_writer.add_vision_use_gelu(True)

downsample_ratio = hparams.get("downsample_ratio", 0.5)

self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio))

def tensor_force_quant(self, name, new_name, bid, n_dims):

if "sound_encoder" in name or new_name.startswith("mm.a."):

if "bias" in new_name or "norm" in new_name:

return gguf.GGMLQuantizationType.F32

if "conv" in new_name and "weight" in new_name:

return gguf.GGMLQuantizationType.F32

return super().tensor_force_quant(name, new_name, bid, n_dims)

@classmethod

def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:

if (titem := super().filter_tensors(item)) is None:

return None

name, gen = titem

if "input_conditioner" in name:

return None

# mtmd does not support video yet so skip tensors related to video.

if "radio_model.model.patch_generator.video_embedder" in name:

return None

if not name.startswith(("vision_model.radio_model.model.", "mlp1.", "sound_encoder.", "sound_projection.")):

return None

if "patch_generator.pos_embed" in name:

if not name.endswith(".weight"):

name += ".weight"

# num_batches is only used for training not inference.

if "conv.norm" in name and "num_batches" in name:

return None

return name, gen

def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:

# RADIO's pos_embed doesn't have .weight suffix, but clip.cpp expects it

if "patch_generator.pos_embed" in name:

# Downsample position embeddings for fixed 512x512 image size

import torch.nn.functional as F

n_embd = self.hparams["hidden_size"]

image_size = self.global_config.get("force_image_size", 512)

patch_size = self.hparams["patch_size"]

target_patches_per_side = image_size // patch_size # 32

max_patches_per_side = int((data_torch.shape[1]) ** 0.5) # 128

if target_patches_per_side != max_patches_per_side:

# Reshape to grid, interpolate, flatten back

data_torch = data_torch.reshape(1, max_patches_per_side, max_patches_per_side, n_embd)

data_torch = data_torch.permute(0, 3, 1, 2).float() # [1, n_embd, 128, 128]

data_torch = F.interpolate(data_torch, size=(target_patches_per_side, target_patches_per_side),

mode='bilinear', align_corners=True)

data_torch = data_torch.permute(0, 2, 3, 1) # [1, 32, 32, n_embd]

data_torch = data_torch.reshape(1, target_patches_per_side * target_patches_per_side, n_embd)

# Reshape linear patch embedding to conv2d format for ggml_conv_2d

# From [n_embd, patch_size*patch_size*3] to [n_embd, 3, patch_size, patch_size]

if "patch_generator.embedder" in name:

patch_size = self.hparams["patch_size"]

n_embd = self.hparams["hidden_size"]

data_torch = data_torch.reshape(n_embd, 3, patch_size, patch_size)

if "depthwise_conv.weight" in name:

data_torch = data_torch.unsqueeze(-1)

data_torch = data_torch.permute(3, 1, 0, 2).contiguous()

if "pointwise_conv" in name and name.endswith(".weight"):

if len(data_torch.shape) == 3 and data_torch.shape[2] == 1:

data_torch = data_torch.reshape(data_torch.shape[0], data_torch.shape[1])

if "subsampling.layers" in name and name.endswith(".bias"):

if len(data_torch.shape) == 1:

data_torch = data_torch.reshape(1, -1, 1, 1)

if "pointwise_conv" in name and name.endswith(".bias"):

if len(data_torch.shape) == 1:

data_torch = data_torch.reshape(1, -1, 1, 1)

for mapped_name, tensor in super().modify_tensors(data_torch, name, bid):

if name.startswith("sound_projection.") and mapped_name.startswith("mm.model.mlp."):

mapped_name = mapped_name.replace("mm.model.mlp.", "mm.a.mlp.")

model_arch = gguf.MODEL_ARCH.NEMOTRON

def set_vocab(self):

self._set_vocab_sentencepiece()

self.gguf_writer.add_pad_token_id(0)

self.gguf_writer.add_unk_token_id(1)

def set_gguf_parameters(self):

super().set_gguf_parameters()

hparams = self.hparams

self.gguf_writer.add_vocab_size(hparams["vocab_size"])

f_norm_eps = self.find_hparam(["layer_norm_eps", "layer_norm_epsilon", "norm_epsilon", "norm_eps"])

self.gguf_writer.add_layer_norm_eps(f_norm_eps)

# * Partial RoPE

rot_pct = self.rope_parameters["partial_rotary_factor"]

n_embd = self.find_hparam(["hidden_size", "n_embd"])

n_head = self.find_hparam(["num_attention_heads", "n_head"])

self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head)

# * RopeScaling for Nemotron

factor = self.hparams.get("factor") or self.rope_parameters.get("factor")

if factor is None:

self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)

else:

self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR)

self.gguf_writer.add_rope_scaling_factor(factor)

def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:

# * Adding +1 to LayerNorm's weights here to implement layernorm1p w/o changing anything on the GGML engine side

# model.layers.{l}.input_layernorm.weight

# model.layers.{l}.post_attention_layernorm.weight

# model.norm.weight

if name.endswith("norm.weight"):

data_torch = data_torch + 1

"""Hybrid mamba2/attention model from NVIDIA"""

model_arch = gguf.MODEL_ARCH.NEMOTRON_H

is_moe: bool = False

supports_mtp_export = True

_experts: list[dict[str, Tensor]] | None = None

_SSM_LAYER_TYPES = {"mamba", "linear_attention"}

_ATTN_LAYER_TYPES = {"attention", "full_attention"}

_MLP_LAYER_TYPES = {"moe"}

def __init__(self, *args, **kwargs):

# We have to determine the correct model architecture (MoE vs non-MoE) before

# calling the parent __init__. This is because the parent constructor

# uses self.model_arch to build the tensor name map, and all MoE-specific

# mappings would be missed if it were called with the default non-MoE arch.

hparams = kwargs.pop("hparams", None)

if hparams is None:

hparams = ModelBase.load_hparams(args[0], self.is_mistral_format)

has_moe_params = (

"num_experts_per_tok" in hparams

or (isinstance(hparams.get("llm_config"), dict) and "num_experts_per_tok" in hparams["llm_config"])

)

if has_moe_params:

self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE

self.is_moe = True

layers_block_type = hparams.get("layers_block_type")

if layers_block_type is not None:

hparams["num_hidden_layers"] = len(layers_block_type)

super().__init__(*args, hparams=hparams, **kwargs)

# Save the top-level head_dim for later

self.head_dim = self.hparams.get("head_dim", self.hparams.get("attention_head_dim"))

assert self.head_dim is not None, "Could not find the attention head dim in config"

# Don't use expand to calculate d_inner

self.d_inner = self.find_hparam(["num_heads"]) * self.d_model

# Update the ssm / attn / mlp layers

# M: Mamba2, *: Attention, -: MLP

# MoE:

# M: Mamba2, *: Attention, E: Expert

pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")

if pattern is None:

self._ssm_layers = []

self._mlp_layers = []

elif isinstance(pattern, str):

self._ssm_layers = [i for i, val in enumerate(pattern) if val == "M"]

self._mlp_layers = [i for i, val in enumerate(pattern) if val == ("E" if self.is_moe else "-")]

else:

self._ssm_layers = [i for i, val in enumerate(pattern) if val in self._SSM_LAYER_TYPES]

self._mlp_layers = [i for i, val in enumerate(pattern) if val in self._MLP_LAYER_TYPES]

# `--no-mtp` drops it entirely; `--mtp` exports only the MTP head

self._mtp_bid: int | None = None

if self.is_moe and not self.no_mtp:

n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0

if n_nextn > 0:

assert n_nextn == 1, (

"NemotronH MTP conversion currently supports num_nextn_predict_layers == 1"

)

self._mtp_bid = self.block_count

self.block_count += 1

# The folded MTP block carries both an attention sub-layer and a

# MoE sub-layer, so register it as both so the per-layer metadata arrays cover it

self._attn_layers.append(self._mtp_bid)

self._mlp_layers.append(self._mtp_bid)

self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)

if self.mtp_only and self._mtp_bid is None:

raise ValueError("--mtp was requested, but this model does not contain a supported MTP head")

def get_attn_layers(self):

pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")

if pattern is None:

return []

assert len(pattern) == self.block_count, f"Mismatch between pattern ({len(pattern)}) and block_count ({self.block_count})!"

if isinstance(pattern, str):

return [i for i, val in enumerate(pattern) if val == "*"]

return [i for i, val in enumerate(pattern) if val in self._ATTN_LAYER_TYPES]

@classmethod

def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:

name, gen = item

if name.startswith("mtp."):

# --no-mtp: drop the MTP head entirely

if cls.no_mtp:

return None

elif cls.mtp_only:

# --mtp: export the MTP head plus the tensors it shares with the target model

# Include lm_head scale sidecars so NVFP4 packing sees them.

keep = name in (

"backbone.embeddings.weight",

"backbone.norm_f.weight",

"lm_head.weight",

"lm_head.weight_scale",

"lm_head.weight_scale_2",

"lm_head.weight_scale_inv",

"lm_head.input_scale",

"lm_head.input_global_scale",

"lm_head.weight_global_scale",

"lm_head.weight_packed",

)

if not keep:

return None

# PEFT names adapter tensors using model.layers.*, while Nemotron-H checkpoints

# and the GGUF tensor map use backbone.layers.*

if name.startswith("model.layers.") and ".mixer." in name:

name = name.replace("model.layers.", "backbone.layers.", 1)

return super().filter_tensors((name, gen))

def prepare_metadata(self, vocab_only: bool):

from_dir = self.fname_out.is_dir()

super().prepare_metadata(vocab_only=vocab_only)

if not self.mtp_only or not from_dir:

return

output_type: str = self.ftype.name.partition("_")[2]

fname_default: str = gguf.naming_convention(

self.metadata.name, self.metadata.basename, self.metadata.finetune,

self.metadata.version, size_label=None, output_type=output_type, model_type=None)

self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"

def set_gguf_parameters(self):

super().set_gguf_parameters()

head_dim = self.head_dim

if head_dim is None:

raise ValueError("Could not find the attention head dim in config")

self.gguf_writer.add_key_length(head_dim)

self.gguf_writer.add_value_length(head_dim)

# Set feed_forward_length

# NOTE: This will trigger an override warning. This is preferable to

# duplicating all the parent logic

if not self.is_moe:

n_ff = self.find_hparam(["intermediate_size", "n_inner", "hidden_dim"])

self.gguf_writer.add_feed_forward_length([

n_ff if i in self._mlp_layers else 0 for i in range(self.block_count)

])

else:

moe_intermediate_size = self.hparams["moe_intermediate_size"]

self.gguf_writer.add_feed_forward_length([

moe_intermediate_size if i in self._mlp_layers else 0 for i in range(self.block_count)

])

self.gguf_writer.add_expert_used_count(self.hparams["num_experts_per_tok"])

self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])

self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"])

self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])

self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])

self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])

self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])

self.gguf_writer.add_expert_group_count(self.hparams["n_group"])

# number of experts used per token (top-k)

if (n_experts_used := self.hparams.get("num_experts_per_tok")) is not None:

self.gguf_writer.add_expert_used_count(n_experts_used)

if (latent_size := self.hparams.get("moe_latent_size")) is not None:

self.gguf_writer.add_moe_latent_size(latent_size)

# MTP head: number of trailing NextN blocks

if self._mtp_bid is not None:

self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"])

def set_vocab(self):

# The NemotronH config uses pattern characters (e.g. '-') that may not

# be supported by the installed transformers version. AutoTokenizer

# internally calls AutoConfig which triggers this parsing failure.

# Using trust_remote_code=True to load the model's own config class.

tokens: list[str] = []

toktypes: list[int] = []

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)

# Pad vocab size (from Mamba2Model/GraniteHybridModel)

self.hparams["pad_vocab_size_multiple"] = 8 # Setting this here since GraniteHybridModel.set_vocab() isn't being invoked now.

# From Mamba2Model.set_vocab():

vocab_size = self.hparams["vocab_size"]

pad_vocab = self.hparams.get("pad_vocab_size_multiple", 16)

# ref: https://stackoverflow.com/a/17511341/22827863

vocab_size = -(vocab_size // -pad_vocab) * pad_vocab

self.hparams["vocab_size"] = vocab_size

assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute]

tokpre = self.get_vocab_base_pre(tokenizer)

reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} # ty: ignore[unresolved-attribute]

added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute]

added_tokens_decoder = tokenizer.added_tokens_decoder # ty: ignore[unresolved-attribute]

for i in range(vocab_size):

if i not in reverse_vocab:

tokens.append(f"[PAD{i}]")

toktypes.append(gguf.TokenType.UNUSED)

else:

token: str = reverse_vocab[i]

if token in added_vocab:

if not added_tokens_decoder[i].normalized:

previous_token = token

token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) # ty: ignore[unresolved-attribute, invalid-assignment]

if previous_token != token:

logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")

if added_tokens_decoder[i].special or self.does_token_look_special(token):

toktypes.append(gguf.TokenType.CONTROL)

else:

token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ") # pre-normalize user-defined spaces

toktypes.append(gguf.TokenType.USER_DEFINED)

else:

toktypes.append(gguf.TokenType.NORMAL)

tokens.append(token)

# From TextModel.set_vocab_gpt2():

self.gguf_writer.add_tokenizer_model("gpt2")

self.gguf_writer.add_tokenizer_pre(tokpre)

self.gguf_writer.add_token_list(tokens)

self.gguf_writer.add_token_types(toktypes)

special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)

special_vocab.add_to_gguf(self.gguf_writer)

# The tokenizer _does_ add a BOS token (via post_processor type

# TemplateProcessing) but does not set add_bos_token to true in the

# config, so we need to explicitly override it here.

if not self.is_moe:

self.gguf_writer.add_add_bos_token(True)

_MTP_SPECIAL_RENAMES = {

"mtp.layers.0.enorm.weight": "model.layers.{bid}.enorm.weight",

"mtp.layers.0.hnorm.weight": "model.layers.{bid}.hnorm.weight",

"mtp.layers.0.eh_proj.weight": "model.layers.{bid}.eh_proj.weight",

"mtp.layers.1.norm.weight": "model.layers.{bid}.post_attention_layernorm.weight",

"mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight",

}

def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:

# mtp.layers.0: NextN input fusion + attention

# mtp.layers.1: MoE + final head norm

if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")):

suffix = name.split(".", 3)[3]

bid = self._mtp_bid

renamed = self._MTP_SPECIAL_RENAMES.get(name)

name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}"

if self.is_moe and bid is not None:

if name.endswith("mixer.gate.e_score_correction.bias"):

yield from ModelBase.modify_tensors(self, data_torch, name, bid)

return

if name.endswith("mixer.dt_bias"):

new_name = name.replace("dt_bias", "dt.bias")

yield from ModelBase.modify_tensors(self, data_torch, new_name, bid)

return

if name.endswith("mixer.conv1d.weight"):

squeezed_data = data_torch.squeeze()

yield from ModelBase.modify_tensors(self, squeezed_data, name, bid)

return

if name.endswith("mixer.A_log"):

transformed_data = -torch.exp(data_torch)

reshaped_data = transformed_data.squeeze().reshape(-1, 1)

yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)

return

if name.endswith("mixer.D"):

reshaped_data = data_torch.squeeze().reshape(-1, 1)

yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)

return

if name.endswith("mixer.norm.weight"):

reshaped_data = data_torch.reshape(self.n_group, -1)

yield from ModelBase.modify_tensors(self, reshaped_data, name, bid)

return

if name.find("mixer.experts") != -1:

n_experts = self.hparams["n_routed_experts"]

assert bid is not None

if self._experts is None:

self._experts = [{} for _ in range(self.block_count)]

self._experts[bid][name] = data_torch

if len(self._experts[bid]) >= n_experts * 2:

# merge the experts into a single tensor

for w_name in ["down_proj", "up_proj"]:

datas: list[Tensor] = []

for xid in range(n_experts):

ename = f"backbone.layers.{bid}.mixer.experts.{xid}.{w_name}.weight"

datas.append(self._experts[bid][ename])

del self._experts[bid][ename]

data_torch = torch.stack(datas, dim=0)

merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"

yield from ModelBase.modify_tensors(self, data_torch, merged_name, bid)

return

else:

return

yield from super().modify_tensors(data_torch, name, bid)

def prepare_tensors(self):

super().prepare_tensors()

if self._experts is not None:

# flatten `list[dict[str, Tensor]]` into `list[str]`

experts = [k for d in self._experts for k in d.keys()]

if len(experts) > 0:

"""NVIDIA Puzzle: NemotronH with a per-block MoE config (block_configs).

model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE

is_moe: bool = True

supports_mtp_export = False

def __init__(self, dir_model: "Path", *args, **kwargs):

hparams = dict(kwargs.pop("hparams", None) or ModelBase.load_hparams(dir_model, self.is_mistral_format))

self.block_configs: list[dict] = hparams["block_configs"]

self.n_layer_trunk = len(self.block_configs)

# block_configs carries the per-block MoE shape, and is the authority on the

# block pattern too: the layers_block_type the HF config wrapper computes is

# not sized to it.

hparams["num_hidden_layers"] = self.n_layer_trunk

hparams["layers_block_type"] = [bc["block_type"] for bc in self.block_configs]

self.model_arch = gguf.MODEL_ARCH.NEMOTRON_H_MOE

# Bypass NemotronHModel.__init__: it assumes a flat num_experts_per_tok /

# moe_intermediate_size and a layers_block_type sized to block_count, neither

# of which hold for Puzzle's per-block config.

GraniteHybridModel.__init__(self, dir_model, *args, hparams=hparams, **kwargs)

self.head_dim = self.find_hparam(["head_dim", "attention_head_dim"])

self.d_inner = self.find_hparam(["num_heads"]) * self.d_model

# NemotronHModel.__init__ folds an MTP block into block_count when the

# config carries num_nextn_predict_layers; Puzzle's config does, but its

# head has a different layout and no inference path, so stay opted out.

self._mtp_bid = None

def set_gguf_parameters(self):

GraniteHybridModel.set_gguf_parameters(self)

head_dim = self.head_dim

if head_dim is None:

raise ValueError("Could not find the attention head dim in config")

self.gguf_writer.add_key_length(head_dim)

self.gguf_writer.add_value_length(head_dim)

ffn_lengths = [bc.get("moe_intermediate_size") or 0 for bc in self.block_configs]

experts_used = [bc.get("num_experts_per_tok") or 0 for bc in self.block_configs]

self.gguf_writer.add_feed_forward_length(ffn_lengths)

self.gguf_writer.add_expert_feed_forward_length(ffn_lengths)

self.gguf_writer.add_expert_used_count(experts_used)

self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["moe_shared_expert_intermediate_size"])

self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"])

self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"])

self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"])

self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])

self.gguf_writer.add_expert_group_count(self.hparams["n_group"])

self.gguf_writer.add_moe_latent_size(self.hparams["moe_latent_size"])

def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:

# The official BF16 checkpoint (NVIDIA-Nemotron-Labs-3-Puzzle-75B-A9B-BF16)

# names the trunk "model.*" (model.layers.*, model.embeddings, model.norm_f)

# where the original release used the NemotronH-style "backbone.*", and spells

# the router bias "e_score_correction_bias" instead of "e_score_correction.bias";

# normalize so both convert identically.

if name.startswith("model."):

name = "backbone." + name[len("model."):]

if name.endswith("mixer.gate.e_score_correction_bias"):

name = name[: -len("e_score_correction_bias")] + "e_score_correction.bias"

yield from super().modify_tensors(data_torch, name, bid)

@classmethod

def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:

# Drop the MTP head unconditionally; see the class docstring.

if item[0].startswith("mtp."):

return None