Felipe97/llama-cpp-compiled
01.1k
1#!/usr/bin/env python32# -*- coding: utf-8 -*-3 4from __future__ import annotations5 6import ast7import logging8import contextlib9import json10import os11import re12import sys13from enum import IntEnum14from pathlib import Path15from hashlib import sha25616from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Iterator, Literal, Sequence, TypeVar, cast17from itertools import chain18from transformers import AutoConfig19 20import numpy as np21import torch22 23if TYPE_CHECKING:24 from torch import Tensor25 26if 'NO_LOCAL_GGUF' not in os.environ:27 sys.path.insert(1, str(Path(__file__).parent.parent / 'gguf-py'))28import gguf29from gguf.vocab import MistralTokenizerType, MistralVocab30 31try:32 from mistral_common.tokens.tokenizers.base import TokenizerVersion # type: ignore[import-not-found, ty:unresolved-import]33 from mistral_common.tokens.tokenizers.multimodal import DATASET_MEAN as _MISTRAL_COMMON_DATASET_MEAN, DATASET_STD as _MISTRAL_COMMON_DATASET_STD # type: ignore[import-not-found, ty:unresolved-import]34 from mistral_common.tokens.tokenizers.tekken import Tekkenizer # type: ignore[import-not-found, ty:unresolved-import]35 from mistral_common.tokens.tokenizers.sentencepiece import ( # type: ignore[import-not-found, ty:unresolved-import]36 SentencePieceTokenizer,37 )38 39 _mistral_common_installed = True40 _mistral_import_error_msg = ""41except ImportError:42 _MISTRAL_COMMON_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)43 _MISTRAL_COMMON_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)44 45 _mistral_common_installed = False46 TokenizerVersion: Any = None47 Tekkenizer: Any = None48 SentencePieceTokenizer: Any = None49 _mistral_import_error_msg = (50 "Mistral format requires `mistral-common` to be installed. Please run "51 "`pip install mistral-common[image,audio]` to install it."52 )53 54 55logger = logging.getLogger("hf-to-gguf")56 57 58AnyModel = TypeVar("AnyModel", bound="type[ModelBase]")59 60 61# for checkpoints that ship no config.json, we will try to provide a synthetic one62HparamsMatcher = Callable[[Path], bool]63HparamsLoader = Callable[[Path], dict[str, Any]]64 65 66class SentencePieceTokenTypes(IntEnum):67 NORMAL = 168 UNKNOWN = 269 CONTROL = 370 USER_DEFINED = 471 UNUSED = 572 BYTE = 673 74 75class ModelType(IntEnum):76 TEXT = 177 MMPROJ = 278 79 80class ModelBase:81 _model_classes: dict[ModelType, dict[str, type[ModelBase]]] = {82 ModelType.TEXT: {},83 ModelType.MMPROJ: {},84 }85 _hparams_loaders: list[tuple[HparamsMatcher, HparamsLoader]] = []86 87 dir_model: Path88 ftype: gguf.LlamaFileType89 fname_out: Path90 is_big_endian: bool91 endianess: gguf.GGUFEndian92 use_temp_file: bool93 lazy: bool94 dry_run: bool95 hparams: dict[str, Any]96 model_tensors: dict[str, Callable[[], Tensor]]97 gguf_writer: gguf.GGUFWriter98 model_name: str | None99 metadata_override: Path | None100 metadata: gguf.Metadata101 dir_model_card: Path102 remote_hf_model_id: str | None103 target_model_dir: Path | None104 105 # subclasses should define this!106 model_arch: gguf.MODEL_ARCH107 108 # subclasses should initialize this!109 block_count: int110 tensor_map: gguf.TensorNameMap111 112 # Mistral format specifics113 is_mistral_format: bool = False114 disable_mistral_community_chat_template: bool = False115 sentence_transformers_dense_modules: bool = False116 117 # MTP (multi-token prediction) export modes; set by main() before instantiation.118 # Architectures that implement the filtering/export behavior opt in by119 # setting supports_mtp_export = True on their model class or a mixin.120 supports_mtp_export: bool = False121 mtp_only: bool = False122 no_mtp: bool = False123 124 def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, *, is_big_endian: bool = False,125 use_temp_file: bool = False, eager: bool = False,126 metadata_override: Path | None = None, model_name: str | None = None,127 split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False,128 small_first_shard: bool = False, hparams: dict[str, Any] | None = None, remote_hf_model_id: str | None = None,129 disable_mistral_community_chat_template: bool = False,130 sentence_transformers_dense_modules: bool = False,131 target_model_dir: Path | None = None,132 fuse_gate_up_exps: bool = False,133 fp8_as_q8: bool = False,134 fuse_qkv: bool = False):135 if type(self) is ModelBase or \136 type(self) is TextModel or \137 type(self) is MmprojModel:138 raise TypeError(f"{type(self).__name__!r} should not be directly instantiated")139 140 if self.is_mistral_format and not _mistral_common_installed:141 raise ImportError(_mistral_import_error_msg)142 143 self.dir_model = dir_model144 self.ftype = ftype145 self.fname_out = fname_out146 self.is_big_endian = is_big_endian147 self.endianess = gguf.GGUFEndian.BIG if is_big_endian else gguf.GGUFEndian.LITTLE148 self.use_temp_file = use_temp_file149 self.lazy = not eager or (remote_hf_model_id is not None)150 self.dry_run = dry_run151 self.remote_hf_model_id = remote_hf_model_id152 self.sentence_transformers_dense_modules = sentence_transformers_dense_modules153 self.target_model_dir = target_model_dir154 self.fuse_gate_up_exps = fuse_gate_up_exps155 self._gate_exp_buffer: dict[int, Tensor] = {}156 self._up_exp_buffer: dict[int, Tensor] = {}157 self.fuse_qkv = fuse_qkv158 self._q_buffer: dict[int, Tensor] = {}159 self._k_buffer: dict[int, Tensor] = {}160 self._v_buffer: dict[int, Tensor] = {}161 self._q_bias_buffer: dict[int, Tensor] = {}162 self._k_bias_buffer: dict[int, Tensor] = {}163 self._v_bias_buffer: dict[int, Tensor] = {}164 self._fusable_qkv_weight_layers: set[int] = set()165 self._fusable_qkv_bias_layers: set[int] = set()166 self.hparams = ModelBase.load_hparams(self.dir_model, self.is_mistral_format) if hparams is None else hparams167 self.model_tensors = self.index_tensors(remote_hf_model_id=remote_hf_model_id)168 self.metadata_override = metadata_override169 self.model_name = model_name170 self.dir_model_card = dir_model # overridden in convert_lora_to_gguf.py171 self._is_nvfp4 = False172 self._is_mxfp4 = False173 self._fp8_as_q8 = fp8_as_q8174 self._fp8_dequantized: set[str] = set()175 176 # Apply heuristics to figure out typical tensor encoding based on first tensor's dtype177 # NOTE: can't use field "torch_dtype" in config.json, because some finetunes lie.178 if self.ftype == gguf.LlamaFileType.GUESSED:179 for _, tensor in self.get_tensors():180 if tensor.dim() < 2:181 continue182 183 if tensor.dtype == torch.bfloat16:184 self.ftype = gguf.LlamaFileType.MOSTLY_BF16185 logger.info("heuristics detected bfloat16 tensor dtype, setting --outtype bf16")186 break187 elif tensor.dtype == torch.float16:188 self.ftype = gguf.LlamaFileType.MOSTLY_F16189 logger.info("heuristics detected float16 tensor dtype, setting --outtype f16")190 break191 else:192 self.ftype = gguf.LlamaFileType.MOSTLY_F16193 logger.info("heuristics unable to detect tensor dtype, defaulting to --outtype f16")194 195 # Configure GGUF Writer196 self.gguf_writer = gguf.GGUFWriter(path=None, arch=gguf.MODEL_ARCH_NAMES[self.model_arch], endianess=self.endianess, use_temp_file=self.use_temp_file,197 split_max_tensors=split_max_tensors, split_max_size=split_max_size, dry_run=dry_run, small_first_shard=small_first_shard)198 199 # Mistral specific200 self.disable_mistral_community_chat_template = disable_mistral_community_chat_template201 202 @classmethod203 def add_prefix_to_filename(cls, path: Path, prefix: str) -> Path:204 stem, suffix = path.stem, path.suffix205 new_name = f"{prefix}{stem}{suffix}"206 return path.with_name(new_name)207 208 def find_hparam(self, keys: Iterable[str], optional: bool = False) -> Any:209 key = next((k for k in keys if k in self.hparams), None)210 if key is not None:211 return self.hparams[key]212 if optional:213 return None214 raise KeyError(f"could not find any of: {keys}")215 216 def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:217 tensors: dict[str, Callable[[], Tensor]] = {}218 219 if remote_hf_model_id is not None:220 is_safetensors = True221 222 logger.info(f"Using remote model with HuggingFace id: {remote_hf_model_id}")223 remote_tensors = gguf.utility.SafetensorRemote.get_list_tensors_hf_model(remote_hf_model_id)224 for name, remote_tensor in remote_tensors.items():225 data_gen = lambda r=remote_tensor: LazyTorchTensor.from_remote_tensor(r) # noqa: E731226 if titem := self.filter_tensors((name, data_gen)):227 tname, tgen = titem228 tensors[tname] = tgen229 230 return tensors231 232 prefix = "model" if not self.is_mistral_format else "consolidated"233 part_names: list[str] = ModelBase.get_model_part_names(self.dir_model, prefix, ".safetensors")234 is_safetensors: bool = len(part_names) > 0235 if not is_safetensors:236 part_names = ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin")237 238 tensor_names_from_index: set[str] = set()239 tensor_names_from_parts: set[str] = set()240 241 if not self.is_mistral_format:242 index_name = "model.safetensors" if is_safetensors else "pytorch_model.bin"243 index_name += ".index.json"244 index_file = self.dir_model / index_name245 246 if index_file.is_file():247 logger.info(f"gguf: loading model weight map from '{index_name}'")248 with open(index_file, "r", encoding="utf-8") as f:249 index: dict[str, Any] = json.load(f)250 weight_map = index.get("weight_map")251 if weight_map is None or not isinstance(weight_map, dict):252 raise ValueError(f"Can't load 'weight_map' from {index_name!r}")253 tensor_names_from_index.update(weight_map.keys())254 part_dict: dict[str, None] = dict.fromkeys(weight_map.values(), None) # ty: ignore[invalid-assignment]255 part_names = sorted(part_dict.keys())256 else:257 weight_map = {}258 else:259 weight_map = {}260 261 for part_name in part_names:262 logger.info(f"gguf: indexing model part '{part_name}'")263 ctx: ContextManager[Any]264 if is_safetensors:265 ctx = cast(ContextManager[Any], gguf.utility.SafetensorsLocal(self.dir_model / part_name))266 else:267 ctx = contextlib.nullcontext(torch.load(str(self.dir_model / part_name), map_location="cpu", mmap=True, weights_only=True))268 269 with ctx as model_part:270 assert model_part is not None271 272 for name in model_part.keys():273 tensor_names_from_parts.add(name)274 if is_safetensors:275 data: gguf.utility.LocalTensor = model_part[name]276 if self.lazy:277 data_gen = lambda data=data: LazyTorchTensor.from_local_tensor(data) # noqa: E731278 else:279 dtype = LazyTorchTensor._dtype_str_map[data.dtype]280 data_gen = lambda data=data, dtype=dtype: torch.from_numpy(data.mmap_bytes()).view(dtype).reshape(data.shape) # noqa: E731281 else:282 data_torch: Tensor = model_part[name]283 if self.lazy:284 data_gen = lambda data=data_torch: LazyTorchTensor.from_eager(data) # noqa: E731285 else:286 data_gen = lambda data=data_torch: data # noqa: E731287 if titem := self.filter_tensors((name, data_gen)):288 tname, tgen = titem289 tensors[tname] = tgen290 291 # verify tensor name presence and identify potentially missing files292 if len(tensor_names_from_index) > 0:293 if len(tensor_names_from_parts.symmetric_difference(tensor_names_from_index)) > 0:294 missing = sorted(tensor_names_from_index.difference(tensor_names_from_parts))295 extra = sorted(tensor_names_from_parts.difference(tensor_names_from_index))296 missing_files = sorted(set(weight_map[n] for n in missing if n in weight_map))297 if len(extra) == 0 and len(missing_files) > 0:298 raise ValueError(f"Missing or incomplete model files: {missing_files}\n"299 f"Missing tensors: {missing}")300 else:301 raise ValueError("Mismatch between weight map and model parts for tensor names:\n"302 f"Missing tensors: {missing}\n"303 f"Extra tensors: {extra}")304 305 return tensors306 307 @staticmethod308 def _scale_is_trivial(scale: Tensor) -> bool:309 return scale.numel() <= 1 and abs(float(scale.float().sum()) - 1.0) < 1e-6310 311 def _write_scale_tensor(self, scale_name: str, scale: Tensor):312 if not self._scale_is_trivial(scale):313 scale_f32 = scale.float().numpy().flatten()314 logger.info(f" + {scale_name} (per-tensor scale, shape [{scale_f32.size}])")315 self.gguf_writer.add_tensor(scale_name, scale_f32)316 317 def _write_scales_tensor(self, scale_name: str, scales: list[float]):318 if not np.allclose(scales, 1.0, atol=1e-6):319 scale_vals = np.array(scales, dtype=np.float32)320 logger.info(f" + {scale_name} (per-expert scale, shape [{len(scales)}])")321 self.gguf_writer.add_tensor(scale_name, scale_vals)322 323 def dequant_model(self):324 # If all quantized tensors were already handled (e.g. pure NVFP4), skip325 if self._is_nvfp4 and not any(k.endswith((".weight_scale", ".weight_scale_inv")) for k in self.model_tensors):326 return327 328 tensors_to_remove: list[str] = []329 new_tensors: dict[str, Callable[[], Tensor]] = {}330 331 if (quant_config := self.hparams.get("quantization_config")) and isinstance(quant_config, dict):332 quant_method = quant_config.get("quant_method")333 334 def dequant_bitnet(weight: Tensor, scale: Tensor) -> Tensor:335 weight = weight.view(torch.uint8)336 orig_shape = weight.shape337 338 shift = torch.tensor([0, 2, 4, 6], dtype=torch.uint8).reshape((4, *(1 for _ in range(len(orig_shape)))))339 data = weight.unsqueeze(0).expand((4, *orig_shape)) >> shift340 data = data & 3341 data = (data.float() - 1).reshape((orig_shape[0] * 4, *orig_shape[1:]))342 343 # The scale is inverted344 return data / scale.float()345 346 def dequant_simple(weight: Tensor, scale: Tensor, block_size: Sequence[int] | None = None) -> Tensor:347 scale = scale.float()348 349 if block_size is not None:350 dim_offset = scale.ndim - len(block_size)351 for i, size in enumerate(block_size):352 scale = scale.repeat_interleave(size, dim_offset + i)353 # unpad the scale (e.g. when the tensor size isn't a multiple of the block size)354 scale = scale[tuple(slice(0, size) for size in weight.shape)]355 356 # align scale dims to weight for correct broadcasting (e.g. [128] -> [128, 1, 1])357 while scale.ndim < weight.ndim:358 scale = scale.unsqueeze(-1)359 360 return weight.float() * scale361 362 # ref: https://github.com/ModelCloud/GPTQModel/blob/037c5c0f6c9e33c500d975b038d02e7ca437546d/gptqmodel/nn_modules/qlinear/__init__.py#L437-L476363 def dequant_gptq(g_idx: Tensor, qweight: Tensor, qzeros: Tensor, scales: Tensor) -> Tensor:364 bits = quant_config["bits"]365 assert bits in (2, 3, 4, 8)366 assert qweight.dtype == qzeros.dtype367 maxq = (2 ** bits) - 1368 weight = None369 zeros = None370 pack_dtype_bits = qweight.dtype.itemsize * 8371 372 if bits in [2, 4, 8]:373 pack_factor = pack_dtype_bits // bits374 wf = torch.tensor(list(range(0, pack_dtype_bits, bits)), dtype=torch.int32).unsqueeze(0)375 if self.lazy:376 wf = LazyTorchTensor.from_eager(wf)377 378 zeros = torch.bitwise_right_shift(379 qzeros.unsqueeze(2).expand(-1, -1, pack_factor),380 wf.unsqueeze(0)381 ).to(torch.int16 if bits == 8 else torch.int8)382 zeros = torch.bitwise_and(zeros, maxq).reshape(scales.shape)383 384 weight = torch.bitwise_and(385 torch.bitwise_right_shift(386 qweight.unsqueeze(1).expand(-1, pack_factor, -1),387 wf.unsqueeze(-1)388 ).to(torch.int16 if bits == 8 else torch.int8),389 maxq390 )391 elif bits == 3:392 raise NotImplementedError("3-bit gptq dequantization is not yet implemented")393 394 assert weight is not None395 assert zeros is not None396 397 weight = weight.reshape(weight.shape[0] * weight.shape[1], weight.shape[2])398 399 # gptq_v2 doesn't need to offset zeros400 if quant_config.get("checkpoint_format", "gptq") == "gptq":401 zeros += 1402 403 return (scales[g_idx].float() * (weight - zeros[g_idx]).float()).T404 405 def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: Tensor | None, num_bits: int, group_size: int):406 assert w.dtype == torch.int32407 shape = tuple(shape_tensor.tolist())408 assert len(shape) == 2409 mask = (1 << num_bits) - 1410 411 shifts = torch.arange(0, 32 - (num_bits - 1), num_bits, dtype=torch.int32)412 if self.lazy:413 shifts = LazyTorchTensor.from_eager(shifts)414 415 if zero_point is None:416 offset = 1 << (num_bits - 1)417 else:418 assert len(zero_point.shape) == 2419 offset = (zero_point.unsqueeze(1) >> shifts.reshape(1, -1, 1)) & mask420 offset = offset.reshape(-1, zero_point.shape[1])421 # trim padding, and prepare for broadcast422 # NOTE: the zero-point is packed along dim 0423 offset = offset[:shape[0], :].unsqueeze(-1)424 425 # extract values426 # NOTE: the weights are packed along dim 1427 unpacked = (w.unsqueeze(-1) >> shifts.reshape(1, 1, -1)) & mask428 unpacked = unpacked.reshape(shape[0], -1)429 430 # trim padding431 unpacked = unpacked[:, :shape[1]]432 433 # prepare for broadcast of the scale434 unpacked = unpacked.reshape(shape[0], (unpacked.shape[-1] + group_size - 1) // group_size, group_size)435 unpacked = unpacked - offset436 437 return (unpacked * scale.unsqueeze(-1).float()).reshape(shape)438 439 if quant_method == "bitnet":440 for name in self.model_tensors.keys():441 if name.endswith(".weight_scale"):442 weight_name = name.removesuffix("_scale")443 w = self.model_tensors[weight_name]444 s = self.model_tensors[name]445 self.model_tensors[weight_name] = lambda w=w, s=s: dequant_bitnet(w(), s())446 tensors_to_remove.append(name)447 elif quant_method == "fp8":448 block_size = quant_config.get("weight_block_size")449 for name in self.model_tensors.keys():450 if name.endswith("_scale_inv"):451 weight_name = name.removesuffix("_scale_inv")452 w = self.model_tensors[weight_name]453 s = self.model_tensors[name]454 self.model_tensors[weight_name] = lambda w=w, s=s, bs=block_size: dequant_simple(w(), s(), bs)455 tensors_to_remove.append(name)456 if self._fp8_as_q8:457 self._fp8_dequantized.add(weight_name)458 if name.endswith(".activation_scale"): # unused459 tensors_to_remove.append(name)460 if name.endswith("_activation_scale"): # Mistral-Small-4-119B-2602, unused461 tensors_to_remove.append(name)462 # mistral format463 if name.endswith(".qscale_weight"):464 weight_name = name.removesuffix("qscale_weight") + "weight"465 w = self.model_tensors[weight_name]466 s = self.model_tensors[name]467 self.model_tensors[weight_name] = lambda w=w, s=s, bs=block_size: dequant_simple(w(), s(), bs)468 tensors_to_remove.append(name)469 if self._fp8_as_q8:470 self._fp8_dequantized.add(weight_name)471 if name.endswith(".qscale_act"):472 tensors_to_remove.append(name)473 elif quant_method == "gptq":474 for name in self.model_tensors.keys():475 if name.endswith(".qweight"):476 base_name = name.removesuffix(".qweight")477 g_idx = self.model_tensors[base_name + ".g_idx"]478 qweight = self.model_tensors[base_name + ".qweight"]479 qzeros = self.model_tensors[base_name + ".qzeros"]480 scales = self.model_tensors[base_name + ".scales"]481 new_tensors[base_name + ".weight"] = (482 lambda g=g_idx, z=qzeros, w=qweight, s=scales: dequant_gptq(483 g(), w(), z(), s()484 )485 )486 tensors_to_remove += [487 base_name + n488 for n in (489 ".g_idx",490 ".qzeros",491 ".qweight",492 ".scales",493 )494 ]495 elif quant_method == "compressed-tensors":496 quant_format = quant_config["format"]497 groups = quant_config["config_groups"]498 nvfp4_compressed_tensors = (499 quant_format == "nvfp4-pack-quantized"500 or quant_format == "mixed-precision"501 and bool(groups)502 and all(g.get("format") == "nvfp4-pack-quantized" for g in groups.values() if isinstance(g, dict))503 )504 505 if len(groups) > 1 and not nvfp4_compressed_tensors:506 raise NotImplementedError("Can't handle multiple config groups for compressed-tensors yet")507 weight_config = tuple(groups.values())[0]["weights"]508 509 if quant_format == "float-quantized" or quant_format == "int-quantized" or quant_format == "naive-quantized":510 block_size = weight_config.get("block_structure", None)511 strategy = weight_config.get("strategy")512 assert strategy == "channel" or strategy == "block"513 assert weight_config.get("group_size") is None # didn't find a model using this yet514 is_fp8 = (515 quant_format == "float-quantized"516 and weight_config.get("type") == "float"517 and weight_config.get("num_bits") == 8518 )519 for name in self.model_tensors.keys():520 if name.endswith(".weight_scale"):521 weight_name = name.removesuffix("_scale")522 w = self.model_tensors[weight_name]523 s = self.model_tensors[name]524 self.model_tensors[weight_name] = lambda w=w, s=s: dequant_simple(w(), s(), block_size)525 tensors_to_remove.append(name)526 if self._fp8_as_q8 and is_fp8:527 self._fp8_dequantized.add(weight_name)528 elif quant_format == "pack-quantized":529 assert weight_config.get("strategy") == "group"530 assert weight_config.get("type", "int") == "int"531 num_bits = weight_config.get("num_bits")532 group_size = weight_config.get("group_size")533 assert isinstance(num_bits, int)534 assert isinstance(group_size, int)535 for name in self.model_tensors.keys():536 if name.endswith(".weight_packed"):537 base_name = name.removesuffix("_packed")538 w = self.model_tensors[name]539 scale = self.model_tensors[base_name + "_scale"]540 shape = self.model_tensors[base_name + "_shape"]541 zero_point = self.model_tensors.get(base_name + "_zero_point", lambda: None)542 new_tensors[base_name] = (543 lambda w=w, scale=scale, shape=shape, zero_point=zero_point: dequant_packed(544 w(), scale(), shape(), zero_point(), num_bits, group_size,545 )546 )547 tensors_to_remove += [base_name + n for n in ("_packed", "_shape", "_scale")]548 if (base_name + "_zero_point") in self.model_tensors:549 tensors_to_remove.append(base_name + "_zero_point")550 elif nvfp4_compressed_tensors:551 # Don't error from compressed-tensors, we'll handle them in _generate_nvfp4_tensors552 pass553 else:554 raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported")555 elif quant_method == "modelopt":556 # Mixed-precision ModelOpt models: NVFP4 tensors are handled by557 # _generate_nvfp4_tensors; FP8 tensors have 1D weight_scale and558 # are dequantized here. k/v scale tensors are unused.559 for name in self.model_tensors.keys():560 if name.endswith(".weight_scale"):561 weight_name = name.removesuffix("_scale")562 if weight_name not in self.model_tensors:563 tensors_to_remove.append(name)564 continue565 w = self.model_tensors[weight_name]566 s = self.model_tensors[name]567 is_fp8_weight = False568 if self._fp8_as_q8:569 is_fp8_weight = w().dtype in (torch.float8_e4m3fn, torch.float8_e5m2)570 self.model_tensors[weight_name] = lambda w=w, s=s: dequant_simple(w(), s(), None)571 tensors_to_remove.append(name)572 if is_fp8_weight:573 self._fp8_dequantized.add(weight_name)574 if name.endswith((".input_scale", ".k_scale", ".v_scale")):575 tensors_to_remove.append(name)576 elif quant_method is not None:577 raise NotImplementedError(f"Quant method is not yet supported: {quant_method!r}")578 579 for name in tensors_to_remove:580 if name in self.model_tensors:581 del self.model_tensors[name]582 583 for name, value in new_tensors.items():584 self.model_tensors[name] = value585 586 @classmethod587 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:588 name, gen = item589 590 if name.endswith("e_score_correction_bias"):591 name = name.replace("e_score_correction_bias", "e_score_correction.bias")592 593 if "language_model." in name:594 name = name.replace("language_model.", "")595 596 return name, gen597 598 def get_tensors(self) -> Iterator[tuple[str, Tensor]]:599 for name, gen in self.model_tensors.items():600 yield name, gen()601 602 def format_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None = None, suffix: str = ".weight") -> str:603 if key not in gguf.MODEL_TENSORS[self.model_arch]:604 raise ValueError(f"Missing {key!r} for MODEL_TENSORS of {self.model_arch!r}")605 name: str = gguf.TENSOR_NAMES[key]606 if "{bid}" in name:607 assert bid is not None608 name = name.format(bid=bid)609 return name + suffix610 611 def match_model_tensor_name(self, name: str, key: gguf.MODEL_TENSOR, bid: int | None, suffix: str = ".weight") -> bool:612 if key not in gguf.MODEL_TENSORS[self.model_arch]:613 return False614 key_name: str = gguf.TENSOR_NAMES[key]615 if "{bid}" in key_name:616 if bid is None:617 return False618 key_name = key_name.format(bid=bid)619 else:620 if bid is not None:621 return False622 return name == (key_name + suffix)623 624 def map_tensor_name(self, name: str, try_suffixes: Sequence[str] = (".weight", ".bias")) -> str:625 new_name = self.tensor_map.get_name(key=name, try_suffixes=try_suffixes)626 if new_name is None:627 raise ValueError(f"Can not map tensor {name!r}")628 return new_name629 630 def prepare_qkv_fusion(self) -> None:631 self._fusable_qkv_weight_layers.clear()632 self._fusable_qkv_bias_layers.clear()633 if not self.fuse_qkv or gguf.MODEL_TENSOR.ATTN_QKV not in gguf.MODEL_TENSORS[self.model_arch]:634 return635 636 qkv_types = {637 gguf.MODEL_TENSOR.ATTN_Q,638 gguf.MODEL_TENSOR.ATTN_K,639 gguf.MODEL_TENSOR.ATTN_V,640 }641 weights: dict[int, set[gguf.MODEL_TENSOR]] = {}642 biases: dict[int, set[gguf.MODEL_TENSOR]] = {}643 644 for name in self.model_tensors:645 mapped = self.tensor_map.get_type_and_name(name, try_suffixes=(".weight", ".bias"))646 if mapped is None:647 continue648 tensor_type, new_name = mapped649 if tensor_type not in qkv_types:650 continue651 652 bid = next((int(part) for part in new_name.split(".") if part.isdecimal()), None)653 if bid is None:654 continue655 if new_name.endswith(".weight"):656 weights.setdefault(bid, set()).add(tensor_type)657 elif new_name.endswith(".bias"):658 biases.setdefault(bid, set()).add(tensor_type)659 660 for bid, weight_types in weights.items():661 bias_types = biases.get(bid, set())662 if weight_types == qkv_types and (not bias_types or bias_types == qkv_types):663 self._fusable_qkv_weight_layers.add(bid)664 if bias_types:665 self._fusable_qkv_bias_layers.add(bid)666 667 def set_gguf_parameters(self):668 raise NotImplementedError("set_gguf_parameters() must be implemented in subclasses")669 670 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:671 new_name = self.map_tensor_name(name)672 673 # Handle gate/up expert tensor fusion if enabled674 if self.fuse_gate_up_exps and bid is not None:675 if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_GATE_EXP, bid):676 self._gate_exp_buffer[bid] = data_torch677 elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_UP_EXP, bid):678 self._up_exp_buffer[bid] = data_torch679 680 # Check if both gate and up are buffered for this layer681 if bid in self._gate_exp_buffer and bid in self._up_exp_buffer:682 gate_data = self._gate_exp_buffer.pop(bid)683 up_data = self._up_exp_buffer.pop(bid)684 # gate/up shape: (n_expert, n_ff, n_embd), concatenate to (n_expert, n_ff*2, n_embd)685 fused_data = torch.cat([gate_data, up_data], dim=1)686 fused_name = self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid)687 logger.info(f"Fused gate_exps and up_exps for layer {bid}")688 return [(fused_name, fused_data)]689 690 # If we buffered a gate/up tensor, wait for the other691 if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_GATE_EXP, bid) or \692 self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_UP_EXP, bid):693 return []694 695 # Handle Q/K/V tensor fusion if enabled696 qkv_bid = next((int(part) for part in new_name.split(".") if part.isdecimal()), None) if self.fuse_qkv else None697 if qkv_bid is not None:698 is_bias = new_name.endswith('.bias')699 suffix = '.bias' if is_bias else '.weight'700 fusable_layers = self._fusable_qkv_bias_layers if is_bias else self._fusable_qkv_weight_layers701 if qkv_bid not in fusable_layers:702 return [(new_name, data_torch)]703 704 buf_q = self._q_bias_buffer if is_bias else self._q_buffer705 buf_k = self._k_bias_buffer if is_bias else self._k_buffer706 buf_v = self._v_bias_buffer if is_bias else self._v_buffer707 708 if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_Q, qkv_bid, suffix):709 buf_q[qkv_bid] = data_torch710 elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_K, qkv_bid, suffix):711 buf_k[qkv_bid] = data_torch712 elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_V, qkv_bid, suffix):713 buf_v[qkv_bid] = data_torch714 715 if qkv_bid in buf_q and qkv_bid in buf_k and qkv_bid in buf_v:716 q_data = buf_q.pop(qkv_bid)717 k_data = buf_k.pop(qkv_bid)718 v_data = buf_v.pop(qkv_bid)719 fused_data = torch.cat([q_data, k_data, v_data], dim=0)720 fused_name = self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV, qkv_bid, suffix=suffix)721 logger.info(f"Fused Q, K, V {suffix[1:]} into QKV for layer {qkv_bid}")722 return [(fused_name, fused_data)]723 724 if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_Q, qkv_bid, suffix) or \725 self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_K, qkv_bid, suffix) or \726 self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_V, qkv_bid, suffix):727 return []728 729 return [(new_name, data_torch)]730 731 def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:732 del new_name, bid # unused733 # Force FP8-original tensors to Q8_0 when requested; Q8_0 is faster than F16/BF16.734 if self._fp8_as_q8 and name in self._fp8_dequantized and n_dims >= 2:735 return gguf.GGMLQuantizationType.Q8_0736 return False737 738 # some models need extra generated tensors (like rope_freqs)739 def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:740 return ()741 742 @staticmethod743 def repack_mxfp4_blocks(packed: Tensor, scale: Tensor) -> np.ndarray:744 """745 Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - only moves bits.746 747 Source (compressed-tensors "mxfp4-pack-quantized", also used by DeepSeek-V4):748 packed uint8 [rows, cols/2] element 2i in the low nibble, 2i+1 in the high one749 scale uint8 [rows, cols/32] one E8M0 biased exponent per 32-element group750 751 Destination, per group: one scale byte then 16 code bytes, where byte j holds752 element j in the low nibble and element j+16 in the high one.753 754 The 4-bit codes need no remapping: both sides index into ggml's kvalues_mxfp4755 order. ggml doubles the kvalues and halves the scale, so the value is the same.756 """757 p = packed.contiguous().view(torch.uint8)758 s = scale.contiguous().view(torch.uint8)759 760 rows, packed_cols = p.shape761 cols = packed_cols * 2762 if cols % 32 != 0:763 raise ValueError(f"MXFP4 source row has {cols} values, expected a multiple of 32")764 765 n_blocks = cols // 32766 if tuple(s.shape) != (rows, n_blocks):767 raise ValueError(f"MXFP4 scale shape {tuple(s.shape)} does not match {(rows, n_blocks)}")768 769 src = p.reshape(rows, n_blocks, 16)770 lo = src & 0x0F # elements 0, 2, 4, ...771 hi = (src >> 4) & 0x0F # elements 1, 3, 5, ...772 773 vals = torch.stack((lo, hi), dim=-1).reshape(rows, n_blocks, 32)774 qs = vals[:, :, :16] | (vals[:, :, 16:] << 4)775 776 raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1)777 return raw.reshape(rows, n_blocks * 17).cpu().numpy()778 779 @staticmethod780 def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]:781 """Repack NVFP4 ModelOpt tensors into ggml super-block layout.782 Preserves original E4M3 scale bits as UE4M3 (strip sign bit).783 The per-tensor scale2 factor is stored as a separate tensor and applied at inference time via ggml_mul().784 Returns (raw_data, logical_shape)."""785 786 out_features = weight.shape[0]787 n_blocks = scale.shape[1]788 789 # Unpack ModelOpt nibble-packed weights790 w = weight.reshape(out_features, n_blocks, 8)791 vals = torch.stack([w & 0x0F, w >> 4], dim=-1).reshape(out_features, n_blocks, 16)792 793 # Preserve original E4M3 scale bits as UE4M3 (strip sign bit)794 d_ue = scale.view(torch.uint8).numpy().reshape(out_features, n_blocks) & 0x7F795 qs = (vals[:, :, :8] | (vals[:, :, 8:] << 4)).to(torch.uint8).numpy()796 797 # Pack into super-blocks: [4 UE4M3 scales, 32 qs bytes] = 36 bytes per 64 elements798 n_super = n_blocks // 4799 d_grouped = d_ue.reshape(out_features, n_super, 4)800 qs_grouped = qs.reshape(out_features, n_super, 4, 8).reshape(out_features, n_super, 32)801 raw = np.concatenate([d_grouped, qs_grouped], axis=-1).reshape(out_features, n_super * 36)802 return raw, [out_features, n_super * 64]803 804 def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor, input_scale: Tensor):805 new_name = self.map_tensor_name(name)806 807 raw, shape = self._nvfp4_pack(weight, scale)808 logger.info(f"Repacked {new_name} with shape {shape} and quantization NVFP4")809 self.gguf_writer.add_tensor(new_name, raw, raw_dtype=gguf.GGMLQuantizationType.NVFP4)810 811 self._write_scale_tensor(new_name.replace(".weight", ".scale"), scale2)812 self._write_scale_tensor(new_name.replace(".weight", ".input_scale"), input_scale)813 814 def _generate_nvfp4_tensors(self):815 # Per-layer expert merging to avoid holding all experts in memory816 expert_blocks: dict[tuple[int, str], list[tuple[int, np.ndarray]]] = {}817 expert_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}818 expert_input_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}819 expert_shapes: dict[tuple[int, str], list[int]] = {}820 n_experts = self.find_hparam(["num_local_experts", "num_experts"], optional=True) or 0821 consumed: list[str] = []822 823 for name in self.model_tensors.keys():824 if not name.endswith(".weight"):825 continue826 scale_name = name.replace(".weight", ".weight_scale")827 scale2_name = name.replace(".weight", ".weight_scale_2")828 input_scale_name = name.replace(".weight", ".input_scale")829 if scale_name not in self.model_tensors:830 continue831 # Force eager materialization of lazy tensors832 weight = LazyTorchTensor.to_eager(self.model_tensors[name]())833 scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]())834 835 # Skip non-NVFP4 tensors (e.g. FP8 with per-channel 1D scales)836 if scale.ndim < 2:837 continue838 839 scale2 = LazyTorchTensor.to_eager(self.model_tensors.get(scale2_name, lambda: torch.tensor(1.0))())840 input_scale = LazyTorchTensor.to_eager(self.model_tensors.get(input_scale_name, lambda: torch.tensor(1.0))())841 842 # Mark tensors for removal from model_tensors (already written to gguf)843 consumed.extend([name, scale_name])844 if scale2_name in self.model_tensors:845 consumed.append(scale2_name)846 if input_scale_name in self.model_tensors:847 consumed.append(input_scale_name)848 849 # Check if this is a per-expert tensor850 m = re.search(r'\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$', name)851 if m:852 expert_id = int(m.group(1))853 proj_type = m.group(2)854 bid_m = re.search(r'\.layers\.(\d+)\.', name)855 bid = int(bid_m.group(1)) if bid_m else 0856 key = (bid, proj_type)857 858 raw, shape = self._nvfp4_pack(weight, scale)859 860 if key not in expert_blocks:861 expert_blocks[key] = []862 expert_scales[key] = []863 expert_input_scales[key] = []864 expert_shapes[key] = shape865 expert_blocks[key].append((expert_id, raw.copy()))866 # Collect per-expert scale2 (scalar per expert)867 expert_scales[key].append((expert_id, float(scale2.float().sum())))868 # Collect per-expert input_scale (scalar per expert)869 expert_input_scales[key].append((expert_id, float(input_scale.float().sum())))870 871 # Flush when all experts for this (layer, proj) are collected872 if n_experts > 0 and len(expert_blocks[key]) >= n_experts:873 self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)874 else:875 self._repack_nvfp4(name, weight, scale, scale2, input_scale)876 877 # Flush any remaining experts (fallback if n_experts was unknown)878 for bid, proj_type in list(expert_blocks.keys()):879 self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)880 881 # Remove consumed tensors so get_tensors/modify_tensors won't see them882 for name in consumed:883 self.model_tensors.pop(name, None)884 885 # Remove any remaining unused auxiliary tensors886 for name in list(self.model_tensors.keys()):887 if name.endswith((".k_scale", ".v_scale")):888 del self.model_tensors[name]889 890 def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type):891 experts = expert_blocks.pop(key)892 scales = expert_scales.pop(key)893 input_scales = expert_input_scales.pop(key)894 shape = expert_shapes.pop(key)895 896 experts.sort(key=lambda x: x[0])897 merged = np.stack([e[1] for e in experts], axis=0)898 merged_name = f"model.layers.{bid}.mlp.experts.{proj_type}.weight"899 new_name = self.map_tensor_name(merged_name)900 logger.info(f"Repacked {new_name} with shape [{len(experts)}, {shape[0]}, {shape[1]}] and quantization NVFP4")901 self.gguf_writer.add_tensor(new_name, merged, raw_dtype=gguf.GGMLQuantizationType.NVFP4)902 903 scales.sort(key=lambda x: x[0])904 self._write_scales_tensor(new_name.replace(".weight", ".scale"), [s[1] for s in scales])905 906 input_scales.sort(key=lambda x: x[0])907 self._write_scales_tensor(new_name.replace(".weight", ".input_scale"), [s[1] for s in input_scales])908 909 del experts, merged910 911 def prepare_tensors(self):912 # detect NVFP4 quantization (ModelOpt and Compressed-tensors formats)913 quantization_config = self.hparams.get("quantization_config") or {}914 quant_algo = quantization_config.get("quant_algo")915 quant_method = quantization_config.get("quant_method")916 quant_format = quantization_config.get("format")917 quant_groups = quantization_config.get("config_groups") or {}918 quant_layers = quantization_config.get("quantized_layers") or {}919 quant_config_file = self.dir_model / "hf_quant_config.json"920 921 if (not quant_algo or not quant_layers) and quant_config_file.is_file():922 with open(quant_config_file, "r", encoding="utf-8") as f:923 hf_quant_config = json.load(f)924 quant_config = hf_quant_config.get("quantization") or {}925 producer = hf_quant_config.get("producer") or {}926 producer_name = (producer.get("name") or "").lower()927 if quant_method is None:928 self.hparams.setdefault("quantization_config", {})["quant_method"] = producer_name929 quant_method = producer_name930 quant_algo = quant_config.get("quant_algo", quant_algo)931 quant_method = quant_config.get("quant_method", quant_method)932 quant_format = quant_config.get("format", quant_format)933 quant_groups = quant_config.get("config_groups", quant_groups) or {}934 quant_layers = quant_config.get("quantized_layers", quant_layers) or {}935 936 # Some models use per-tensor quant_algo (e.g. "MIXED_PRECISION" with937 # per-layer NVFP4/FP8) instead of a single global "NVFP4" value.938 nvfp4_compressed_tensors = quant_method == "compressed-tensors" and (939 quant_format == "nvfp4-pack-quantized"940 or quant_format == "mixed-precision"941 and bool(quant_groups)942 and all(g.get("format") == "nvfp4-pack-quantized" for g in quant_groups.values() if isinstance(g, dict))943 )944 if quant_algo != "NVFP4":945 if nvfp4_compressed_tensors:946 quant_algo = "NVFP4"947 elif any(str(v.get("quant_algo")).endswith("NVFP4") for v in quant_layers.values() if isinstance(v, dict)):948 quant_algo = "NVFP4"949 950 self._is_nvfp4 = quant_algo in ("NVFP4", "W4A16_NVFP4")951 self._is_mxfp4 = quant_method == "mxfp4"952 953 # NVFP4 weights are repacked and written directly to gguf_writer.954 # This must run before dequant_model so NVFP4 tensors are removed955 # from model_tensors, leaving only non-NVFP4 (e.g. FP8) for dequant.956 if self._is_nvfp4:957 if nvfp4_compressed_tensors:958 # Convert compressed-tensors 'global' scales into the reciprocal959 def inverse_scale(gen):960 def load():961 scale = LazyTorchTensor.to_eager(gen()).float()962 return 1.0 / scale963 return load964 965 # Change the compressed-tensors names to the ModelOpt names for handling consistently later966 for name in list(self.model_tensors.keys()):967 if name.endswith(".weight_packed"):968 weight_name = name.removesuffix("_packed")969 if weight_name not in self.model_tensors:970 self.model_tensors[weight_name] = self.model_tensors.pop(name)971 elif name.endswith(".weight_global_scale"):972 scale2_name = name.replace(".weight_global_scale", ".weight_scale_2")973 if scale2_name not in self.model_tensors:974 self.model_tensors[scale2_name] = inverse_scale(self.model_tensors.pop(name))975 elif name.endswith(".input_global_scale"):976 input_scale_name = name.replace(".input_global_scale", ".input_scale")977 if input_scale_name not in self.model_tensors:978 self.model_tensors[input_scale_name] = inverse_scale(self.model_tensors.pop(name))979 self._generate_nvfp4_tensors()980 981 self.dequant_model()982 983 self.prepare_qkv_fusion()984 985 # Handle empty tensor_map for models with block_count=0 (like MobileNetV5)986 if self.tensor_map.mapping:987 max_name_len = max(len(s) for _, s in self.tensor_map.mapping.values()) + len(".weight,")988 else:989 max_name_len = len("vision_encoder.weight,") # Default reasonable length990 991 for name, data_torch in chain(self.generate_extra_tensors(), self.get_tensors()):992 # we don't need these993 if name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")):994 continue995 996 old_dtype = data_torch.dtype997 998 # convert any unsupported data types to float32999 if data_torch.dtype not in (torch.float16, torch.float32):1000 data_torch = data_torch.to(torch.float32)1001 1002 # use the first number-like part of the tensor name as the block id1003 bid = None1004 for part in name.split("."):1005 if part.isdecimal():1006 bid = int(part)1007 break1008 1009 for new_name, data_torch in (self.modify_tensors(data_torch, name, bid)):1010 # TODO: why do we squeeze here?1011 # data = data_torch.squeeze().numpy()1012 data = data_torch.numpy()1013 1014 n_dims = len(data.shape)1015 data_qtype: gguf.GGMLQuantizationType | bool = self.tensor_force_quant(name, new_name, bid, n_dims)1016 1017 # Most of the codebase that takes in 1D tensors or norms only handles F32 tensors1018 if n_dims <= 1 or new_name.endswith("_norm.weight"):1019 data_qtype = gguf.GGMLQuantizationType.F321020 1021 # Conditions should closely match those in llama_model_quantize_internal in llama.cpp1022 # Some tensor types are always in float321023 if data_qtype is False and (1024 any(1025 self.match_model_tensor_name(new_name, key, bid)1026 for key in (1027 gguf.MODEL_TENSOR.FFN_GATE_INP,1028 gguf.MODEL_TENSOR.FFN_GATE_INP_SHEXP,1029 gguf.MODEL_TENSOR.POS_EMBD,1030 gguf.MODEL_TENSOR.TOKEN_TYPES,1031 gguf.MODEL_TENSOR.SSM_CONV1D,1032 gguf.MODEL_TENSOR.SHORTCONV_CONV,1033 gguf.MODEL_TENSOR.TIME_MIX_FIRST,1034 gguf.MODEL_TENSOR.TIME_MIX_W1,1035 gguf.MODEL_TENSOR.TIME_MIX_W2,1036 gguf.MODEL_TENSOR.TIME_MIX_DECAY_W1,1037 gguf.MODEL_TENSOR.TIME_MIX_DECAY_W2,1038 gguf.MODEL_TENSOR.TIME_MIX_LERP_FUSED,1039 gguf.MODEL_TENSOR.POSNET_NORM1,1040 gguf.MODEL_TENSOR.POSNET_NORM2,1041 gguf.MODEL_TENSOR.V_ENC_EMBD_POS,1042 gguf.MODEL_TENSOR.A_ENC_EMBD_POS,1043 gguf.MODEL_TENSOR.ALTUP_CORRECT_COEF,1044 gguf.MODEL_TENSOR.ALTUP_PREDICT_COEF,1045 # Kimi KDA conv weights should be F321046 gguf.MODEL_TENSOR.SSM_CONV1D_Q,1047 gguf.MODEL_TENSOR.SSM_CONV1D_K,1048 gguf.MODEL_TENSOR.SSM_CONV1D_V,1049 # DSA indexer weights should be F321050 gguf.MODEL_TENSOR.INDEXER_PROJ,1051 )1052 )1053 or new_name[-7:] not in (".weight", ".lora_a", ".lora_b")1054 ):1055 data_qtype = gguf.GGMLQuantizationType.F321056 1057 if data_qtype is False and any(1058 self.match_model_tensor_name(new_name, key, bid)1059 for key in (1060 gguf.MODEL_TENSOR.TOKEN_EMBD,1061 gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD,1062 gguf.MODEL_TENSOR.OUTPUT,1063 gguf.MODEL_TENSOR.ALTUP_ROUTER,1064 gguf.MODEL_TENSOR.LAUREL_L,1065 gguf.MODEL_TENSOR.LAUREL_R,1066 )1067 ):1068 if self.ftype in (1069 gguf.LlamaFileType.MOSTLY_TQ1_0,1070 gguf.LlamaFileType.MOSTLY_TQ2_0,1071 ):1072 # TODO: use Q4_K and Q6_K1073 data_qtype = gguf.GGMLQuantizationType.F161074 1075 # No override (data_qtype is False), or wants to be quantized (data_qtype is True)1076 if isinstance(data_qtype, bool):1077 if self.ftype == gguf.LlamaFileType.ALL_F32:1078 data_qtype = gguf.GGMLQuantizationType.F321079 elif self.ftype == gguf.LlamaFileType.MOSTLY_F16:1080 data_qtype = gguf.GGMLQuantizationType.F161081 elif self.ftype == gguf.LlamaFileType.MOSTLY_BF16:1082 data_qtype = gguf.GGMLQuantizationType.BF161083 elif self.ftype == gguf.LlamaFileType.MOSTLY_Q8_0:1084 data_qtype = gguf.GGMLQuantizationType.Q8_01085 elif self.ftype == gguf.LlamaFileType.MOSTLY_TQ1_0:1086 data_qtype = gguf.GGMLQuantizationType.TQ1_01087 elif self.ftype == gguf.LlamaFileType.MOSTLY_TQ2_0:1088 data_qtype = gguf.GGMLQuantizationType.TQ2_01089 else:1090 raise ValueError(f"Unknown file type: {self.ftype.name}")1091 1092 # a chunked tensor quantizes as one chunk at a time, while it is written1093 quantize = data.quantize if isinstance(data, gguf.LazyChunkedTensor) else (1094 lambda qtype, d=data: gguf.quants.quantize(d, qtype))1095 1096 try:1097 data = quantize(data_qtype)1098 except gguf.QuantError as e:1099 logger.warning("%s, %s", e, "falling back to F16")1100 data_qtype = gguf.GGMLQuantizationType.F161101 data = quantize(data_qtype)1102 1103 shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape1104 1105 # reverse shape to make it similar to the internal ggml dimension order1106 shape_str = f"{{{', '.join(str(n) for n in reversed(shape))}}}"1107 1108 # n_dims is implicit in the shape1109 logger.info(f"{f'%-{max_name_len}s' % f'{new_name},'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")1110 1111 self.gguf_writer.add_tensor(new_name, data, raw_dtype=data_qtype)1112 1113 qkv_buffers = (1114 self._q_buffer, self._k_buffer, self._v_buffer,1115 self._q_bias_buffer, self._k_bias_buffer, self._v_bias_buffer,1116 )1117 if any(qkv_buffers):1118 raise ValueError("QKV fusion did not consume all buffered tensors")1119 1120 def set_type(self):1121 self.gguf_writer.add_type(gguf.GGUFType.MODEL)1122 1123 def prepare_metadata(self, vocab_only: bool):1124 1125 total_params, shared_params, expert_params, expert_count = self.gguf_writer.get_total_parameter_count()1126 1127 self.metadata = gguf.Metadata.load(self.metadata_override, self.dir_model_card, self.model_name, total_params)1128 1129 # If we are using HF model id, set the metadata name to the model id1130 if self.remote_hf_model_id:1131 self.metadata.name = self.remote_hf_model_id1132 1133 # Fallback to model directory name if metadata name is still missing1134 if self.metadata.name is None:1135 self.metadata.name = self.dir_model.name1136 1137 if self.ftype in (gguf.LlamaFileType.ALL_F32, gguf.LlamaFileType.MOSTLY_F16, gguf.LlamaFileType.MOSTLY_BF16):1138 if self._is_nvfp4:1139 self.ftype = gguf.LlamaFileType.MOSTLY_NVFP41140 elif self._is_mxfp4:1141 self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE1142 1143 # Generate parameter weight class (useful for leader boards) if not yet determined1144 if self.metadata.size_label is None and total_params > 0:1145 self.metadata.size_label = gguf.size_label(total_params, shared_params, expert_params, expert_count)1146 1147 self.set_type()1148 1149 logger.info("Set meta model")1150 self.metadata.set_gguf_meta_model(self.gguf_writer)1151 1152 logger.info("Set model parameters")1153 self.set_gguf_parameters()1154 1155 logger.info("Set model quantization version")1156 self.gguf_writer.add_quantization_version(gguf.GGML_QUANT_VERSION)1157 1158 def write_vocab(self):1159 raise NotImplementedError("write_vocab() must be implemented in subclasses")1160 1161 def write(self):1162 self.prepare_tensors()1163 self.prepare_metadata(vocab_only=False)1164 self.gguf_writer.write_header_to_file(path=self.fname_out)1165 self.gguf_writer.write_kv_data_to_file()1166 self.gguf_writer.write_tensors_to_file(progress=True)1167 self.gguf_writer.close()1168 1169 @staticmethod1170 def get_model_part_names(dir_model: Path, prefix: str, suffix: str) -> list[str]:1171 part_names: list[str] = []1172 for filename in os.listdir(dir_model):1173 if filename.startswith(prefix) and filename.endswith(suffix):1174 part_names.append(filename)1175 1176 part_names.sort()1177 1178 return part_names1179 1180 @staticmethod1181 def load_hparams_guess(dir_model: Path) -> dict[str, Any] | None:1182 # some models ship no config.json, will try to guess them1183 from conversion import load_all_models1184 load_all_models()1185 1186 for matcher, loader in ModelBase._hparams_loaders:1187 if matcher(dir_model):1188 return loader(dir_model)1189 return None1190 1191 @classmethod1192 def register_hparams_loader(cls, matcher: HparamsMatcher) -> Callable[[HparamsLoader], HparamsLoader]:1193 def inner(loader: HparamsLoader) -> HparamsLoader:1194 cls._hparams_loaders.append((matcher, loader))1195 return loader1196 return inner1197 1198 @staticmethod1199 def load_hparams(dir_model: Path, is_mistral_format: bool):1200 if is_mistral_format: