diff --git a/docs/source/en/gguf.md b/docs/source/en/gguf.md index 0d3343ef22e5..cb4b1e5427e5 100644 --- a/docs/source/en/gguf.md +++ b/docs/source/en/gguf.md @@ -24,10 +24,10 @@ rendered properly in your Markdown viewer. The GGUF format also supports many quantized data types (refer to [quantization type table](https://hf.co/docs/hub/en/gguf#quantization-types) for a complete list of supported quantization types) which saves a significant amount of memory, making inference with large models like Whisper and Llama feasible on local and edge devices. -Transformers supports loading models stored in the GGUF format for further training or finetuning. The GGUF checkpoint is **dequantized to fp32** where the full model weights are available and compatible with PyTorch. +Transformers supports loading models stored in the GGUF format for further training or finetuning. The GGUF checkpoint is **dequantized lazily on the target device** by the conversion-mapping pipeline (see [Adding GGUF support for a new architecture](#adding-gguf-support-for-a-new-architecture) below for how that's plugged in). > [!TIP] -> Models that support GGUF include Llama, Mistral, Qwen2, Qwen2Moe, Phi3, Bloom, Falcon, StableLM, GPT2, Starcoder2, and [more](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/ggml.py) +> Architectures wired up for GGUF loading include Llama, Mistral, Phi3, Cohere, Qwen2, Qwen3, Deci, StableLM, Starcoder2, Nemotron, Gemma2, Gemma3 (text + multimodal), Bloom, GPT2, Mamba, LFM2, Falcon, Qwen2-MoE, Qwen3-MoE, MiniMax-M2, GPT-OSS, T5, UMT5. The full list can be fetched with `_GGUF_ARCH_CONVERTERS.keys()` in [`modeling_gguf_pytorch_utils.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_gguf_pytorch_utils.py). Add the `gguf_file` parameter to [`~PreTrainedModel.from_pretrained`] to specify the GGUF file to load. @@ -52,3 +52,82 @@ model.save_pretrained("directory") !python ${path_to_llama_cpp}/convert-hf-to-gguf.py ${directory} ``` + +## Adding GGUF support for a new architecture + +GGUF loading plugs into the same weight-conversion framework that the rest of `transformers` uses for safetensors / Fp8 / etc. To support a new architecture you only need to: + +1. **Write a list of `WeightRenaming` / `WeightConverter` rules** in [`src/transformers/modeling_gguf_pytorch_utils.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_gguf_pytorch_utils.py) and register it in `_GGUF_ARCH_CONVERTERS`. You can also just combine existing rules whenever relevant. You then map the `model_type` to the rules. +2. **Add the `model_type` to `EXPECTED_MODEL_TYPES`** in [`tests/quantization/ggml/test_gguf_arch_coverage.py`](https://github.com/huggingface/transformers/blob/main/tests/quantization/ggml/test_gguf_arch_coverage.py) so the fast regression test enforces it. + +### Rule types + +For more details you should check the doc [dynamic weight loading](https://huggingface.co/docs/transformers/weightconverter). +As a reminder it looks like this: + +```python +from transformers.core_model_loading import WeightConverter, WeightRenaming + +# Pure key rename (cheap — no op chain). Substring `re.sub` style: escape dots, +# use alternation when the same target template covers several source names. +WeightRenaming(r"^blk\.", "model.layers.") +WeightRenaming(r"\.ffn_(gate|up|down)\.weight", r".mlp.\1_proj.weight") + +# With an actual tensor transform (transpose, permute, concat, etc.). +WeightConverter( + source_patterns=r"\.attn_q\.weight", + target_patterns=".self_attn.q_proj.weight", + operations=[ReversePermuteAttnQ()], +) +``` + +`WeightRenaming`s are applied **sequentially** (each one operates on the previous rule's output), so structural prefix renames (`^blk\.` → `model.layers.`, etc.) should come first and per-tensor renames after. `WeightConverter`s are evaluated after all the renames; the first one whose source pattern matches is selected. + +`GGUFDequantize` is automatically prepended to every `WeightConverter`'s op chain by `GGUFQuantizer.update_weight_conversions`, updating the keys to dequantize. + +### Existing transform ops + +We added a bunch of weight conversion operations that are for now sufficient to cover all the cases. They are in [`src/transformers/gguf_conversion_ops.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/gguf_conversion_ops.py). + +| Op | Used for | +|----|----------| +| `ReversePermuteAttnQ`, `ReversePermuteAttnK` | Llama-family rope-permuted Q/K weights | +| `SubtractOne` | Gemma / Nemotron norm weights (stored as `weight + 1` in GGUF) | +| `Unsqueeze(dim)` | Mamba conv1d, LFM2 shortconv — add a singleton dim | +| `LogNegate` | Mamba SSM-A: `log(-x)` on load | +| `BloomReshapeQKVWeight`, `BloomReshapeQKVBias` | Bloom interleaved QKV | +| `Concatenate(dim)` (from `core_model_loading`) | MoE gate+up merge into `gate_up_proj` | +| `Transpose` (from `core_model_loading`) | GPT-2 transposed `c_attn` / `c_proj` / `c_fc` | + +If your architecture needs a new tensor transform, add it to `gguf_conversion_ops.py` as a `ConversionOps` subclass with `convert(input_dict, source_patterns, target_patterns, **kwargs)` and a `reverse_op` property. + +### Example: adding a new Llama-like architecture + +For a model whose HF `model_type` is `"my_llama"` and which uses the standard rope-permuted Q/K, plain Llama norms, and `mlp.gate_proj` / `up_proj` / `down_proj`, the rule list is already provided as `_LLAMA_CONVERTERS`. You only need: + +```python +# src/transformers/modeling_gguf_pytorch_utils.py +_GGUF_ARCH_CONVERTERS = { + ... + "my_llama": _LLAMA_CONVERTERS, +} +``` + +For a variant with `q_norm` / `k_norm` and attn biases (StableLM style) you compose: + +```python +_MY_LLAMA_CONVERTERS = _LLAMA_CONVERTERS + [ + WeightRenaming(r"\.attn_(q|k)_norm\.weight", r".self_attn.\1_norm.weight"), + WeightRenaming(r"\.attn_(q|k|v)\.bias", r".self_attn.\1_proj.bias"), +] +``` + +For a wholly new layout, start from one of the existing entries (`_BLOOM_CONVERTERS`, `_GPT2_CONVERTERS`, `_T5_CONVERTERS`, …) and adapt the renames + add any required `WeightConverter` ops. + +### Dequantization performance + +The actual byte-level dequant happens inside the `GGUFDequantize` op via [`integrations/gguf_dequant.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/gguf_dequant.py), a pure-PyTorch port of [city96/ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) (same kernels diffusers uses). + +It supports Q4_0/Q4_1, Q5_0/Q5_1, Q8_0, Q2_K…Q6_K, IQ4_NL, IQ4_XS, BF16, F16, F32 and runs on CPU or GPU/MPS. + +If you add a new quant type, add the corresponding `_dq_` kernel to `gguf_dequant.py` and register it in `_build_dispatch()`. diff --git a/docs/source/en/main_classes/quantization.md b/docs/source/en/main_classes/quantization.md index 13540d4d928e..306e7ab08ee5 100755 --- a/docs/source/en/main_classes/quantization.md +++ b/docs/source/en/main_classes/quantization.md @@ -117,3 +117,7 @@ Learn how to quantize models in the [Quantization](../quantization) guide. ## SinqConfig [[autodoc]] SinqConfig + +## GgufQuantizeConfig + +[[autodoc]] GgufQuantizeConfig diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 9b05c92dabf1..2dde14201c56 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -259,6 +259,7 @@ "FineGrainedFP8Config", "FourOverSixConfig", "FPQuantConfig", + "GgufQuantizeConfig", "GPTQConfig", "HiggsConfig", "HqqConfig", @@ -775,6 +776,7 @@ from .utils.quantization_config import FineGrainedFP8Config as FineGrainedFP8Config from .utils.quantization_config import FourOverSixConfig as FourOverSixConfig from .utils.quantization_config import FPQuantConfig as FPQuantConfig + from .utils.quantization_config import GgufQuantizeConfig as GgufQuantizeConfig from .utils.quantization_config import GPTQConfig as GPTQConfig from .utils.quantization_config import HiggsConfig as HiggsConfig from .utils.quantization_config import HqqConfig as HqqConfig diff --git a/src/transformers/cli/serving/model_manager.py b/src/transformers/cli/serving/model_manager.py index 56db99698e5a..ed54f8a2360f 100644 --- a/src/transformers/cli/serving/model_manager.py +++ b/src/transformers/cli/serving/model_manager.py @@ -228,10 +228,22 @@ def _load_processor(self, model_id_and_revision: str) -> "ProcessorMixin | PreTr Args: model_id_and_revision: Model ID in ``'model_id@revision'`` format. """ + import os as _os + from transformers import AutoProcessor + from transformers.modeling_utils import _maybe_auto_detect_gguf model_id, revision = model_id_and_revision.split("@", 1) - return AutoProcessor.from_pretrained(model_id, revision=revision, trust_remote_code=self.trust_remote_code) + kwargs = {"revision": revision, "trust_remote_code": self.trust_remote_code} + # GGUF: the tokenizer lives inside the .gguf; point the loader at its dir + filename. + if model_id.lower().endswith(".gguf") and _os.path.isfile(model_id): + kwargs["gguf_file"] = _os.path.basename(model_id) + model_id = _os.path.dirname(model_id) + else: + _gguf = _maybe_auto_detect_gguf(model_id, revision=revision) + if _gguf is not None: + kwargs["gguf_file"] = _gguf + return AutoProcessor.from_pretrained(model_id, **kwargs) def _load_model( self, model_id_and_revision: str, tqdm_class: type | None = None, progress_callback: Callable | None = None @@ -260,6 +272,20 @@ def _load_model( "tqdm_class": tqdm_class, } + # GGUF: serve receives only a model id. Resolve a single .gguf + pass gguf_file so + # AutoConfig + from_pretrained read config/weights from it. Loaders expect a repo/dir + # + filename, so for a local .gguf path point model_id at its directory. + import os as _os + from transformers.modeling_utils import _maybe_auto_detect_gguf + + if model_id.lower().endswith(".gguf") and _os.path.isfile(model_id): + model_kwargs["gguf_file"] = _os.path.basename(model_id) + model_id = _os.path.dirname(model_id) + else: + _gguf = _maybe_auto_detect_gguf(model_id, revision=revision) + if _gguf is not None: + model_kwargs["gguf_file"] = _gguf + if progress_callback is not None: progress_callback({"status": "loading", "model": model_id_and_revision, "stage": "config"}) config = AutoConfig.from_pretrained(model_id, **model_kwargs) @@ -271,6 +297,11 @@ def _load_model( return AutoModelForMultimodalLM.from_pretrained(model_id, **model_kwargs) + if not config.architectures: # GGUF configs often omit architectures; route by model_type + from transformers import AutoModelForCausalLM + + return AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs) + architecture = getattr(transformers, config.architectures[0]) return architecture.from_pretrained(model_id, **model_kwargs) diff --git a/src/transformers/gguf_conversion_ops.py b/src/transformers/gguf_conversion_ops.py new file mode 100644 index 000000000000..ffb71d155fdb --- /dev/null +++ b/src/transformers/gguf_conversion_ops.py @@ -0,0 +1,416 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING + +from .core_model_loading import ConversionOps + + +if TYPE_CHECKING: + import torch + + +def _single_input_target(input_dict, source_patterns, target_patterns): + """Return the output key for a single-input conversion op.""" + if len(input_dict) != 1: + raise ValueError("Undefined Operation encountered!") + if len(target_patterns) > 1: + if len(source_patterns) == 1: + return source_patterns[0] + raise ValueError("Undefined Operation encountered!") + target = target_patterns[0] + if r"\1" in target: + return next(iter(input_dict)) + return target + + +class GGUFDequantize(ConversionOps): + """ + Reads `quant_type` from each input :class:`GGUFQuantizedTensor` and + dequantizes the raw uint8 bytes to a floating-point `torch.Tensor` using + the pure-torch kernels in `integrations/gguf_dequant.py` + """ + + def __init__(self, hf_quantizer=None): + # Quantizer reference is only needed for the reverse op (re-quantize on save). + self.hf_quantizer = hf_quantizer + + def __deepcopy__(self, memo): + # `WeightConverter` deep-copies its op list during loading. The default + # deepcopy would walk `self.hf_quantizer` → `gguf_tensors` (which + # contains :class:`GGUFQuantizedTensor` subclasses that lack `__deepcopy__`). + # We don't need a deep copy of the quantizer anyway — share the reference. + return GGUFDequantize(self.hf_quantizer) + + def convert(self, input_dict, source_patterns, target_patterns, **kwargs): + from .integrations.gguf_dequant import GGUFQuantizedTensor, dequantize_gguf_tensor + + out = {} + for key, tensors in input_dict.items(): + tensors_list = tensors if isinstance(tensors, list) else [tensors] + processed = [ + dequantize_gguf_tensor(t, t.quant_type, device=t.device) if isinstance(t, GGUFQuantizedTensor) else t + for t in tensors_list + ] + out[key] = processed if isinstance(tensors, list) else processed[0] + return out + + @property + def reverse_op(self): + return GGUFQuantize(self.hf_quantizer) + + +class GGUFQuantize(ConversionOps): + """Reverse op of :class:`GGUFDequantize`. Re-packs a floating-point tensor + into GGUF block bytes via `gguf.quants.quantize`. Used on the save path + to round-trip a dequantized / on-the-fly-quantized model back to a `.gguf` + file. Mirrors :class:`~transformers.integrations.finegrained_fp8.Fp8Quantize`. + + Today gguf-py only ships a Python quantizer for `Q4_0` and `Q8_0`; other + quant types (K-quants, IQ4) are read-only upstream, so attempting to use + them here raises with a clear message. + """ + + def __init__(self, hf_quantizer=None): + self.hf_quantizer = hf_quantizer + + def __deepcopy__(self, memo): + return GGUFQuantize(self.hf_quantizer) + + # TODO support all quantization types. + _GGUF_PY_QUANTIZE_SUPPORTED = ("Q4_0", "Q8_0") + + def _quantize_one(self, key: str, value): + import gguf + import torch + + quant_type = getattr(getattr(self.hf_quantizer, "quantization_config", None), "quant_type", None) or "Q4_0" + if quant_type not in self._GGUF_PY_QUANTIZE_SUPPORTED: + raise ValueError( + f"On-the-fly GGUF quantize only supports {self._GGUF_PY_QUANTIZE_SUPPORTED} (gguf-py limit); " + f"got quant_type={quant_type!r}. Pick Q4_0 or Q8_0, or load an existing .gguf via `gguf_file=`." + ) + ggml_type = getattr(gguf.GGMLQuantizationType, quant_type) + # gguf-py expects a numpy array in fp32. + fp = value.detach().to("cpu", dtype=torch.float32).contiguous().numpy() + packed = gguf.quants.quantize(fp, ggml_type) + return {key: torch.from_numpy(packed.view("uint8").reshape(-1))} + + def convert(self, input_dict, **kwargs): + result: dict = {} + for key, value in input_dict.items(): + tensor = value[0] if isinstance(value, list) else value + result.update(self._quantize_one(key, tensor)) + return result + + @property + def reverse_op(self): + return GGUFDequantize(self.hf_quantizer) + + +class Unsqueeze(ConversionOps): + """Unsqueeze a tensor along `dim`.""" + + def __init__(self, dim: int = 0): + self.dim = dim + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + **kwargs, + ) -> dict[str, torch.Tensor]: + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + return {target_pattern: tensor.unsqueeze(self.dim)} + + @property + def reverse_op(self) -> ConversionOps: + return Squeeze(self.dim) + + +class Squeeze(ConversionOps): + """Squeeze a tensor along `dim`. Inverse of :class:`Unsqueeze`.""" + + def __init__(self, dim: int = 0): + self.dim = dim + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + **kwargs, + ) -> dict[str, torch.Tensor]: + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + return {target_pattern: tensor.squeeze(self.dim)} + + @property + def reverse_op(self) -> ConversionOps: + return Unsqueeze(self.dim) + + +class SubtractOne(ConversionOps): + """Subtract 1 from GGUF norm weights stored as `w + 1` (Gemma/Nemotron). + + The `-1` must run in fp32: `cast_to_fp16(w + 1) - 1` drops ~1 ULP near + `w = 1` and the norm weights end up ~5e-4 off vs HF. The loader's + `_materialize_copy` casts to the target dtype *before* the converter chain + runs, so this op alone would see fp16. `GGUFQuantizer` keeps these norms in + fp32 for the load (via `_keep_in_fp32_modules_strict`, keyed off the + `SubtractOne` converter targets), so by the time `convert` runs the tensor is + fp32 and the subtraction is exact. + """ + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + **kwargs, + ) -> dict[str, torch.Tensor]: + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + return {target_pattern: tensor - 1} + + @property + def reverse_op(self) -> ConversionOps: + return AddOne() + + +class AddOne(ConversionOps): + """Add 1 to a tensor. Inverse of :class:`SubtractOne`.""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + **kwargs, + ) -> dict[str, torch.Tensor]: + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + return {target_pattern: tensor + 1} + + @property + def reverse_op(self) -> ConversionOps: + return SubtractOne() + + +class LogNegate(ConversionOps): + """Apply `log(-tensor)` (used for GGUF Mamba SSM-A de-transform).""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + **kwargs, + ) -> dict[str, torch.Tensor]: + import torch + + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + return {target_pattern: torch.log(-tensor)} + + @property + def reverse_op(self) -> ConversionOps: + raise NotImplementedError("LogNegate is not easily reversible") + + +class ReversePermuteAttnQ(ConversionOps): + """Reverse Q-projection GGUF permutation. Reads `config.num_attention_heads` at convert time.""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + config=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + if not hasattr(config, "num_attention_heads"): + raise ValueError("Config does not have `num_attention_heads`, you can't use `ReversePermuteAttnQ`.") + num_heads = config.num_attention_heads + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + dim = tensor.shape[0] // num_heads // 2 + return { + target_pattern: tensor.reshape(num_heads, dim, 2, *tensor.shape[1:]).swapaxes(2, 1).reshape(tensor.shape) + } + + @property + def reverse_op(self) -> ConversionOps: + return PermuteAttnQ() + + +class PermuteAttnQ(ConversionOps): + """Forward Q-projection GGUF permutation — inverse of :class:`ReversePermuteAttnQ`, + exposed as its :attr:`reverse_op` so the conversion is invertible. Re-emits Q + weights in llama.cpp's expected layout; reads `config.num_attention_heads` at + convert time.""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + config=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + if not hasattr(config, "num_attention_heads"): + raise ValueError("Config does not have `num_attention_heads`, you can't use `ReversePermuteAttnQ`.") + num_heads = config.num_attention_heads + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + dim = tensor.shape[0] // num_heads // 2 + return { + target_pattern: tensor.reshape(num_heads, 2, dim, *tensor.shape[1:]).swapaxes(2, 1).reshape(tensor.shape) + } + + @property + def reverse_op(self) -> ConversionOps: + return ReversePermuteAttnQ() + + +class ReversePermuteAttnK(ConversionOps): + """Reverse K-projection GGUF permutation. Reads `config.num_attention_heads` and + `config.num_key_value_heads` at convert time.""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + config=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + if not hasattr(config, "num_attention_heads"): + raise ValueError("Config does not have `num_attention_heads`, you can't use `ReversePermuteAttnQ`.") + num_kv_heads = config.num_key_value_heads + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + dim = tensor.shape[0] // num_kv_heads // 2 + return { + target_pattern: tensor.reshape(num_kv_heads, dim, 2, *tensor.shape[1:]) + .swapaxes(2, 1) + .reshape(tensor.shape) + } + + @property + def reverse_op(self) -> ConversionOps: + return PermuteAttnK() + + +class PermuteAttnK(ConversionOps): + """Forward K-projection GGUF permutation — inverse of :class:`ReversePermuteAttnK`.""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + config=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + if not hasattr(config, "num_attention_heads"): + raise ValueError("Config does not have `num_attention_heads`, you can't use `ReversePermuteAttnQ`.") + num_kv_heads = config.num_key_value_heads + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + tensor = tensors[0] if isinstance(tensors, list) else tensors + dim = tensor.shape[0] // num_kv_heads // 2 + return { + target_pattern: tensor.reshape(num_kv_heads, 2, dim, *tensor.shape[1:]) + .swapaxes(2, 1) + .reshape(tensor.shape) + } + + @property + def reverse_op(self) -> ConversionOps: + return ReversePermuteAttnK() + + +class BloomReshapeQKVWeight(ConversionOps): + """Reverse Bloom QKV weight reshape. Reads `config.n_head` and `config.hidden_size` at convert time.""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + config=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + import torch + + if not hasattr(config, "n_head") or not hasattr(config, "hidden_size"): + raise ValueError("Config does not have `num_attention_heads`, you can't use `ReversePermuteAttnQ`.") + n_head = config.n_head + n_embed = config.hidden_size + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + w = tensors[0] if isinstance(tensors, list) else tensors + q, k, v = torch.chunk(w, 3, dim=0) + q = q.reshape(n_head, n_embed // n_head, n_embed) + k = k.reshape(n_head, n_embed // n_head, n_embed) + v = v.reshape(n_head, n_embed // n_head, n_embed) + qkv = torch.stack([q, k, v], dim=1) + return {target_pattern: qkv.reshape(n_head * 3 * (n_embed // n_head), n_embed)} + + @property + def reverse_op(self) -> ConversionOps: + raise NotImplementedError("BloomReshapeQKVWeight is one-way") + + +class BloomReshapeQKVBias(ConversionOps): + """Reverse Bloom QKV bias reshape. Reads `config.n_head` and `config.hidden_size` at convert time.""" + + def convert( + self, + input_dict: dict[str, torch.Tensor], + source_patterns: list[str], + target_patterns: list[str], + config=None, + **kwargs, + ) -> dict[str, torch.Tensor]: + import torch + + if not hasattr(config, "n_head") or not hasattr(config, "hidden_size"): + raise ValueError("Config does not have `num_attention_heads`, you can't use `ReversePermuteAttnQ`.") + n_head = config.n_head + n_embed = config.hidden_size + target_pattern = _single_input_target(input_dict, source_patterns, target_patterns) + tensors = next(iter(input_dict.values())) + w = tensors[0] if isinstance(tensors, list) else tensors + q, k, v = torch.chunk(w, 3) + q = q.reshape(n_head, n_embed // n_head) + k = k.reshape(n_head, n_embed // n_head) + v = v.reshape(n_head, n_embed // n_head) + return {target_pattern: torch.stack([q, k, v], dim=1).flatten()} + + @property + def reverse_op(self) -> ConversionOps: + raise NotImplementedError("BloomReshapeQKVBias is one-way") diff --git a/src/transformers/integrations/gguf_dequant.py b/src/transformers/integrations/gguf_dequant.py new file mode 100644 index 000000000000..40d140b3b6ac --- /dev/null +++ b/src/transformers/integrations/gguf_dequant.py @@ -0,0 +1,533 @@ +# Copyright 2026 The HuggingFace Inc. team and City96. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +""" +!!!! THIS ENTIRE CODE IS FULLY VENDORED !!!!!!!!!!!!! +SOURCES: +https://github.com/city96/ComfyUI-GGUF/blob/main/dequant.py +https://github.com/huggingface/diffusers/blob/main/src/diffusers/quantizers/gguf/utils.py + +The reference dequant in `gguf-py` is pure NumPy and operates on grouped +row slices, which is ~5–20x slower than the same logic expressed as +`torch` ops on a `uint8` view of the raw bytes (and also avoids the +`__array_finalize__` overhead from memmap slicing). The ops here run on +CPU or GPU unchanged — pass an already-on-device tensor as input. + +`GGUFQuantizedTensor` is a `torch.Tensor` subclass that carries the +`quant_type` metadata alongside the raw uint8 bytes, so the standard +loader can move it to the target device with `.to(device)` and the +`GGUFDequantize` op in the weight-conversion chain can dequant on-device +without any GGUF-specific hook in `core_model_loading`. +""" + +from __future__ import annotations + +import torch + + +# Re-exported lazily so `import transformers` does not require gguf installed. +QK_K = 256 +K_SCALE_SIZE = 12 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _to_uint32(x): + import torch + + x = x.view(torch.uint8).to(torch.int32) + return (x[:, 0] | x[:, 1] << 8 | x[:, 2] << 16 | x[:, 3] << 24).unsqueeze(1) + + +def _split_block_dims(blocks, *args): + import torch + + n_max = blocks.shape[1] + dims = list(args) + [n_max - sum(args)] + return torch.split(blocks, dims, dim=1) + + +def _get_scale_min(scales): + import torch + + n_blocks = scales.shape[0] + scales = scales.view(torch.uint8).reshape((n_blocks, 3, 4)) + d, m, m_d = torch.split(scales, scales.shape[-2] // 3, dim=-2) + sc = torch.cat([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], dim=-1) + mn = torch.cat([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], dim=-1) + return sc.reshape((n_blocks, 8)), mn.reshape((n_blocks, 8)) + + +# --------------------------------------------------------------------------- +# Per-format kernels +# --------------------------------------------------------------------------- + + +def _dq_Q8_0(blocks, block_size, type_size, dtype=None): + import torch + + d, x = _split_block_dims(blocks, 2) + d = d.view(torch.float16).to(dtype) + x = x.view(torch.int8) + return d * x + + +def _dq_Q5_1(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + d, m, qh, qs = _split_block_dims(blocks, 2, 2, 4) + d = d.view(torch.float16).to(dtype) + m = m.view(torch.float16).to(dtype) + qh = _to_uint32(qh) + qh = qh.reshape((n_blocks, 1)) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) + ql = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor( + [0, 4], device=d.device, dtype=torch.uint8 + ).reshape(1, 1, 2, 1) + qh = (qh & 1).to(torch.uint8) + ql = (ql & 0x0F).reshape((n_blocks, -1)) + qs = ql | (qh << 4) + return (d * qs) + m + + +def _dq_Q5_0(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + d, qh, qs = _split_block_dims(blocks, 2, 4) + d = d.view(torch.float16).to(dtype) + qh = _to_uint32(qh) + qh = qh.reshape(n_blocks, 1) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) + ql = qs.reshape(n_blocks, -1, 1, block_size // 2) >> torch.tensor( + [0, 4], device=d.device, dtype=torch.uint8 + ).reshape(1, 1, 2, 1) + qh = (qh & 1).to(torch.uint8) + ql = (ql & 0x0F).reshape(n_blocks, -1) + qs = (ql | (qh << 4)).to(torch.int8) - 16 + return d * qs + + +def _dq_Q4_1(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + d, m, qs = _split_block_dims(blocks, 2, 2) + d = d.view(torch.float16).to(dtype) + m = m.view(torch.float16).to(dtype) + qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor( + [0, 4], device=d.device, dtype=torch.uint8 + ).reshape(1, 1, 2, 1) + qs = (qs & 0x0F).reshape(n_blocks, -1) + return (d * qs) + m + + +def _dq_Q4_0(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + d, qs = _split_block_dims(blocks, 2) + d = d.view(torch.float16).to(dtype) + qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor( + [0, 4], device=d.device, dtype=torch.uint8 + ).reshape((1, 1, 2, 1)) + qs = (qs & 0x0F).reshape((n_blocks, -1)).to(torch.int8) - 8 + return d * qs + + +def _dq_Q6_K(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + ql, qh, scales, d = _split_block_dims(blocks, QK_K // 2, QK_K // 4, QK_K // 16) + scales = scales.view(torch.int8).to(dtype) + d = d.view(torch.float16).to(dtype) + d = (d * scales).reshape((n_blocks, QK_K // 16, 1)) + ql = ql.reshape((n_blocks, -1, 1, 64)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( + (1, 1, 2, 1) + ) + ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) + qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape( + (1, 1, 4, 1) + ) + qh = (qh & 0x03).reshape((n_blocks, -1, 32)) + q = (ql | (qh << 4)).to(torch.int8) - 32 + q = q.reshape((n_blocks, QK_K // 16, -1)) + return (d * q).reshape((n_blocks, QK_K)) + + +def _dq_Q5_K(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + d, dmin, scales, qh, qs = _split_block_dims(blocks, 2, 2, K_SCALE_SIZE, QK_K // 8) + d = d.view(torch.float16).to(dtype) + dmin = dmin.view(torch.float16).to(dtype) + sc, m = _get_scale_min(scales) + d = (d * sc).reshape((n_blocks, -1, 1)) + dm = (dmin * m).reshape((n_blocks, -1, 1)) + ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( + (1, 1, 2, 1) + ) + qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.arange(0, 8, device=d.device, dtype=torch.uint8).reshape( + (1, 1, 8, 1) + ) + ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) + qh = (qh & 0x01).reshape((n_blocks, -1, 32)) + q = ql | (qh << 4) + return (d * q - dm).reshape((n_blocks, QK_K)) + + +def _dq_Q4_K(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + d, dmin, scales, qs = _split_block_dims(blocks, 2, 2, K_SCALE_SIZE) + d = d.view(torch.float16).to(dtype) + dmin = dmin.view(torch.float16).to(dtype) + sc, m = _get_scale_min(scales) + d = (d * sc).reshape((n_blocks, -1, 1)) + dm = (dmin * m).reshape((n_blocks, -1, 1)) + qs = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( + (1, 1, 2, 1) + ) + qs = (qs & 0x0F).reshape((n_blocks, -1, 32)) + return (d * qs - dm).reshape((n_blocks, QK_K)) + + +def _dq_Q3_K(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + hmask, qs, scales, d = _split_block_dims(blocks, QK_K // 8, QK_K // 4, 12) + d = d.view(torch.float16).to(dtype) + lscales, hscales = scales[:, :8], scales[:, 8:] + lscales = lscales.reshape((n_blocks, 1, 8)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape( + (1, 2, 1) + ) + lscales = lscales.reshape((n_blocks, 16)) + hscales = hscales.reshape((n_blocks, 1, 4)) >> torch.tensor( + [0, 2, 4, 6], device=d.device, dtype=torch.uint8 + ).reshape((1, 4, 1)) + hscales = hscales.reshape((n_blocks, 16)) + scales = (lscales & 0x0F) | ((hscales & 0x03) << 4) + scales = scales.to(torch.int8) - 32 + dl = (d * scales).reshape((n_blocks, 16, 1)) + ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape( + (1, 1, 4, 1) + ) + qh = hmask.reshape(n_blocks, -1, 1, 32) >> torch.arange(0, 8, device=d.device, dtype=torch.uint8).reshape( + (1, 1, 8, 1) + ) + ql = ql.reshape((n_blocks, 16, QK_K // 16)) & 3 + qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & 1) ^ 1 + q = ql.to(torch.int8) - (qh << 2).to(torch.int8) + return (dl * q).reshape((n_blocks, QK_K)) + + +def _dq_Q2_K(blocks, block_size, type_size, dtype=None): + import torch + + n_blocks = blocks.shape[0] + scales, qs, d, dmin = _split_block_dims(blocks, QK_K // 16, QK_K // 4, 2) + d = d.view(torch.float16).to(dtype) + dmin = dmin.view(torch.float16).to(dtype) + dl = (d * (scales & 0xF)).reshape((n_blocks, QK_K // 16, 1)) + ml = (dmin * (scales >> 4)).reshape((n_blocks, QK_K // 16, 1)) + shift = torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) + qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & 3 + qs = qs.reshape((n_blocks, QK_K // 16, 16)) + qs = dl * qs - ml + return qs.reshape((n_blocks, -1)) + + +def _dq_BF16(blocks, block_size, type_size, dtype=None): + import torch + + return (blocks.view(torch.int16).to(torch.int32) << 16).view(torch.float32) + + +_IQ4_KVALUES = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113) + + +def _dq_IQ4_NL(blocks, block_size, type_size, dtype=None): + import torch + + kvalues = torch.tensor(_IQ4_KVALUES, dtype=torch.float32, device=blocks.device) + n_blocks = blocks.shape[0] + d, qs = _split_block_dims(blocks, 2) + d = d.view(torch.float16).to(dtype) + qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor( + [0, 4], device=blocks.device, dtype=torch.uint8 + ).reshape((1, 1, 2, 1)) + qs = (qs & 15).reshape((n_blocks, -1)).to(torch.int64) + kvalues = kvalues.view(1, 1, 16) + qs = qs.unsqueeze(-1) + qs = torch.gather(kvalues.expand(qs.shape[0], qs.shape[1], 16), 2, qs) + qs = qs.squeeze(-1).to(dtype) + return d * qs + + +def _dq_IQ4_XS(blocks, block_size, type_size, dtype=None): + import torch + + kvalues = torch.tensor(_IQ4_KVALUES, dtype=torch.float32, device=blocks.device) + n_blocks = blocks.shape[0] + d, scales_h, scales_l, qs = _split_block_dims(blocks, 2, 2, QK_K // 64) + d = d.view(torch.float16).to(dtype) + scales_h = scales_h.view(torch.int16) + scales_l = scales_l.reshape((n_blocks, -1, 1)) >> torch.tensor( + [0, 4], device=blocks.device, dtype=torch.uint8 + ).reshape((1, 1, 2)) + scales_h = scales_h.reshape((n_blocks, 1, -1)) >> torch.tensor( + [2 * i for i in range(QK_K // 32)], device=blocks.device, dtype=torch.uint8 + ).reshape((1, -1, 1)) + scales_l = scales_l.reshape((n_blocks, -1)) & 0x0F + scales_h = scales_h.reshape((n_blocks, -1)) & 0x03 + scales = (scales_l | (scales_h << 4)) - 32 + dl = (d * scales.to(dtype)).reshape((n_blocks, -1, 1)) + shifts_q = torch.tensor([0, 4], device=blocks.device, dtype=torch.uint8).reshape(1, 1, 2, 1) + qs = qs.reshape((n_blocks, -1, 1, 16)) >> shifts_q + qs = (qs & 15).reshape((n_blocks, -1, 32)).to(torch.int64) + kvalues = kvalues.view(1, 1, 1, 16) + qs = qs.unsqueeze(-1) + qs = torch.gather(kvalues.expand(qs.shape[0], qs.shape[1], qs.shape[2], 16), 3, qs) + qs = qs.squeeze(-1).to(dtype) + return (dl * qs).reshape(n_blocks, -1) + + +def _build_dispatch(): + """Build the {GGMLQuantizationType: kernel} dispatch table lazily so that + importing this module does not require `gguf` to be installed.""" + import gguf + + return { + gguf.GGMLQuantizationType.IQ4_NL: _dq_IQ4_NL, + gguf.GGMLQuantizationType.IQ4_XS: _dq_IQ4_XS, + gguf.GGMLQuantizationType.BF16: _dq_BF16, + gguf.GGMLQuantizationType.Q8_0: _dq_Q8_0, + gguf.GGMLQuantizationType.Q5_1: _dq_Q5_1, + gguf.GGMLQuantizationType.Q5_0: _dq_Q5_0, + gguf.GGMLQuantizationType.Q4_1: _dq_Q4_1, + gguf.GGMLQuantizationType.Q4_0: _dq_Q4_0, + gguf.GGMLQuantizationType.Q6_K: _dq_Q6_K, + gguf.GGMLQuantizationType.Q5_K: _dq_Q5_K, + gguf.GGMLQuantizationType.Q4_K: _dq_Q4_K, + gguf.GGMLQuantizationType.Q3_K: _dq_Q3_K, + gguf.GGMLQuantizationType.Q2_K: _dq_Q2_K, + } + + +_DISPATCH = None + + +def supported_quant_types(): + global _DISPATCH + if _DISPATCH is None: + _DISPATCH = _build_dispatch() + return set(_DISPATCH.keys()) + + +class GGUFQuantizedTensor(torch.Tensor): + """`torch.Tensor` subclass that carries the GGUF `quant_type` alongside + raw uint8 bytes. + + Three small affordances make the hot path through `spawn_materialize` fast + *without* touching `core_model_loading`: + + * `__getitem__(Ellipsis)` short-circuits — `tensor[...]` is a no-op for + already-loaded torch bytes; the default goes through `__torch_function__` + dispatch for no reason. + * `.to(...)` defaults `non_blocking=True`, letting the loader queue the + next transfer / dequant on the same MPS/CUDA stream while the previous + copy is still in flight. + * `__torch_function__` only re-wraps on `Tensor.to` — the one place we + need `quant_type` to survive (so the :class:`GGUFDequantize` op in the + conversion chain can read it on the device side). All other ops return + plain tensors, which avoids the per-op wrap overhead. + + Inspired by `GGUFParameter` in diffusers. + """ + + # Class-level default so subclass instances spawned by torch's default + # `__torch_function__` path (which doesn't invoke our `__new__`) still + # have the attribute defined. + quant_type = None + + @staticmethod + def __new__(cls, data, quant_type=None): + data = data if data is not None else torch.empty(0) + instance = torch.Tensor._make_subclass(cls, data, require_grad=False) + instance.quant_type = quant_type + return instance + + def __getitem__(self, key): + # `_materialize_copy` does `tensor = tensor[...]` to pull a memmap + # safetensors slice into RAM. Our bytes are already a torch.Tensor view, + # so this is a no-op; short-circuit before torch dispatches via + # `__torch_function__` (which would re-wrap into a fresh subclass). + if key is Ellipsis: + return self + return super().__getitem__(key) + + def to(self, *args, **kwargs): + # The loader queues a dequant op on the destination device right after + # this transfer (same MPS/CUDA stream), so `non_blocking=True` overlaps + # the bytes copy with the next `.to(device)` call from another worker. + # Skip the inject if `non_blocking` was already supplied positionally + # (3rd positional after device + dtype) — torch's `nn.Module._apply` + # calls `Tensor.to` with positional args on some torch versions and + # would otherwise raise a duplicate-kwarg TypeError. + if "non_blocking" not in kwargs and len(args) < 3: + kwargs["non_blocking"] = True + return super().to(*args, **kwargs) + + @staticmethod + def _extract_quant_type(args): + for arg in args: + if isinstance(arg, GGUFQuantizedTensor) and arg.quant_type is not None: + return arg.quant_type + if isinstance(arg, (list, tuple)) and arg and isinstance(arg[0], GGUFQuantizedTensor): + if arg[0].quant_type is not None: + return arg[0].quant_type + return None + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + result = super().__torch_function__(func, types, args, kwargs) + # Carry `quant_type` through every subclass-preserving op so it survives the + # byte-manipulation chain that builds the swapped buffers — the loader's + # `.to` move, the Q/K permute (reshape/swapaxes), and the gate/up merge + # (`torch.cat`) — and `GgufLinear`/`GgufExperts.bind_after_load` can read it + # back. These subclasses only exist during loading (buffers are unwrapped to + # plain tensors at bind time), so this never runs on the forward hot path. + quant_type = cls._extract_quant_type(args) + if quant_type is None: + return result + if isinstance(result, cls): + if getattr(result, "quant_type", None) is None: + result.quant_type = quant_type + return result + # `.to` may return a plain tensor (e.g. device-only move); keep it tagged so + # the downstream dequant op can still read the type. + if func is torch.Tensor.to and isinstance(result, torch.Tensor): + return cls(result, quant_type=quant_type) + return result + + +def dequantize_gguf_tensor(data, quant_type, dtype=None, device=None) -> torch.Tensor: + """Dequantize a GGUF tensor to a `torch.Tensor` using torch ops. + + Args: + data: the `ReaderTensor.data` numpy array from `gguf.GGUFReader`. + For quantized types this is shaped as the **byte** shape + (rows × byte-cols, uint8); for F16/F32 it is already the + logical shape with the matching dtype. + quant_type: `gguf.GGMLQuantizationType` enum value. + dtype: target floating-point dtype for the dequantized output + (defaults to `torch.float32`). + device: device to run the dequant kernel on. The raw uint8 input is + transferred to `device` *before* the kernel runs so the big + float output is produced on-device — this is dramatically faster + on MPS / CUDA than dequantizing on CPU and copying the + fully-expanded tensor afterwards. + + Returns: + Tensor with the **logical** shape recovered from the byte shape + (matching what `gguf.dequantize` returns). + + Raises: + NotImplementedError: when `quant_type` has no torch kernel. + """ + import gguf + import torch + + global _DISPATCH + if _DISPATCH is None: + _DISPATCH = _build_dispatch() + + if dtype is None: + dtype = torch.float32 + target_device = torch.device(device) if device is not None else None + + # Dequant output must be a *plain* `torch.Tensor`, never a `GGUFQuantizedTensor`. + # If the subclass leaks through, the fp32 dequant values still carry `quant_type` + # and the subclass's `.to` / `__torch_function__` override re-wraps on the + # loader's device+dtype cast — on MPS this reinterprets the storage as the target + # dtype (raw bytes → fp16 garbage like `28512.0`; fp32 happens to survive). Unwrap + # up-front so every downstream view/reshape/kernel op yields base tensors. + if isinstance(data, GGUFQuantizedTensor): + data = data.as_subclass(torch.Tensor) + + # F32 / F16 / BF16 sources arrive as regular floating-point tensors when + # wrapped via `GGUFQuantizedTensor(torch.from_numpy(...))` — no byte-view + # round-trip needed. The uint8 reinterpret path below is destructive on + # MPS for non-quantized fp tensors. + if isinstance(data, torch.Tensor) and data.is_floating_point(): + out = data.to(dtype=dtype) if dtype is not None and data.dtype != dtype else data + if target_device is not None and out.device != target_device: + out = out.to(target_device, non_blocking=True) + return out + + # Accept either a numpy array (gguf reader) or a torch.Tensor (e.g. a + # `GGUFQuantizedTensor` already moved to the target device). For numpy + # we materialise the mmap view into RAM and wrap zero-copy; for torch we + # just take a uint8 view of the raw bytes. + if isinstance(data, torch.Tensor): + bytes_view = data.contiguous().view(torch.uint8) + byte_shape = tuple(bytes_view.shape) + flat = bytes_view.reshape(-1) + else: + import numpy as np + + arr = np.ascontiguousarray(data) + byte_shape = arr.shape + flat = torch.from_numpy(arr.reshape(-1)) + + if target_device is not None and flat.device != target_device: + flat = flat.to(target_device, non_blocking=True) + + # Already-float types coming in as raw uint8 bytes (e.g. from `numpy` + # path above): reinterpret as the matching float dtype. + if quant_type == gguf.GGMLQuantizationType.F32: + out = flat.view(torch.float32).reshape(byte_shape[:-1] + (byte_shape[-1] // 4,)) + return out.to(dtype) if dtype != torch.float32 else out + if quant_type == gguf.GGMLQuantizationType.F16: + out = flat.view(torch.float16).reshape(byte_shape[:-1] + (byte_shape[-1] // 2,)) + return out.to(dtype) if dtype != torch.float16 else out + + block_size, type_size = gguf.GGML_QUANT_SIZES[quant_type] + # logical_shape recovers the per-element shape from the byte shape (last + # dim shrinks from `n_blocks * type_size` bytes → `n_blocks * block_size` elements). + logical_shape = (*byte_shape[:-1], byte_shape[-1] // type_size * block_size) + + kernel = _DISPATCH.get(quant_type) + if kernel is None: + # Fall back to gguf-py's numpy dequant for types without a fast torch + # kernel (IQ1_S/M, IQ2_*, IQ3_*). These are uncommon, CPU-side, and + # avoid us having to port every llama.cpp i-quant by hand. + import numpy as np + + cpu_bytes = flat.detach().to("cpu", non_blocking=True).numpy() + dequanted = gguf.quants.dequantize(cpu_bytes.reshape(-1, type_size), quant_type) + out = torch.from_numpy(np.ascontiguousarray(dequanted)).to(dtype=dtype) + if target_device is not None and out.device != target_device: + out = out.to(target_device, non_blocking=True) + return out.reshape(logical_shape) + + blocks = flat.reshape((-1, type_size)) + out = kernel(blocks, block_size, type_size, dtype=dtype) + return out.reshape(logical_shape) diff --git a/src/transformers/integrations/gguf_kernels.py b/src/transformers/integrations/gguf_kernels.py new file mode 100644 index 000000000000..3eb5a5f030fb --- /dev/null +++ b/src/transformers/integrations/gguf_kernels.py @@ -0,0 +1,64 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + + +_GGUF_METAL_KERNELS = None +_GGUF_METAL_LOADED = False + + +# Quant-type → kernel-name suffix table. Kernel names like +# `mul_mat__f32` / `mul_mat_vec__f32` / `mul_mat_id__f32` +# are deterministic from this, so consumers can build the full name once at +# construction. +_KERNEL_FMT: dict[str, str] = { + "Q4_0": "q4_0", + "Q5_0": "q5_0", + "Q5_1": "q5_1", + "Q8_0": "q8_0", + "Q4_K": "q4_K", + "Q5_K": "q5_K", + "Q6_K": "q6_K", + "IQ4_NL": "iq4_nl", + "IQ4_XS": "iq4_xs", +} + + +# TODO Change the repo name to kernel community once merged. +def ensure_metal_kernels(repo: str = "ArthurZ/gguf-kernels"): + """Return the loaded kernels handle, caching across calls. Raises + `RuntimeError` on failure — no fallback.""" + global _GGUF_METAL_KERNELS, _GGUF_METAL_LOADED + if _GGUF_METAL_LOADED: + if _GGUF_METAL_KERNELS is None: + raise RuntimeError( + f"GGUF metal kernels failed to load earlier in this process. Reinstall / pre-fetch {repo!r} and retry." + ) + return _GGUF_METAL_KERNELS + _GGUF_METAL_LOADED = True + try: + from kernels import get_kernel + except ImportError as exc: + raise RuntimeError( + "The GGUF fast path requires the `kernels` package. Install it with `pip install kernels`." + ) from exc + try: + _GGUF_METAL_KERNELS = get_kernel(repo) + except Exception as exc: + raise RuntimeError( + f"Could not load GGUF metal kernels from {repo!r}: {exc!r}. The kernel repo must be " + f"registered with the kernels-package backend (typically under kernels-community/); " + f"once migrated this path Just Works." + ) from exc + return _GGUF_METAL_KERNELS diff --git a/src/transformers/integrations/gguf_linear.py b/src/transformers/integrations/gguf_linear.py new file mode 100644 index 000000000000..d67716688f62 --- /dev/null +++ b/src/transformers/integrations/gguf_linear.py @@ -0,0 +1,423 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import re + +import torch +import torch.nn as nn + +from .gguf_kernels import _KERNEL_FMT, ensure_metal_kernels +from .moe import ExpertsInterface + + +# Per quant type: (block bytes, block elems). The format string side of the +# table (kernel names like mul_mat_q4_K_f32) lives in gguf_kernels._KERNEL_FMT. +_QUANT_INFO: dict[str, tuple[int, int]] = { + "Q4_0": (18, 32), + "Q5_0": (22, 32), + "Q5_1": (24, 32), + "Q8_0": (34, 32), + "Q4_K": (144, 256), + "Q5_K": (176, 256), + "Q6_K": (210, 256), + "IQ4_NL": (18, 32), + "IQ4_XS": (136, 256), +} + + +# Q4_K_M averages ~4.5 bits/weight. Used to size the meta-time placeholder +# buffers when the exact quant type isn't known yet (uniform swap): the loader +# replaces them with the real-shaped bytes at load, and this keeps the +# device_map memory estimate (numel × 1 byte for a uint8 buffer) in the right +# ballpark without needing per-tensor quant types up front. +_AVG_BYTES_PER_ELEM = 4.5 / 8 + + +def _quant_type_name(quant_type) -> str | None: + if quant_type is None: + return None + return quant_type.name if hasattr(quant_type, "name") else str(quant_type) + + +def gguf_linear_supports(quant_type) -> bool: + """Return True if quant_type has a matching mul_mat__f32 kernel.""" + return _quant_type_name(quant_type) in _QUANT_INFO + + +def _bytes_per_row(in_features: int, quant_type: str | None) -> int: + """Row byte-width for a GGUF-quantized weight. Exact when the quant type is + known; otherwise the Q4_K_M average (placeholder for the meta-time uniform + swap, corrected to the real width when the bytes load).""" + if quant_type in _QUANT_INFO: + block_bytes, block_elems = _QUANT_INFO[quant_type] + return (in_features // block_elems) * block_bytes + return max(1, round(in_features * _AVG_BYTES_PER_ELEM)) + + +class GgufLinear(nn.Module): + """Linear layer with GGUF-quantized weights stored as raw block bytes (uint8). + + Forward dispatches to mul_mat_vec__f32 for batch 1 and mul_mat__f32 + for batched inputs. Inference-only — no dequant fallback. Callers that want + a CPU/CUDA dequant path should let the loader's GGUFDequantize op produce + fp32 weights into a standard nn.Linear instead. + """ + + def __init__(self, in_features: int, out_features: int, quant_type: str | None = None, bias: bool = False): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.quant_type = quant_type + # 2D `(out_features, bytes_per_row)` — matches the source byte shape coming + # out of the GGUF reader, so the standard loader assigns it in without any + # reshape. `quant_type` may be `None` here: the uniform meta-time swap + # creates these placeholders without knowing the per-tensor type, sizes the + # buffer from the Q4_K_M average, and the loader overwrites it with the + # real-shaped bytes. Kernels + exact quant type are bound post-load via + # `bind_after_load`. + self.register_buffer( + "weight", torch.empty(out_features, _bytes_per_row(in_features, quant_type), dtype=torch.uint8) + ) + if bias: + self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.float32), requires_grad=False) + else: + self.register_parameter("bias", None) + self._mv_op = self._mat_op = None + if quant_type is not None: + self._bind_kernels(quant_type) + + def _bind_kernels(self, quant_type: str) -> None: + """Resolve the metal matvec/matmul ops for `quant_type` and cache them so + the forward is branch-free (compile-friendly). `ensure_metal_kernels()` + raises a clear RuntimeError if the kernels can't load — no fallback.""" + if quant_type not in _QUANT_INFO: + raise NotImplementedError(f"GgufLinear only supports {sorted(_QUANT_INFO)} today, got {quant_type}") + block_elems = _QUANT_INFO[quant_type][1] + if self.in_features % block_elems != 0: + raise ValueError( + f"in_features must be a multiple of {block_elems} for {quant_type}, got {self.in_features}" + ) + self.quant_type = quant_type + ops = ensure_metal_kernels()._ops + fmt = _KERNEL_FMT[quant_type] + self._mv_op = getattr(ops, f"mul_mat_vec_{fmt}_f32").default + self._mat_op = getattr(ops, f"mul_mat_{fmt}_f32").default + + def bind_after_load(self, quant_type: str | None = None) -> None: + """Post-load: bind kernels once the real bytes are in `weight`. `quant_type` + comes from the saved plan on a safetensors reload; otherwise it's read off + the loaded buffer (a `GGUFQuantizedTensor` carries it). The buffer is + unwrapped to a plain uint8 tensor so the subclass stays out of the + compiled forward.""" + if quant_type is None: + quant_type = _quant_type_name(getattr(self.weight, "quant_type", None)) + if type(self.weight) is not torch.Tensor: + self.weight = self.weight.as_subclass(torch.Tensor) + self._bind_kernels(quant_type) + + def extra_repr(self) -> str: + return ( + f"in_features={self.in_features}, out_features={self.out_features}, " + f"quant_type={self.quant_type}, bias={self.bias is not None}" + ) + + @torch.no_grad() + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[-1] != self.in_features: + raise ValueError(f"Last dim of x must be {self.in_features}, got {x.shape[-1]}") + if x.device.type != "mps": + raise RuntimeError( + "GgufLinear runs only on MPS with the metal kernels loaded. Re-load with " + "`dtype=torch.bfloat16` (or any explicit dtype) to dequantize to a normal nn.Linear." + ) + + batch_shape = x.shape[:-1] + x_flat = x.reshape(-1, self.in_features).contiguous().to(torch.float32) + N = x_flat.shape[0] + # Buffer is 2D `(out, bytes_per_row)` to match the GGUF source layout; + # kernels want a flat byte array, so view here at the call boundary. + qw = self.weight.view(-1) + # Pick the kernel by batch size: matvec when N==1, matmul otherwise. The + # matmul kernel requires N % 32 == 0; non-aligned batches loop per-row + # over matvec (rare in practice — most call sites land aligned). + if N == 1: + y = torch.empty(self.out_features, dtype=torch.float32, device=x.device) + self._mv_op(qw, x_flat.view(-1), y) + elif N % 32 == 0: + y = torch.empty(N * self.out_features, dtype=torch.float32, device=x.device) + self._mat_op(qw, x_flat, y) + else: + y = torch.empty(N * self.out_features, dtype=torch.float32, device=x.device) + for i in range(N): + self._mv_op(qw, x_flat[i].reshape(-1), y[i * self.out_features : (i + 1) * self.out_features]) + y = y.view(N, self.out_features).reshape(*batch_shape, self.out_features) + if self.bias is not None: + y = y + self.bias.to(y.device).to(y.dtype) + # Metal matvec / matmul emit fp32. Cast back to the caller's dtype so the + # KV cache stays in the model's chosen precision instead of silently widening. + if y.dtype != x.dtype: + y = y.to(x.dtype) + return y + + +class GgufExperts(nn.Module): + """Drop-in for fused-expert MoE modules with the Mixtral / Qwen2MoE / + DeepSeek-V3 layout (`gate_up_proj` shape `(num_experts, 2*intermediate, + hidden)`, `down_proj` shape `(num_experts, hidden, intermediate)`, + no per-expert biases), but storing raw uint8 GGUF bytes instead of fp + weights. Mixed quant types per layer are fine — Q4_K_M ships gate/up = + Q4_K, down = Q8_0. + + Archs with a different layout (transposed params, biases — gpt_oss) are + not swapped to this class; `GGUFDequantize` lands bf16 into their + original fused-expert module and `moe.py`'s `batched_mm_experts_forward` + handles them via its existing `is_transposed` / `has_bias` flags.""" + + # ExpertsInterface contract flags. + has_gate = True + has_bias = False + is_transposed = False + is_concatenated = True + + def __init__(self, config, gate_up_quant: str | None = None, down_quant: str | None = None): + super().__init__() + self.config = config + self.num_experts = config.num_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.moe_intermediate_size + self.gate_up_quant, self.down_quant = gate_up_quant, down_quant + + # uint8 byte buffers, sized from the per-tensor quant type when known and + # from the Q4_K_M average otherwise (uniform meta-time swap); the loader + # overwrites them with the real-shaped bytes. Kernels bind post-load. + self.register_buffer( + "gate_up_proj", + torch.empty( + self.num_experts, + 2 * self.intermediate_dim, + _bytes_per_row(self.hidden_dim, gate_up_quant), + dtype=torch.uint8, + ), + ) + self.register_buffer( + "down_proj", + torch.empty( + self.num_experts, + self.hidden_dim, + _bytes_per_row(self.intermediate_dim, down_quant), + dtype=torch.uint8, + ), + ) + + from ..activations import ACT2FN + + self.act_fn = ACT2FN[config.hidden_act] + self._id_op_gate_up = self._id_op_down = None + if gate_up_quant is not None and down_quant is not None: + self._bind_kernels(gate_up_quant, down_quant) + + decode_size = getattr(config, "num_experts_per_tok", None) or 4 + # Pre-allocated scratch for the decode hot path. At decode there is a single token, so the + # kernels always see the same `(num_experts_per_tok, ...)` shapes; allocating the gate/up/out + # buffers once at init (instead of `torch.empty` inside `forward`) keeps the forward graph + # static so torch.compile / dynamo does not re-trace on every step. Prefill (variable token + # counts) falls back to fresh allocations — see `gguf_bmm_experts_forward`. + self.register_buffer( + "_gate_buf", torch.empty(decode_size, self.intermediate_dim, dtype=torch.float32), persistent=False + ) + self.register_buffer( + "_up_buf", torch.empty(decode_size, self.intermediate_dim, dtype=torch.float32), persistent=False + ) + self.register_buffer( + "_pair_out", torch.empty(decode_size, self.hidden_dim, dtype=torch.float32), persistent=False + ) + self._decode_size = decode_size + + def _bind_kernels(self, gate_up_quant: str, down_quant: str) -> None: + """Resolve and cache the `mul_mat_id` ops for the gate/up and down quant + types so the forward is branch-free.""" + for label, qt in (("gate_up", gate_up_quant), ("down", down_quant)): + if qt not in _QUANT_INFO: + raise NotImplementedError(f"GgufExperts {label}_quant {qt} not in {sorted(_QUANT_INFO)}") + self.gate_up_quant, self.down_quant = gate_up_quant, down_quant + ops = ensure_metal_kernels()._ops + self._id_op_gate_up = getattr(ops, f"mul_mat_id_{_KERNEL_FMT[gate_up_quant]}_f32").default + self._id_op_down = getattr(ops, f"mul_mat_id_{_KERNEL_FMT[down_quant]}_f32").default + + def bind_after_load(self, gate_up_quant: str | None = None, down_quant: str | None = None) -> None: + """Post-load: bind kernels from the saved plan (reload) or the loaded + buffers' carried quant types, and unwrap the byte buffers to plain uint8.""" + if gate_up_quant is None: + gate_up_quant = _quant_type_name(getattr(self.gate_up_proj, "quant_type", None)) + if down_quant is None: + down_quant = _quant_type_name(getattr(self.down_proj, "quant_type", None)) + if type(self.gate_up_proj) is not torch.Tensor: + self.gate_up_proj = self.gate_up_proj.as_subclass(torch.Tensor) + if type(self.down_proj) is not torch.Tensor: + self.down_proj = self.down_proj.as_subclass(torch.Tensor) + self._bind_kernels(gate_up_quant, down_quant) + + def extra_repr(self) -> str: + return ( + f"num_experts={self.num_experts}, hidden_dim={self.hidden_dim}, " + f"intermediate_dim={self.intermediate_dim}, " + f"gate_up={self.gate_up_quant}, down={self.down_quant}" + ) + + @torch.no_grad() + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + """Dispatch through ALL_GGUF_EXPERTS_FUNCTIONS (GGUF-specific sibling + of ALL_EXPERTS_FUNCTIONS). Default is the bmm path.""" + impl = getattr(self, "_experts_implementation", None) or "bmm" + return ALL_GGUF_EXPERTS_FUNCTIONS.get_interface(impl, gguf_bmm_experts_forward)( + self, hidden_states, top_k_index, top_k_weights + ) + + +MODEL_TYPE_TO_GGUF_EXPERTS: dict[str, type[nn.Module]] = { + # Same layout as MixtralExperts: `gate_up_proj` shape + # `(num_experts, 2*intermediate, hidden)`. + "qwen2_moe": GgufExperts, + "qwen3_moe": GgufExperts, + "minimax_m2": GgufExperts, + "mixtral": GgufExperts, + "deepseek_v3": GgufExperts, + # gpt_oss is intentionally NOT here. Its experts store `gate_up_proj` + # transposed — `(num_experts, hidden, 2*intermediate)` — and carry + # per-expert biases. The Metal `mul_mat_id` kernels we ship assume rows = + # output axis, so a byte-passthrough GgufExperts subclass would have a + # broken forward. The FP8 pattern handles archs without specialised kernels + # the same way: skip the swap, let the dequant op land bf16 into the + # original `GptOssExperts`, and reuse `moe.py`'s + # `batched_mm_experts_forward` (which already handles + # `is_transposed=True` + `has_bias=True`). Re-add gpt_oss here once a + # transposed-aware `mul_mat_id` kernel exists. +} + + +@torch.no_grad() +def gguf_bmm_experts_forward( + self: nn.Module, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, +) -> torch.Tensor: + num_top_k = top_k_index.size(-1) + num_tokens = hidden_states.size(0) + hidden_dim = self.hidden_dim + intermediate_dim = self.intermediate_dim + device = hidden_states.device + # S = total token-expert pairs across the batch (each routed token contributes + # num_top_k work items). The mul_mat_id kernel writes one output row per pair. + S = num_tokens * num_top_k + + selected = hidden_states.repeat_interleave(num_top_k, dim=0).to(torch.float32).contiguous() + ids32 = top_k_index.reshape(-1).to(torch.int32) + sample_weights = top_k_weights.reshape(-1).to(torch.float32) + + # Avoids `torch.empty` inside the compiled graph — without this, + # dynamo would re-trace on the second call + if num_tokens == 1 and self._decode_size == S: + # this was alloceted at init time + gate_buf, up_buf, out = self._gate_buf, self._up_buf, self._pair_out + else: + # this one is for prefill + gate_buf = torch.empty(S, intermediate_dim, dtype=torch.float32, device=device) + up_buf = torch.empty(S, intermediate_dim, dtype=torch.float32, device=device) + out = torch.empty(S, hidden_dim, dtype=torch.float32, device=device) + + # TODO: gate_up should be a single call passing gate_up (no slice, contiguous, view here) + gate_up_bytes = self.gate_up_proj.view(self.num_experts, -1) + half = gate_up_bytes.shape[1] // 2 + gate_qw = gate_up_bytes[:, :half].contiguous().view(-1) + up_qw = gate_up_bytes[:, half:].contiguous().view(-1) + down_qw = self.down_proj.view(-1) + self._id_op_gate_up(gate_qw, selected, ids32, gate_buf) + self._id_op_gate_up(up_qw, selected, ids32, up_buf) + inter = (self.act_fn(gate_buf) * up_buf).contiguous() + self._id_op_down(down_qw, inter, ids32, out) + + weighted = out * sample_weights.unsqueeze(-1) + return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype) + + +class GgufExpertsInterface(ExpertsInterface): + _global_mapping = { + "bmm": gguf_bmm_experts_forward, + } + + +ALL_GGUF_EXPERTS_FUNCTIONS = GgufExpertsInterface() + + +def replace_with_gguf_linear( + model: nn.Module, + modules_to_not_convert: set[str] | None = None, + modules_to_convert: list[str] | None = None, + quant_type: str | None = None, +) -> int: + """Uniformly swap every `nn.Linear` (and matching fused-expert module) to a + `GgufLinear` / `GgufExperts` placeholder at meta time — no source→target + rename. + + gguf_file / reload (`quant_type=None`): the per-module quant type is unknown + here; the buffer is sized from the Q4_K_M average and the exact type + metal + kernels are bound post-load from the bytes that arrive (`bind_after_load`). + Linears whose GGUF source turns out to be float (e.g. MoE routers) are + reverted to plain `nn.Linear` then. + + on-the-fly (`quant_type` set): a single config quant type applies to every + converted module, so it's bound at construction; `modules_to_convert` (glob + includes) restricts which Linears are quantized. + + `modules_to_not_convert` is a user/skip list; the GGUF file itself decides + what is quantized, so no default skip-list is applied.""" + import fnmatch + + skip_re = re.compile("|".join(re.escape(p) for p in modules_to_not_convert)) if modules_to_not_convert else None + experts_cls = MODEL_TYPE_TO_GGUF_EXPERTS.get(getattr(model.config, "model_type", None)) + swapped = 0 + for name, mod in list(model.named_modules()): + if skip_re is not None and skip_re.search(name): + continue + if modules_to_convert is not None and not any(fnmatch.fnmatchcase(name, pat) for pat in modules_to_convert): + continue + new_module: nn.Module | None = None + with torch.device("meta"): + if isinstance(mod, nn.Linear): + new_module = GgufLinear(mod.in_features, mod.out_features, quant_type=quant_type, bias=mod.bias is not None) + elif experts_cls is not None and hasattr(mod, "gate_up_proj") and hasattr(mod, "down_proj"): + new_module = experts_cls(config=getattr(mod, "config", model.config), gate_up_quant=quant_type, down_quant=quant_type) + if new_module is None: + continue + model.set_submodule(name, new_module) + swapped += 1 + return swapped + + +__all__: list[str] = [ + "ALL_GGUF_EXPERTS_FUNCTIONS", + "GgufLinear", + "GgufExperts", + "MODEL_TYPE_TO_GGUF_EXPERTS", + "gguf_bmm_experts_forward", + "gguf_linear_supports", + "replace_with_gguf_linear", +] diff --git a/src/transformers/modeling_gguf_pytorch_utils.py b/src/transformers/modeling_gguf_pytorch_utils.py index 2de6cc13fc85..a35f6acea92a 100644 --- a/src/transformers/modeling_gguf_pytorch_utils.py +++ b/src/transformers/modeling_gguf_pytorch_utils.py @@ -13,12 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import re -from typing import NamedTuple - -import numpy as np -from tqdm.auto import tqdm - from .integrations import ( GGUF_CONFIG_DEFAULTS_MAPPING, GGUF_CONFIG_MAPPING, @@ -30,9 +24,6 @@ from .utils.logging import get_logger -if is_torch_available(): - import torch - logger = get_logger(__name__) @@ -53,422 +44,427 @@ GGUF_SUPPORTED_ARCHITECTURES = list(GGUF_TO_TRANSFORMERS_MAPPING["config"].keys()) -class GGUFTensor(NamedTuple): - weights: np.ndarray - name: str - metadata: dict +# --- Llama-family (rope-permuted Q/K, plain Llama norms) -------------------- +# Constructing the rename rules pulls in torch via WeightRenaming/WeightConverter, +# so we keep that behind `is_torch_available()` — `load_gguf_checkpoint` is +# imported by `tokenization_utils_tokenizers` even in tokenizer-only installs. +_GGUF_ARCH_CONVERTERS: dict[str, list] = {} +if is_torch_available(): + from .core_model_loading import Concatenate, Transpose, WeightConverter, WeightRenaming + from .gguf_conversion_ops import ( + BloomReshapeQKVBias, + BloomReshapeQKVWeight, + LogNegate, + ReversePermuteAttnK, + ReversePermuteAttnQ, + SubtractOne, + Unsqueeze, + ) + from .integrations.gguf_dequant import GGUFQuantizedTensor # noqa: F401 (referenced by load_gguf_checkpoint) + + _BLK_PREFIX = WeightRenaming(r"^blk\.", "model.layers.") + _LLAMA_SHARED_RENAMES = [ + _BLK_PREFIX, + WeightRenaming(r"^token_embd\.weight", "model.embed_tokens.weight"), + WeightRenaming(r"^output_norm\.weight", "model.norm.weight"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"\.attn_v\.weight", ".self_attn.v_proj.weight"), + WeightRenaming(r"\.attn_output\.weight", ".self_attn.o_proj.weight"), + WeightRenaming(r"\.ffn_(gate|up|down)\.weight", r".mlp.\1_proj.weight"), + # Attn biases exist on some Llama-family archs (Qwen2/3, StableLM, …) and not others; + # the rules are no-ops when the GGUF doesn't ship those tensors. + WeightRenaming(r"\.attn_(q|k|v)\.bias", r".self_attn.\1_proj.bias"), + ] + _ROPE_ATTN_CONVERTERS = [ + WeightConverter( + source_patterns=r"\.attn_q\.weight", + target_patterns=".self_attn.q_proj.weight", + operations=[ReversePermuteAttnQ()], + ), + WeightConverter( + source_patterns=r"\.attn_k\.weight", + target_patterns=".self_attn.k_proj.weight", + operations=[ReversePermuteAttnK()], + ), + ] + _NORM_RENAMES = [ + WeightRenaming(r"\.attn_norm\.weight", ".input_layernorm.weight"), + WeightRenaming(r"\.ffn_norm\.weight", ".post_attention_layernorm.weight"), + ] + _NORM_SUBTRACT_ONE_CONVERTERS = [ + # Nemotron / Gemma store every RMSNorm weight as `w + 1` (legacy + # `NemotronTensorProcessor` / `Gemma2TensorProcessor`). The final + # `output_norm` is no exception — apply `SubtractOne` there too. + WeightConverter( + source_patterns=r"^output_norm\.weight", + target_patterns="model.norm.weight", + operations=[SubtractOne()], + ), + WeightConverter( + source_patterns=r"\.attn_norm\.weight", + target_patterns=".input_layernorm.weight", + operations=[SubtractOne()], + ), + WeightConverter( + source_patterns=r"\.ffn_norm\.weight", + target_patterns=".post_attention_layernorm.weight", + operations=[SubtractOne()], + ), + ] + + # Only the `llama` arch itself applies llama.cpp's head-pair Q/K permute. + # Every other "Llama-family" arch (mistral, phi3, cohere, qwen2/3, deci, + # nemotron, gemma2/3, stablelm, lfm2) stores `attn_q` / `attn_k` in + # HF-native head layout — they need plain renames, not a permute, or the + # row shuffle corrupts attention. + _PLAIN_ATTN_QK_RENAMES = [ + WeightRenaming(r"\.attn_q\.weight", ".self_attn.q_proj.weight"), + WeightRenaming(r"\.attn_k\.weight", ".self_attn.k_proj.weight"), + ] + _LLAMA_CONVERTERS = _LLAMA_SHARED_RENAMES + _NORM_RENAMES + _ROPE_ATTN_CONVERTERS + _LLAMA_FAMILY_CONVERTERS = _LLAMA_SHARED_RENAMES + _NORM_RENAMES + _PLAIN_ATTN_QK_RENAMES + # Nemotron handles `output_norm` via the `SubtractOne` converter above, + # so we drop the plain `^output_norm\.weight` rename from the shared + # list to avoid double-mapping the same source key. + _NEMOTRON_SHARED_RENAMES = [r for r in _LLAMA_SHARED_RENAMES if r.source_patterns[0] != r"^output_norm\.weight"] + # Nemotron uses LayerNorm (not RMSNorm), so every norm carries a `.bias` + # alongside the `.weight`. Biases pass through as-is (no `SubtractOne`). + _NEMOTRON_NORM_BIAS_RENAMES = [ + WeightRenaming(r"^output_norm\.bias", "model.norm.bias"), + WeightRenaming(r"\.attn_norm\.bias", ".input_layernorm.bias"), + WeightRenaming(r"\.ffn_norm\.bias", ".post_attention_layernorm.bias"), + ] + _NEMOTRON_CONVERTERS = ( + _NEMOTRON_SHARED_RENAMES + _NORM_SUBTRACT_ONE_CONVERTERS + _NEMOTRON_NORM_BIAS_RENAMES + _PLAIN_ATTN_QK_RENAMES + ) -class TensorProcessor: - def __init__(self, config=None): - self.config = config or {} - - def preprocess_name(self, hf_name: str) -> str: - """ - Preprocesses the tensor name to ease loading the GGUF tensors. - """ - return hf_name - - def perform_fallback_tensor_mapping( - self, gguf_to_hf_name_map: dict[str, str], suffix: str, qual_name: str, hf_name: str - ): - """ - Called when get_gguf_hf_weights_map fails to map a HF parameter - (tensor) and corresponding GGUF one. - - This is particularly useful to resolve one-to-many - HF-GGUF mappings sometimes appear in some MoE models. - """ - pass - - def process(self, weights, name, **kwargs): - return GGUFTensor(weights, name, {}) - - -class LlamaTensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - def process(self, weights, name, **kwargs): - if ".attn_k." in name or ".attn_q." in name: - num_heads = self.config.get("num_attention_heads") - num_kv_heads = self.config.get("num_key_value_heads") - - if None in (num_heads, num_kv_heads): - return GGUFTensor(weights, name, {}) - if ".attn_q." in name: - weights = self._reverse_permute_weights(weights, num_heads, num_heads) - elif ".attn_k." in name: - weights = self._reverse_permute_weights(weights, num_heads, num_kv_heads) - return GGUFTensor(weights, name, {}) - - def _reverse_permute_weights( - self, weights: np.ndarray, n_head: int, num_kv_heads: int | None = None - ) -> np.ndarray: - # Original permutation implementation - # https://github.com/ggerganov/llama.cpp/blob/a38b884c6c4b0c256583acfaaabdf556c62fabea/convert_hf_to_gguf.py#L1402-L1408 - if num_kv_heads is not None and n_head != num_kv_heads: - n_head = num_kv_heads - - dim = weights.shape[0] // n_head // 2 - w = weights.reshape(n_head, dim, 2, *weights.shape[1:]) - return w.swapaxes(2, 1).reshape(weights.shape) - - -class Qwen2MoeTensorProcessor(TensorProcessor): - HF_EXPERT_RENAME_PATTERN = re.compile(r"mlp.experts.\d+.") - HF_MOE_W13_PATTERN = re.compile(r"model\.layers\.(?P\d+)\.mlp\.experts\.gate_up_proj") - GGUF_MOE_WEIGHTS_PATTERN = re.compile(r"(?P.*\.ffn_(?Pgate|down|up)_exps)\.weight$") - - def __init__(self, config=None): - super().__init__(config=config) - - def preprocess_name(self, hf_name: str) -> str: - return re.sub(self.HF_EXPERT_RENAME_PATTERN, "mlp.experts.", hf_name) - - def perform_fallback_tensor_mapping( - self, gguf_to_hf_name_map: dict[str, str], suffix: str, qual_name: str, hf_name: str - ): - # Map merged MoE weights (w1 (gate) and w3 (up)) separately. - if m := re.fullmatch(self.HF_MOE_W13_PATTERN, hf_name): - full_hf_name = qual_name + hf_name - gguf_to_hf_name_map[f"blk.{m['bid']}.ffn_gate_exps{suffix}"] = full_hf_name - gguf_to_hf_name_map[f"blk.{m['bid']}.ffn_up_exps{suffix}"] = full_hf_name - - def process(self, weights, name: str, **kwargs): - if m := re.fullmatch(self.GGUF_MOE_WEIGHTS_PATTERN, name): - tensor_key_mapping = kwargs.get("tensor_key_mapping") - parsed_parameters = kwargs.get("parsed_parameters") - if tensor_key_mapping: - self._set_moe_expert_tensor(weights, parsed_parameters, tensor_key_mapping[m["name"]], m["w"]) - return GGUFTensor(weights, None, {}) - if "ffn_gate_inp_shexp" in name: - # for compatibility tensor shared_expert_gate must be (1, 2048) dim, - # quantized one is (2048) - weights = np.expand_dims(weights, axis=0) - return GGUFTensor(weights, name, {}) - - def _set_moe_expert_tensor(self, weights: np.ndarray, parsed_parameters: dict[str, dict], hf_name: str, w: str): - torch_weights = torch.from_numpy(np.copy(weights)) - if w == "down": - parsed_parameters["tensors"][hf_name] = torch_weights - else: - # Double the size of the second dimension to interleave w1 (gate) and w3 (up) - # weights per expert (which is the first dimension). - # w1 (gate) comes first and w3 (up) comes second. - # ref: https://github.com/vllm-project/vllm/blob/8f8fda261a620234fdeea338f44093d5d8072879/vllm/model_executor/layers/fused_moe/layer.py#L988-L1015 - shape = list(weights.shape) - shard_dim = 1 - shard_size = shape[shard_dim] - shape[shard_dim] = shard_size * 2 - if hf_name not in parsed_parameters["tensors"]: - parsed_parameters["tensors"][hf_name] = torch.zeros(shape, dtype=torch_weights.dtype) - out: torch.Tensor = parsed_parameters["tensors"][hf_name] - if w == "gate": - out = out.narrow(shard_dim, 0, shard_size) - else: # w == "up" - out = out.narrow(shard_dim, shard_size, shard_size) - out.copy_(torch_weights) - - -class GptOssTensorProcessor(TensorProcessor): - """ - Tensor processor for GPT-OSS models (MoE with 128 experts). - Handles: - - Splitting stacked expert tensors (down_proj, gate_proj, up_proj) into individual experts. - - Interleaving gate and up projections if stored in a combined tensor (gate_up_projs). - - Bias tensors (1D) are passed through without transpose. - """ - - # Regex for separate expert tensors: e.g., blk.0.ffn_down_projs.weight - GGUF_MOE_WEIGHTS_PATTERN = re.compile(r"blk\.(?P\d+)\.ffn_(?Pdown|gate|up)_projs\.weight$") - # Regex for combined gate+up tensor: e.g., blk.0.ffn_gate_up_projs.weight - GGUF_MOE_COMBINED_PATTERN = re.compile(r"blk\.(?P\d+)\.ffn_gate_up_projs\.weight$") - - def __init__(self, config=None): - super().__init__(config=config) - - def process(self, weights, name: str, **kwargs): - # 1. Handle separate MoE expert tensors (down, gate, up) - if m := self.GGUF_MOE_WEIGHTS_PATTERN.match(name): - tensor_key_mapping = kwargs.get("tensor_key_mapping") - parsed_parameters = kwargs.get("parsed_parameters") - if tensor_key_mapping and parsed_parameters: - self._split_moe_expert_tensor(weights, parsed_parameters, m["bid"], m["proj"], tensor_key_mapping) - return GGUFTensor(weights, None, {}) # signal handled - - # 2. Handle combined gate+up tensor - if m := self.GGUF_MOE_COMBINED_PATTERN.match(name): - tensor_key_mapping = kwargs.get("tensor_key_mapping") - parsed_parameters = kwargs.get("parsed_parameters") - if tensor_key_mapping and parsed_parameters: - self._interleave_gate_up_tensor(weights, parsed_parameters, m["bid"], tensor_key_mapping) - return GGUFTensor(weights, None, {}) - - # 3. Bias tensors (1D) → no transpose - if ".bias" in name and len(weights.shape) == 1: - return GGUFTensor(weights, name, {}) - - # 4. Default handling for all other tensors - return GGUFTensor(weights, name, {}) - - def _split_moe_expert_tensor( - self, - weights: np.ndarray, - parsed_parameters: dict, - bid: str, - proj: str, - tensor_key_mapping: dict, - ): - """Split a stacked MoE tensor into individual expert tensors.""" - num_experts = self.config.get("num_local_experts", 128) - # Expected shape: [num_experts, hidden_size, intermediate_size] (or swapped). - # We assume the stored order is correct for the projection after splitting. - for i in range(min(num_experts, weights.shape[0])): - expert_weight = weights[i] # shape: [hidden, inter] or [inter, hidden] - # Build HF parameter name - hf_name = f"model.layers.{bid}.block_sparse_moe.experts.{i}.{proj}_proj.weight" - # Apply any user‑provided tensor key mapping - for key, mapped_key in tensor_key_mapping.items(): - if key in hf_name: - hf_name = hf_name.replace(key, mapped_key) - # Store the tensor - parsed_parameters["tensors"][hf_name] = torch.tensor(expert_weight, copy=True) - - def _interleave_gate_up_tensor( - self, - weights: np.ndarray, - parsed_parameters: dict, - bid: str, - tensor_key_mapping: dict, - ): - """ - Process a combined gate+up tensor. - Expected shape: [num_experts, intermediate_size, hidden_size]. - Interleaving: gate occupies first half of intermediate dimension, - up occupies second half. Transpose to [hidden, half_inter] per expert. - """ - num_experts = self.config.get("num_local_experts", 128) - inter_size = weights.shape[1] - half_inter = inter_size // 2 - gate_part = weights[:, :half_inter, :] # [E, half_inter, hidden] - up_part = weights[:, half_inter:, :] # [E, half_inter, hidden] - - for i in range(min(num_experts, weights.shape[0])): - gate_weight = gate_part[i].T # [hidden, half_inter] - up_weight = up_part[i].T # [hidden, half_inter] - - gate_name = f"model.layers.{bid}.block_sparse_moe.experts.{i}.gate_proj.weight" - up_name = f"model.layers.{bid}.block_sparse_moe.experts.{i}.up_proj.weight" - - # Apply mapping - for key, mapped_key in tensor_key_mapping.items(): - if key in gate_name: - gate_name = gate_name.replace(key, mapped_key) - if key in up_name: - up_name = up_name.replace(key, mapped_key) - - parsed_parameters["tensors"][gate_name] = torch.tensor(gate_weight, copy=True) - parsed_parameters["tensors"][up_name] = torch.tensor(up_weight, copy=True) - - -class BloomTensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - def process(self, weights, name, **kwargs): - if "attn_qkv" in name: - num_heads = self.config["n_head"] - n_embed = self.config["hidden_size"] - if "weight" in name: - weights = self._reverse_reshape_weights(weights, num_heads, n_embed) - else: - weights = self._reverse_reshape_bias(weights, num_heads, n_embed) - return GGUFTensor(weights, name, {}) - - def _reverse_reshape_weights(self, weights: np.ndarray, n_head: int, n_embed: int): - # Original reshape implementation - # https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L972-L985 - q, k, v = np.array_split(weights, 3, axis=0) - - q = q.reshape(n_head, n_embed // n_head, n_embed) - k = k.reshape(n_head, n_embed // n_head, n_embed) - v = v.reshape(n_head, n_embed // n_head, n_embed) - qkv_weights = np.stack([q, k, v], axis=1) - - return qkv_weights.reshape(n_head * 3 * (n_embed // n_head), n_embed) - - def _reverse_reshape_bias(self, weights: np.ndarray, n_head: int, n_embed: int): - # Original reshape implementation - # https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L986-L998 - q_bias, k_bias, v_bias = np.array_split(weights, 3) - - q_bias = q_bias.reshape(n_head, n_embed // n_head) - k_bias = k_bias.reshape(n_head, n_embed // n_head) - v_bias = v_bias.reshape(n_head, n_embed // n_head) - - qkv_bias = np.stack([q_bias, k_bias, v_bias], axis=1).flatten() - return qkv_bias - - -class T5TensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - def process(self, weights, name, **kwargs): - bid = None - for chunk in name.split("."): - if chunk.isdigit(): - bid = int(chunk) - break - return GGUFTensor(weights, name, {"bid": bid}) - - -class GPT2TensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - def process(self, weights, name, **kwargs): - # Original transpose implementation - # https://github.com/ggerganov/llama.cpp/blob/a38b884c6c4b0c256583acfaaabdf556c62fabea/convert_hf_to_gguf.py#L2060-L2061 - if ( - "attn_qkv.weight" in name - or "ffn_down.weight" in name - or "ffn_up.weight" in name - or "attn_output.weight" in name - ): - weights = weights.T - - # Handle special case for output.weight - if name == "output.weight": - # output.weight has conflicts with attn_output.weight in name checking - # Store the tensor directly and signal to skip further processing - name = "lm_head.weight" - parsed_parameters = kwargs.get("parsed_parameters", {}) - parsed_parameters["tensors"][name] = torch.from_numpy(np.copy(weights)) - name = None # Signal to skip further processing - return GGUFTensor(weights, name, {}) - - -class MambaTensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - def process(self, weights, name, **kwargs): - if "ssm_conv1d.weight" in name: - # for compatibility tensor ssm_conv1d must be (5120, 1, 4]) dim, - # quantized one is (5120, 4) - weights = np.expand_dims(weights, axis=1) - if "ssm_a" in name: - # Original exponential implementation - # https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L2975-L2977 - weights = np.log(-weights) - return GGUFTensor(weights, name, {}) - - -class NemotronTensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - # ref : https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L4666 - def process(self, weights, name, **kwargs): - if "norm.weight" in name: - weights = weights - 1 - return GGUFTensor(weights, name, {}) - - -class Gemma2TensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - # ref: https://github.com/ggerganov/llama.cpp/blob/d79d8f39b4da6deca4aea8bf130c6034c482b320/convert_hf_to_gguf.py#L3191 - # ref: https://github.com/huggingface/transformers/blob/fc37f38915372c15992b540dfcbbe00a916d4fc6/src/transformers/models/gemma/modeling_gemma.py#L89 - def process(self, weights, name, **kwargs): - if "norm.weight" in name: - weights = weights - 1 - return GGUFTensor(weights, name, {}) - - -class Lfm2TensorProcessor(TensorProcessor): - def __init__(self, config=None): - super().__init__(config=config) - - def process(self, weights, name, **kwargs): - if "shortconv.conv.weight" in name: - ## GGUF shape is [hidden_dim, L_cache], HF expects [hidden_dim, 1, L_cache] - weights = np.expand_dims(weights, axis=1) ## equivalent to unsqueeze(1) - return GGUFTensor(weights, name, {}) - - -class MiniMaxM2TensorProcessor(TensorProcessor): - HF_EXPERT_RENAME_PATTERN = re.compile(r"mlp\.experts\.\d+\.") - HF_MOE_W13_PATTERN = re.compile(r"(?:model\.)?layers\.(?P\d+)\.mlp\.experts\.gate_up_proj") - GGUF_MOE_WEIGHTS_PATTERN = re.compile(r"(?P.*\.ffn_(?Pgate|down|up)_exps)\.weight$") - HF_BIAS_PATTERN = re.compile(r"(?:model\.)?layers\.(?P\d+)\.mlp\.e_score_correction_bias") - - def __init__(self, config=None): - super().__init__(config=config) - - def preprocess_name(self, hf_name: str) -> str: - return re.sub(self.HF_EXPERT_RENAME_PATTERN, "mlp.experts.", hf_name) - - def perform_fallback_tensor_mapping( - self, gguf_to_hf_name_map: dict[str, str], suffix: str, qual_name: str, hf_name: str - ): - # Map merged gate_up_proj to both ffn_gate_exps and ffn_up_exps GGUF tensors. - if m := re.fullmatch(self.HF_MOE_W13_PATTERN, hf_name): - full_hf_name = qual_name + hf_name - gguf_to_hf_name_map[f"blk.{m['bid']}.ffn_gate_exps{suffix}"] = full_hf_name - gguf_to_hf_name_map[f"blk.{m['bid']}.ffn_up_exps{suffix}"] = full_hf_name - # Map e_score_correction_bias to GGUF exp_probs_b.bias. - elif m := re.fullmatch(self.HF_BIAS_PATTERN, hf_name): - gguf_to_hf_name_map[f"blk.{m['bid']}.exp_probs_b.bias"] = qual_name + hf_name - - def process(self, weights, name: str, **kwargs): - if m := re.fullmatch(self.GGUF_MOE_WEIGHTS_PATTERN, name): - tensor_key_mapping = kwargs.get("tensor_key_mapping") - parsed_parameters = kwargs.get("parsed_parameters") - if tensor_key_mapping: - self._set_moe_expert_tensor(weights, parsed_parameters, tensor_key_mapping[m["name"]], m["w"]) - return GGUFTensor(weights, None, {}) - return GGUFTensor(weights, name, {}) - - def _set_moe_expert_tensor(self, weights: np.ndarray, parsed_parameters: dict[str, dict], hf_name: str, w: str): - torch_weights = torch.from_numpy(np.copy(weights)) - if w == "down": - parsed_parameters["tensors"][hf_name] = torch_weights - else: - # Merge gate and up into gate_up_proj [num_experts, 2*intermediate, hidden] - shape = list(weights.shape) - shard_dim = 1 - shard_size = shape[shard_dim] - shape[shard_dim] = shard_size * 2 - if hf_name not in parsed_parameters["tensors"]: - parsed_parameters["tensors"][hf_name] = torch.zeros(shape, dtype=torch_weights.dtype) - out: torch.Tensor = parsed_parameters["tensors"][hf_name] - if w == "gate": - out = out.narrow(shard_dim, 0, shard_size) - else: # w == "up" - out = out.narrow(shard_dim, shard_size, shard_size) - out.copy_(torch_weights) - - -TENSOR_PROCESSORS = { - "llama": LlamaTensorProcessor, - "qwen2moe": Qwen2MoeTensorProcessor, - "gpt_oss": GptOssTensorProcessor, - "qwen3moe": Qwen2MoeTensorProcessor, - "bloom": BloomTensorProcessor, - "t5": T5TensorProcessor, - "t5encoder": T5TensorProcessor, - "gpt2": GPT2TensorProcessor, - "mamba": MambaTensorProcessor, - "nemotron": NemotronTensorProcessor, - "gemma2": Gemma2TensorProcessor, - "gemma3": Gemma2TensorProcessor, - "lfm2": Lfm2TensorProcessor, - "minimax-m2": MiniMaxM2TensorProcessor, -} + # Gemma-2/3 have four layer norms per block (input + post-attention + pre-ffn + post-ffn) + # and store every RMSNorm weight as `w + 1` (so we apply SubtractOne everywhere). + # `_NORM_SUBTRACT_ONE_CONVERTERS` already maps attn_norm → input_layernorm but its + # ffn_norm → post_attention_layernorm rule is wrong for Gemma — drop it and add the + # four-norm mapping explicitly. + _GEMMA_NORM_CONVERTERS = [ + WeightConverter( + source_patterns=r"\.attn_norm\.weight", + target_patterns=".input_layernorm.weight", + operations=[SubtractOne()], + ), + WeightConverter( + source_patterns=r"\.post_attention_norm\.weight", + target_patterns=".post_attention_layernorm.weight", + operations=[SubtractOne()], + ), + WeightConverter( + source_patterns=r"\.ffn_norm\.weight", + target_patterns=".pre_feedforward_layernorm.weight", + operations=[SubtractOne()], + ), + WeightConverter( + source_patterns=r"\.post_ffw_norm\.weight", + target_patterns=".post_feedforward_layernorm.weight", + operations=[SubtractOne()], + ), + # Output norm also stored as `w + 1` for Gemma2/3 — needs `SubtractOne`. + WeightConverter( + source_patterns=r"^output_norm\.weight", + target_patterns="model.norm.weight", + operations=[SubtractOne()], + ), + ] + # Gemma2/3 share `_LLAMA_SHARED_RENAMES` for everything that isn't a norm — + # but the shared block also includes the plain `output_norm.weight → model.norm.weight` + # rename that we override above. Strip it from the shared list to avoid + # double-mapping the same source key. + _GEMMA_SHARED_RENAMES = [r for r in _LLAMA_SHARED_RENAMES if r.source_patterns[0] != r"^output_norm\.weight"] + _GEMMA2_CONVERTERS = _GEMMA_SHARED_RENAMES + _GEMMA_NORM_CONVERTERS + _PLAIN_ATTN_QK_RENAMES + + # Gemma-3 adds per-head q/k RMSNorm inside attention (also stored as `w + 1`). + _GEMMA3_CONVERTERS = _GEMMA2_CONVERTERS + [ + WeightConverter( + source_patterns=r"\.attn_q_norm\.weight", + target_patterns=".self_attn.q_norm.weight", + operations=[SubtractOne()], + ), + WeightConverter( + source_patterns=r"\.attn_k_norm\.weight", + target_patterns=".self_attn.k_norm.weight", + operations=[SubtractOne()], + ), + ] + + _T5_CONVERTERS = [ + WeightRenaming(r"^enc\.blk\.", "encoder.block."), + WeightRenaming(r"^dec\.blk\.", "decoder.block."), + WeightRenaming(r"^enc\.output_norm\.weight", "encoder.final_layer_norm.weight"), + WeightRenaming(r"^dec\.output_norm\.weight", "decoder.final_layer_norm.weight"), + WeightRenaming(r"^token_embd\.weight", "shared.weight"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + # Self-attention (layer.0) — shared by encoder + decoder + WeightRenaming(r"\.attn_(q|k|v|o)\.weight", r".layer.0.SelfAttention.\1.weight"), + WeightRenaming(r"\.attn_rel_b\.weight", ".layer.0.SelfAttention.relative_attention_bias.weight"), + WeightRenaming(r"\.attn_norm\.weight", ".layer.0.layer_norm.weight"), + # Cross-attention (layer.1) — decoder only + WeightRenaming(r"\.cross_attn_(q|k|v|o)\.weight", r".layer.1.EncDecAttention.\1.weight"), + WeightRenaming(r"\.cross_attn_norm\.weight", ".layer.1.layer_norm.weight"), + # FFN — encoder uses layer.1, decoder uses layer.2. Inject the layer index + # via two structural renames so the per-tensor renames below are shared. + WeightRenaming(r"^encoder\.block\.(\d+)\.ffn_", r"encoder.block.\1.layer.1.ffn_"), + WeightRenaming(r"^decoder\.block\.(\d+)\.ffn_", r"decoder.block.\1.layer.2.ffn_"), + WeightRenaming(r"\.ffn_norm\.weight", ".layer_norm.weight"), + WeightRenaming(r"\.ffn_gate\.weight", ".DenseReluDense.wi_0.weight"), + WeightRenaming(r"\.ffn_up\.weight", ".DenseReluDense.wi_1.weight"), + WeightRenaming(r"\.ffn_down\.weight", ".DenseReluDense.wo.weight"), + ] + + # --- StableLM (Llama-family attn — no rope permute — plus q_norm/k_norm and biases) + # StableLM uses LayerNorm (not RMSNorm) so every norm carries a `bias` + # alongside `weight`. The base `_NORM_RENAMES` and `_LLAMA_SHARED_RENAMES` + # only cover `.weight`; add the matching `.bias` renames here. + _STABLELM_CONVERTERS = _LLAMA_FAMILY_CONVERTERS + [ + WeightRenaming(r"^output_norm\.bias", "model.norm.bias"), + WeightRenaming(r"\.attn_norm\.bias", ".input_layernorm.bias"), + WeightRenaming(r"\.ffn_norm\.bias", ".post_attention_layernorm.bias"), + WeightRenaming(r"\.attn_(q|k)_norm\.weight", r".self_attn.\1_layernorm.weight"), + WeightRenaming(r"\.attn_(q|k|v)\.bias", r".self_attn.\1_proj.bias"), + ] + + # --- Starcoder2 (Llama-style attn + single-layer c_fc/c_proj MLP, biases) -- + _STARCODER2_CONVERTERS = [ + _BLK_PREFIX, + WeightRenaming(r"^token_embd\.weight", "model.embed_tokens.weight"), + WeightRenaming(r"^output_norm\.(weight|bias)", r"model.norm.\1"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"\.attn_v\.(weight|bias)", r".self_attn.v_proj.\1"), + WeightRenaming(r"\.attn_output\.(weight|bias)", r".self_attn.o_proj.\1"), + WeightRenaming(r"\.attn_(q|k)\.bias", r".self_attn.\1_proj.bias"), + WeightRenaming(r"\.ffn_up\.(weight|bias)", r".mlp.c_fc.\1"), + WeightRenaming(r"\.ffn_down\.(weight|bias)", r".mlp.c_proj.\1"), + WeightRenaming(r"\.attn_norm\.(weight|bias)", r".input_layernorm.\1"), + WeightRenaming(r"\.ffn_norm\.(weight|bias)", r".post_attention_layernorm.\1"), + # Starcoder2 stores Q/K in HF-native head layout; no rope permute. + *_PLAIN_ATTN_QK_RENAMES, + ] + + # --- Falcon (merged QKV, dense_h_to_4h MLP, 7B vs 40B norm variants) ------- + _FALCON_CONVERTERS = [ + WeightRenaming(r"^blk\.", "transformer.h."), + WeightRenaming(r"^token_embd\.weight", "transformer.word_embeddings.weight"), + WeightRenaming(r"^output_norm\.(weight|bias)", r"transformer.ln_f.\1"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"\.attn_qkv\.weight", ".self_attention.query_key_value.weight"), + WeightRenaming(r"\.attn_output\.weight", ".self_attention.dense.weight"), + WeightRenaming(r"\.ffn_up\.weight", ".mlp.dense_h_to_4h.weight"), + WeightRenaming(r"\.ffn_down\.weight", ".mlp.dense_4h_to_h.weight"), + # 7B style: parallel attn+mlp shares one input_layernorm + WeightRenaming(r"\.attn_norm\.(weight|bias)", r".input_layernorm.\1"), + # 40B style: separate ln_attn / ln_mlp + WeightRenaming(r"\.attn_norm_2\.(weight|bias)", r".ln_mlp.\1"), + ] + + # --- GPT-OSS (MoE: router + post_attention_norm + merged gate_up_proj) ----- + _GPT_OSS_CONVERTERS = [ + _BLK_PREFIX, + WeightRenaming(r"^token_embd\.weight", "model.embed_tokens.weight"), + WeightRenaming(r"^output_norm\.weight", "model.norm.weight"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"\.attn_v\.weight", ".self_attn.v_proj.weight"), + WeightRenaming(r"\.attn_output\.weight", ".self_attn.o_proj.weight"), + WeightRenaming(r"\.attn_norm\.weight", ".input_layernorm.weight"), + WeightRenaming(r"\.post_attention_norm\.weight", ".post_attention_layernorm.weight"), + WeightRenaming(r"\.ffn_norm\.weight", ".post_attention_layernorm.weight"), + WeightRenaming(r"\.ffn_gate_inp\.weight", ".mlp.router.weight"), + WeightRenaming(r"\.ffn_down_exps\.weight", ".mlp.experts.down_proj"), + *_ROPE_ATTN_CONVERTERS, + WeightConverter( + source_patterns=[r"\.ffn_gate_exps\.weight", r"\.ffn_up_exps\.weight"], + target_patterns=".mlp.experts.gate_up_proj", + operations=[Concatenate(dim=1)], + ), + ] + + # --- MiniMax-M2 (Qwen3-MoE-like + e_score_correction_bias) ----------------- + _MINIMAX_M2_CONVERTERS = [ + _BLK_PREFIX, + WeightRenaming(r"^token_embd\.weight", "model.embed_tokens.weight"), + WeightRenaming(r"^output_norm\.weight", "model.norm.weight"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"\.attn_v\.weight", ".self_attn.v_proj.weight"), + WeightRenaming(r"\.attn_output\.weight", ".self_attn.o_proj.weight"), + WeightRenaming(r"\.attn_(q|k)_norm\.weight", r".self_attn.\1_norm.weight"), + WeightRenaming(r"\.attn_norm\.weight", ".input_layernorm.weight"), + WeightRenaming(r"\.ffn_norm\.weight", ".post_attention_layernorm.weight"), + WeightRenaming(r"\.ffn_gate_inp\.weight", ".mlp.gate.weight"), + WeightRenaming(r"\.exp_probs_b\.bias", ".mlp.e_score_correction_bias"), + WeightRenaming(r"\.ffn_down_exps\.weight", ".mlp.experts.down_proj"), + # MiniMax-M2 stores Q/K in HF-native head layout; no rope permute. + *_PLAIN_ATTN_QK_RENAMES, + WeightConverter( + source_patterns=[r"\.ffn_gate_exps\.weight", r"\.ffn_up_exps\.weight"], + target_patterns=".mlp.experts.gate_up_proj", + operations=[Concatenate(dim=1)], + ), + ] + + # --- Qwen2-MoE / Qwen3-MoE (shared experts + merged gate_up_proj) ---------- + # NB: modern Qwen2/3-MoE GGUFs ship Q/K projection weights already in HF + # layout (no llama.cpp-style head permute) — applying `ReversePermuteAttn*` + # here would corrupt the matmul output ~140% relative. Tested against + # gdax/Qwen1.5-MoE-A2.7B_gguf: max abs diff to the HF reference is + # 3.8e-2 (Q4_K quantization noise) without the permute, 1.27 with it. + _QWEN2_MOE_CONVERTERS = [ + _BLK_PREFIX, + WeightRenaming(r"^token_embd\.weight", "model.embed_tokens.weight"), + WeightRenaming(r"^output_norm\.weight", "model.norm.weight"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"\.attn_q\.weight", ".self_attn.q_proj.weight"), + WeightRenaming(r"\.attn_k\.weight", ".self_attn.k_proj.weight"), + WeightRenaming(r"\.attn_v\.weight", ".self_attn.v_proj.weight"), + WeightRenaming(r"\.attn_output\.weight", ".self_attn.o_proj.weight"), + WeightRenaming(r"\.attn_(q|k|v)\.bias", r".self_attn.\1_proj.bias"), + WeightRenaming(r"\.attn_norm\.weight", ".input_layernorm.weight"), + WeightRenaming(r"\.ffn_norm\.weight", ".post_attention_layernorm.weight"), + WeightRenaming(r"\.ffn_(gate|up|down)_shexp\.weight", r".mlp.shared_expert.\1_proj.weight"), + WeightRenaming(r"\.ffn_gate_inp\.weight", ".mlp.gate.weight"), + WeightRenaming(r"\.ffn_down_exps\.weight", ".mlp.experts.down_proj"), + WeightConverter( + source_patterns=r"\.ffn_gate_inp_shexp", + target_patterns=".mlp.shared_expert_gate", + operations=[Unsqueeze(0)], + ), + WeightConverter( + source_patterns=[r"\.ffn_gate_exps\.weight", r"\.ffn_up_exps\.weight"], + target_patterns=".mlp.experts.gate_up_proj", + operations=[Concatenate(dim=1)], + ), + ] + + # --- Bloom ----------------------------------------------------------------- + _BLOOM_CONVERTERS = [ + _BLK_PREFIX, + WeightRenaming(r"^token_embd\.(weight|bias)", r"transformer.word_embeddings.\1"), + WeightRenaming(r"^token_embd_norm\.(weight|bias)", r"transformer.word_embeddings_layernorm.\1"), + WeightRenaming(r"^output_norm\.(weight|bias)", r"transformer.ln_f.\1"), + WeightRenaming(r"^model\.layers\.", "transformer.h."), + WeightRenaming(r"\.attn_norm\.(weight|bias)", r".input_layernorm.\1"), + WeightRenaming(r"\.attn_output\.(weight|bias)", r".self_attention.dense.\1"), + WeightRenaming(r"\.ffn_norm\.(weight|bias)", r".post_attention_layernorm.\1"), + WeightRenaming(r"\.ffn_up\.(weight|bias)", r".mlp.dense_h_to_4h.\1"), + WeightRenaming(r"\.ffn_down\.(weight|bias)", r".mlp.dense_4h_to_h.\1"), + WeightConverter( + source_patterns=r"\.attn_qkv\.weight", + target_patterns=".self_attention.query_key_value.weight", + operations=[BloomReshapeQKVWeight()], + ), + WeightConverter( + source_patterns=r"\.attn_qkv\.bias", + target_patterns=".self_attention.query_key_value.bias", + operations=[BloomReshapeQKVBias()], + ), + ] + + # --- GPT-2 (Transpose on every Linear weight, bias kept as-is) ------------- + # GGUF stores GPT-2 `c_attn`/`c_proj`/`c_fc` weights transposed relative to HF. + # The TARGET names differ per source (attn_qkv→c_attn, attn_output/ffn_down→c_proj, + # ffn_up→c_fc), so the four Transpose converters can't be merged into one rule. + _GPT2_CONVERTERS = [ + _BLK_PREFIX, + WeightRenaming(r"^token_embd\.weight", "transformer.wte.weight"), + WeightRenaming(r"^position_embd\.weight", "transformer.wpe.weight"), + WeightRenaming(r"^output_norm\.(weight|bias)", r"transformer.ln_f.\1"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"^model\.layers\.", "transformer.h."), + WeightRenaming(r"\.attn_qkv\.bias", ".attn.c_attn.bias"), + WeightRenaming(r"\.attn_output\.bias", ".attn.c_proj.bias"), + WeightRenaming(r"\.ffn_up\.bias", ".mlp.c_fc.bias"), + WeightRenaming(r"\.ffn_down\.bias", ".mlp.c_proj.bias"), + WeightRenaming(r"\.attn_norm\.(weight|bias)", r".ln_1.\1"), + WeightRenaming(r"\.ffn_norm\.(weight|bias)", r".ln_2.\1"), + WeightConverter( + source_patterns=r"\.attn_qkv\.weight", target_patterns=".attn.c_attn.weight", operations=[Transpose()] + ), + WeightConverter( + source_patterns=r"\.attn_output\.weight", target_patterns=".attn.c_proj.weight", operations=[Transpose()] + ), + WeightConverter( + source_patterns=r"\.ffn_up\.weight", target_patterns=".mlp.c_fc.weight", operations=[Transpose()] + ), + WeightConverter( + source_patterns=r"\.ffn_down\.weight", target_patterns=".mlp.c_proj.weight", operations=[Transpose()] + ), + ] + + # --- Mamba (SSM) ----------------------------------------------------------- + _MAMBA_CONVERTERS = [ + WeightRenaming(r"^blk\.", "backbone.layers."), + WeightRenaming(r"^token_embd\.weight", "backbone.embeddings.weight"), + WeightRenaming(r"^output_norm\.weight", "backbone.norm_f.weight"), + WeightRenaming(r"^output\.weight", "lm_head.weight"), + WeightRenaming(r"\.ssm_in\.weight", ".mixer.in_proj.weight"), + WeightRenaming(r"\.ssm_conv1d\.bias", ".mixer.conv1d.bias"), + WeightRenaming(r"\.ssm_x\.weight", ".mixer.x_proj.weight"), + WeightRenaming(r"\.ssm_dt\.weight", ".mixer.dt_proj.weight"), + WeightRenaming(r"\.ssm_dt\.bias", ".mixer.dt_proj.bias"), + WeightRenaming(r"\.ssm_d$", ".mixer.D"), + WeightRenaming(r"\.ssm_out\.weight", ".mixer.out_proj.weight"), + WeightRenaming(r"\.attn_norm\.weight", ".norm.weight"), + WeightConverter( + source_patterns=r"\.ssm_conv1d\.weight", + target_patterns=".mixer.conv1d.weight", + operations=[Unsqueeze(1)], + ), + WeightConverter(source_patterns=r"\.ssm_a$", target_patterns=".mixer.A_log", operations=[LogNegate()]), + ] + + # --- LFM2 (Llama-family attn — no rope permute — plus shortconv) ---------- + _LFM2_CONVERTERS = _LLAMA_FAMILY_CONVERTERS + [ + WeightConverter( + source_patterns=r"\.shortconv\.conv\.weight", + target_patterns=".conv.weight", + operations=[Unsqueeze(1)], + ), + ] + + _GGUF_ARCH_CONVERTERS: dict[str, list] = { + # Llama itself is the only arch where llama.cpp applies the head-pair + # Q/K permute (`LlamaTensorProcessor` was the only legacy processor + # using `_reverse_permute_weights`). Everything else uses HF-native + # head layout — plain renames via `_LLAMA_FAMILY_CONVERTERS`. + "llama": _LLAMA_CONVERTERS, + "mistral": _LLAMA_FAMILY_CONVERTERS, + "phi3": _LLAMA_FAMILY_CONVERTERS, + "cohere": _LLAMA_FAMILY_CONVERTERS, + "qwen2": _LLAMA_FAMILY_CONVERTERS, + "qwen3": _LLAMA_FAMILY_CONVERTERS, + "deci": _LLAMA_FAMILY_CONVERTERS, + # Norm-subtract-one variants + "nemotron": _NEMOTRON_CONVERTERS, + "gemma2": _GEMMA2_CONVERTERS, + "gemma3": _GEMMA3_CONVERTERS, + "gemma3_text": _GEMMA3_CONVERTERS, + # Misc archs + "bloom": _BLOOM_CONVERTERS, + "gpt2": _GPT2_CONVERTERS, + "mamba": _MAMBA_CONVERTERS, + "lfm2": _LFM2_CONVERTERS, + "falcon": _FALCON_CONVERTERS, + "stablelm": _STABLELM_CONVERTERS, + "starcoder2": _STARCODER2_CONVERTERS, + # MoE + "qwen2_moe": _QWEN2_MOE_CONVERTERS, + "qwen3_moe": _QWEN2_MOE_CONVERTERS, + "minimax_m2": _MINIMAX_M2_CONVERTERS, + "gpt_oss": _GPT_OSS_CONVERTERS, + # T5 / UMT5 / T5-encoder share the same encoder–decoder mapping + "t5": _T5_CONVERTERS, + "t5encoder": _T5_CONVERTERS, + "umt5": _T5_CONVERTERS, + } + + +def get_gguf_converters(model_type: str) -> list: + """Return the GGUF→HF rename rules (`WeightRenaming` / `WeightConverter`) for a given HF model type.""" + return list(_GGUF_ARCH_CONVERTERS.get(model_type, [])) def read_field(reader, field): @@ -478,96 +474,7 @@ def read_field(reader, field): return [_gguf_parse_value(value.parts[_data_index], value.types) for _data_index in value.data] -# modified from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/loader.py#L1115-L1147 -def get_gguf_hf_weights_map( - hf_model, - processor: TensorProcessor, - model_type: str | None = None, - num_layers: int | None = None, - qual_name: str = "", -): - """ - GGUF uses this naming convention for their tensors from HF checkpoint: - `blk.N.BB.weight` and `blk.N.BB.bias` - where N signifies the block number of a layer, and BB signifies the - attention/mlp layer components. - See "Standardized tensor names" in - https://github.com/ggerganov/ggml/blob/master/docs/gguf.md for details. - """ - if is_gguf_available() and is_torch_available(): - from gguf import MODEL_ARCH_NAMES, get_tensor_name_map - else: - logger.error( - "Loading a GGUF checkpoint in PyTorch, requires both PyTorch and GGUF>=0.10.0 to be installed. Please see " - "https://pytorch.org/ and https://github.com/ggerganov/llama.cpp/tree/master/gguf-py for installation instructions." - ) - raise ImportError("Please install torch and gguf>=0.10.0 to load a GGUF checkpoint in PyTorch.") - - model_type = hf_model.config.model_type if model_type is None else model_type - num_layers = hf_model.config.num_hidden_layers if num_layers is None else num_layers - # hack: ggufs have a different name for cohere - if model_type == "cohere": - model_type = "command-r" - elif model_type == "qwen2_moe": - model_type = "qwen2moe" - elif model_type == "qwen3_moe": - model_type = "qwen3moe" - elif model_type == "gemma3_text": - model_type = "gemma3" - elif model_type == "umt5": - model_type = "t5" - elif model_type == "minimax_m2": - model_type = "minimax-m2" - elif model_type == "gpt_oss": - model_type = "gpt-oss" - arch = None - for key, value in MODEL_ARCH_NAMES.items(): - if value == model_type: - arch = key - break - if arch is None: - raise NotImplementedError( - f"Unknown gguf model_type: {model_type} in gguf-py. " - "This might because you're using an outdated version of gguf-py package, " - "you can install `gguf` package from source refer to " - "https://github.com/ggerganov/llama.cpp/tree/master/gguf-py#development" - ) - name_map = get_tensor_name_map(arch, num_layers) - - # Use a dummy conversion to get the mapping, because - # hf => gguf and gguf => hf mappings are reversed - gguf_to_hf_name_map = {} - state_dict = hf_model.state_dict() - for hf_name in state_dict: - hf_name = processor.preprocess_name(hf_name) - - name, suffix = hf_name, "" - if hf_name.endswith(".weight") or hf_name.endswith(".bias"): - name, suffix = hf_name.rsplit(".", 1) - suffix = "." + suffix - - gguf_name = name_map.get_name(name) - if gguf_name is None: - processor.perform_fallback_tensor_mapping(gguf_to_hf_name_map, suffix, qual_name, hf_name) - continue - - gguf_to_hf_name_map[gguf_name + suffix] = qual_name + hf_name - - # Some model like Bloom converted from BloomModel instead of BloomForCausalLM - # Therefore, we need to check submodule as well to get a correct mapping - if named_children := hf_model.named_children(): - for name, child in named_children: - sub_map = get_gguf_hf_weights_map( - child, processor, model_type, num_layers, qual_name=f"{qual_name}{name}." - ) - # Ignore the keys that are already in the main map to avoid overwriting - sub_map = {k: v for k, v in sub_map.items() if k not in gguf_to_hf_name_map} - gguf_to_hf_name_map.update(sub_map) - - return gguf_to_hf_name_map - - -def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_load=None, torch_dtype=None): +def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False): """ Load a GGUF file and return a dictionary of parsed parameters containing tensors, the parsed tokenizer and config attributes. @@ -576,17 +483,12 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo gguf_checkpoint_path (`str`): The path the to GGUF file to load return_tensors (`bool`, defaults to `False`): - Whether to read the tensors from the file and return them. Not doing so is faster - and only loads the metadata in memory. - model_to_load (`nn.Module`, *optional*): - The model to load the weights into. This is used to map GGUF tensor names to - Transformers parameter names. - torch_dtype (`torch.dtype`, *optional*): - The desired `torch.dtype` for the loaded tensors. If provided, tensors will be - converted to this dtype immediately after dequantization to save memory. + Whether to materialize the tensor data (`"tensors"`) and return it. When `False`, + only the header is read — config, tokenizer, and the cheap `"tensor_quant_types"` / + `"weight_mapping"` metadata — which is faster and enough to plan a module swap. """ if is_gguf_available() and is_torch_available(): - from gguf import GGUFReader, dequantize + from gguf import GGUFReader else: logger.error( "Loading a GGUF checkpoint in PyTorch, requires both PyTorch and GGUF>=0.10.0 to be installed. Please see " @@ -598,7 +500,7 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo fields = reader.fields reader_keys = list(fields.keys()) - parsed_parameters = {k: {} for k in GGUF_TO_TRANSFORMERS_MAPPING} + parsed_parameters: dict = {k: {} for k in GGUF_TO_TRANSFORMERS_MAPPING} architecture = read_field(reader, "general.architecture")[0] # NOTE: Some GGUF checkpoints may miss `general.name` field in metadata @@ -626,12 +528,8 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo if "qwen2moe" in architecture: updated_architecture = "qwen2_moe" - elif "gpt_oss" in architecture or "gpt-oss" in architecture: - updated_architecture = "gpt_oss" elif "qwen3moe" in architecture: updated_architecture = "qwen3_moe" - elif "minimax-m2" in architecture: - updated_architecture = "minimax_m2" # For stablelm architecture, we need to set qkv_bias and use_parallel_residual from tensors # If `qkv_bias=True`, qkv_proj with bias will be present in the tensors @@ -695,13 +593,6 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo if parsed_parameters["config"]["model_type"] == "gemma3": parsed_parameters["config"]["model_type"] = "gemma3_text" - # MiniMax-M2: convert expert_gating_func integer to scoring_func string - if parsed_parameters["config"].get("model_type") == "minimax_m2": - _gating_func_map = {0: "none", 1: "softmax", 2: "sigmoid"} - _scoring = parsed_parameters["config"].get("scoring_func") - if isinstance(_scoring, int): - parsed_parameters["config"]["scoring_func"] = _gating_func_map.get(_scoring, "softmax") - if parsed_parameters["config"]["model_type"] == "lfm2": gguf_num_key_value_heads = parsed_parameters["config"]["num_key_value_heads"] # LFM2 GGUF checkpoint defines num_key_value_heads as a list of integers .e.g [0, 0, 8, 0, 0, 8, 0, 0, 8, 0, 8, 0, 8, 0, 8, 0] but we need to set it to the max value for HF @@ -715,43 +606,59 @@ def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_lo i for i, num_kv_heads in enumerate(gguf_num_key_value_heads) if num_kv_heads > 0 ] + # MiniMax-M2: convert expert_gating_func integer to scoring_func string + if parsed_parameters["config"].get("model_type") == "minimax_m2": + _gating_func_map = {0: "none", 1: "softmax", 2: "sigmoid"} + _scoring = parsed_parameters["config"].get("scoring_func") + if isinstance(_scoring, int): + parsed_parameters["config"]["scoring_func"] = _gating_func_map.get(_scoring, "softmax") + + # Llama-3 family (Meta-Llama-3 / Llama-3.1 / Llama-3.2) ship the + # `llama3` rope_scaling implicitly — the GGUF file doesn't store + # `rope.scaling.*` keys, so we infer it from the well-known signature + # (arch=llama + freq_base=500000 + context_length >= 8192). HF expects + # `rope_scaling.rope_type = "llama3"` with the official factors; without + # this the model decodes as gibberish for any non-trivial prompt. + if ( + updated_architecture == "llama" + and parsed_parameters["config"].get("rope_theta") == 500000.0 + and (parsed_parameters["config"].get("max_position_embeddings", 0) or 0) >= 8192 * 2 + ): + max_pos = parsed_parameters["config"]["max_position_embeddings"] + # Llama-3.0 = 8192 → factor 1 (no scaling); Llama-3.1/3.2 use factor + # max_position_embeddings / original_max_position_embeddings. + original = 8192 + factor = max(1.0, float(max_pos) / float(original)) + if factor > 1.0: + parsed_parameters["config"]["rope_scaling"] = { + "rope_type": "llama3", + "factor": factor, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": original, + } + + # GPT-OSS: reconstruct rope_scaling from the architecture-prefixed metadata keys if updated_architecture == "gpt_oss": - # Helper to read keys with the correct prefix - def read_gpt_key(reader, suffix, default=None): - key = f"gpt-oss.{suffix}" - if key in reader.fields: - val = reader.fields[key].parts[0] - if isinstance(val, bytes): - val = val.decode("utf-8") - return val - return default - - # Reconstruct rope_scaling from GGUF metadata - rope_type = read_gpt_key(reader, "rope.scaling.type") - if rope_type is not None: + rope_type_field = reader.fields.get("gpt-oss.rope.scaling.type") + if rope_type_field is not None: + rope_type = rope_type_field.parts[0] + if isinstance(rope_type, bytes): + rope_type = rope_type.decode("utf-8") rope_scaling = {"rope_type": rope_type} - - # Collect all rope.scaling keys dynamically for key in reader.fields: - if not key.startswith("gpt-oss.rope.scaling."): + if not key.startswith("gpt-oss.rope.scaling.") or key.endswith(".type"): continue suffix = key[len("gpt-oss.rope.scaling.") :] - if suffix == "type": - continue value = reader.fields[key].parts[0] if isinstance(value, bytes): value = value.decode("utf-8") - # Convert to appropriate type if suffix in ("factor", "attention_factor", "beta_fast", "beta_slow"): value = float(value) elif suffix in ("original_context_length", "original_max_position_embeddings"): - # Map GGUF's original_context_length to HF's original_max_position_embeddings suffix = "original_max_position_embeddings" value = int(value) - else: - pass rope_scaling[suffix] = value - parsed_parameters["config"]["rope_scaling"] = rope_scaling # retrieve config vocab_size from tokenizer @@ -766,39 +673,37 @@ def read_gpt_key(reader, suffix, default=None): "This will use default value from model config class and cause unexpected behavior." ) - if return_tensors: - parsed_parameters["tensors"] = {} - - config = parsed_parameters.get("config", {}) - - ProcessorClass = TENSOR_PROCESSORS.get(architecture, TensorProcessor) - processor = ProcessorClass(config=config) - - tensor_key_mapping = get_gguf_hf_weights_map(model_to_load, processor) - - for tensor in tqdm(reader.tensors, desc="Converting and de-quantizing GGUF tensors..."): - name = tensor.name - weights = dequantize(tensor.data, tensor.tensor_type) - - result = processor.process( - weights=weights, - name=name, - tensor_key_mapping=tensor_key_mapping, - parsed_parameters=parsed_parameters, - ) - - weights = result.weights - name = result.name - - if name not in tensor_key_mapping: - continue + # Tensors GGUF ships as metadata for its own runtime but HF models compute on the + # fly (or don't store as a parameter). Skip so they don't show up as "unexpected". + _GGUF_RUNTIME_AUX_TENSORS = frozenset({"rope_freqs.weight"}) + + # Header-only metadata — read straight off the reader without touching any + # tensor data, so it is always cheap and available even when `return_tensors` + # is False. The GGUF quantizer plans its module swap from this alone (names + + # quant types + the rename/convert map): a pure-renaming step, no tensor load. + model_type = parsed_parameters.get("config", {}).get("model_type", architecture) + parsed_parameters["tensor_quant_types"] = { + tensor.name: tensor.tensor_type + for tensor in reader.tensors + if tensor.name not in _GGUF_RUNTIME_AUX_TENSORS + } + parsed_parameters["weight_mapping"] = get_gguf_converters(model_type) - name = tensor_key_mapping[name] - - tensor = torch.from_numpy(np.copy(weights)) - if torch_dtype is not None: - tensor = tensor.to(torch_dtype) - parsed_parameters["tensors"][name] = tensor + if return_tensors: + # Wrap raw uint8 bytes in a `torch.Tensor` subclass that carries `quant_type`. + # `GGUFDequantize` does the actual dequant inside the WeightConverter chain, + # on whatever device the loader has moved the bytes to. + import numpy as np + import torch # local: keep top-of-file import-light when torch isn't required + + # `np.copy(tensor.data)` materializes a writable buffer — the mmap view + # returned by GGUFReader is non-writable, and downstream `.to(device,dtype)` + # casts (especially on MPS) treat the storage as scratch and corrupt it. + parsed_parameters["tensors"] = { + tensor.name: GGUFQuantizedTensor(torch.from_numpy(np.copy(tensor.data)), quant_type=tensor.tensor_type) + for tensor in reader.tensors + if tensor.name not in _GGUF_RUNTIME_AUX_TENSORS + } if len(reader_keys) > 0: logger.info(f"Some keys of the GGUF file were not considered: {reader_keys}") diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 13c06973b4a9..5d5177ee44d1 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -546,6 +546,33 @@ def _add_variant(weights_name: str, variant: str | None = None) -> str: return weights_name +def _maybe_auto_detect_gguf(pretrained_model_name_or_path: str, **download_kwargs) -> str | None: + """Return a GGUF filename if ``pretrained_model_name_or_path`` resolves + unambiguously to one — a direct ``.gguf`` path, a local dir containing + exactly one ``.gguf``, or a Hub repo containing exactly one ``.gguf``. + Otherwise return ``None`` (multiple variants → caller must pass + ``gguf_file=`` explicitly to disambiguate).""" + if pretrained_model_name_or_path.lower().endswith(".gguf"): + return pretrained_model_name_or_path + if os.path.isdir(pretrained_model_name_or_path): + local_ggufs = [f for f in os.listdir(pretrained_model_name_or_path) if f.lower().endswith(".gguf")] + return local_ggufs[0] if len(local_ggufs) == 1 else None + # Hub repo: list files without downloading. + try: + from huggingface_hub import HfApi + + api = HfApi() + files = api.list_repo_files( + pretrained_model_name_or_path, + token=download_kwargs.get("token"), + revision=download_kwargs.get("revision"), + ) + except Exception: + return None + ggufs = [f for f in files if f.lower().endswith(".gguf")] + return ggufs[0] if len(ggufs) == 1 else None + + def _get_resolved_checkpoint_files( pretrained_model_name_or_path: str | os.PathLike | None, variant: str | None, @@ -4092,6 +4119,10 @@ def from_pretrained( # For BC on torch_dtype argument if torch_dtype is not None: dtype = dtype if dtype is not None else torch_dtype + # Track whether the user explicitly passed dtype: relevant for the GGUF + # path, where the default keeps weights native (Metal kernels) and an + # explicit dtype= switches to dequantize-on-load. + dtype_was_explicit = dtype is not None if dtype is None: dtype = "auto" @@ -4126,6 +4157,12 @@ def from_pretrained( tp_plan, tp_size=tp_size, device_mesh=device_mesh, device_map=device_map ) + # Auto-detect ``.gguf`` when the caller passes a Hub repo (or a local + # directory) that contains exactly one GGUF file — saves them from + # having to spell out ``gguf_file=`` for GGUF-only repos. + if gguf_file is None and isinstance(pretrained_model_name_or_path, str): + gguf_file = _maybe_auto_detect_gguf(pretrained_model_name_or_path, **download_kwargs_with_commit) + if gguf_file is not None and not is_accelerate_available(): raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.") @@ -4178,23 +4215,23 @@ def from_pretrained( if "experts_implementation" in kwargs: config._experts_implementation = kwargs.pop("experts_implementation") + # ``gguf_file=`` is shorthand for ``quantization_config=GgufQuantizeConfig(gguf_file=...)``. + # Building that config here lets the standard ``get_hf_quantizer`` dispatch below own + # quantizer construction (and its `validate_environment`) — no GGUF special-case after this. + if gguf_file is not None: + from .utils.quantization_config import GgufQuantizeConfig + + quantization_config = GgufQuantizeConfig.from_gguf_file( + gguf_file, quantization_config, dequantize=dtype_was_explicit + ) + hf_quantizer, config, device_map = get_hf_quantizer( config, quantization_config, device_map, weights_only, user_agent ) - if gguf_file: - if hf_quantizer is not None: - raise ValueError( - "You cannot combine Quantization and loading a model from a GGUF file, try again by making sure you did not passed a `quantization_config` or that you did not load a quantized model from the Hub." - ) - if device_map is not None and ( - (isinstance(device_map, dict) and "disk" in device_map.values()) or "disk" in device_map - ): - raise RuntimeError( - "One or more modules is configured to be mapped to disk. Disk offload is not supported for models " - "loaded from GGUF files." - ) - + # ``use_kernels`` controls the model-wide kernel resolver (attention / + # RMSNorm / ...); it's independent from the GGUF kernels, which load + # via ``integrations.gguf_kernels.ensure_metal_kernels`` regardless. if kernel_config is not None and not use_kernels: logger.warning_once( "A kernel_config was provided but use_kernels is False; setting use_kernels=True automatically. To suppress this warning, explicitly set use_kernels to True." @@ -4213,6 +4250,13 @@ def from_pretrained( tqdm_class=tqdm_class, ) + # Quantizers that own an on-disk state load (today only GGUF) get a hook + # to do it themselves — no transformers-side GGUF branch. The hook only + # fires when the quantizer has a file to load (``gguf_file=`` path); the + # on-the-fly path leaves the standard safetensors loader alone. + if gguf_file is not None and hf_quantizer is not None and hasattr(hf_quantizer, "load_checkpoint_state"): + state_dict = hf_quantizer.load_checkpoint_state(checkpoint_files[0]) + is_quantized = hf_quantizer is not None # Find the correct dtype based on current state @@ -4220,18 +4264,6 @@ def from_pretrained( dtype, checkpoint_files, config, sharded_metadata, state_dict, weights_only, hf_quantizer ) - if gguf_file: - from .modeling_gguf_pytorch_utils import load_gguf_checkpoint - - # we need a dummy model to get the state_dict - for this reason, we keep the state_dict as if it was - # passed directly as a kwarg from now on - with torch.device("meta"): - dummy_model = cls(config) - - state_dict = load_gguf_checkpoint( - checkpoint_files[0], return_tensors=True, model_to_load=dummy_model, torch_dtype=dtype - )["tensors"] - config.name_or_path = pretrained_model_name_or_path # Overwrite `config.fusion_config` if it is provided. diff --git a/src/transformers/models/auto/auto_factory.py b/src/transformers/models/auto/auto_factory.py index c6a457de9e4a..abdcc3f42575 100644 --- a/src/transformers/models/auto/auto_factory.py +++ b/src/transformers/models/auto/auto_factory.py @@ -345,11 +345,15 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], **kwargs, ) - # if torch_dtype=auto was passed here, ensure to pass it on - if kwargs_orig.get("torch_dtype", None) == "auto": - kwargs["torch_dtype"] = "auto" - if kwargs_orig.get("dtype", None) == "auto": - kwargs["dtype"] = "auto" + # Pass user-specified dtype through to the model loader. The config + # also records it (so non-quantized loads work today), but the GGUF + # path needs to distinguish "user explicitly requested a dtype" + # (→ dequantize) from "config.json happens to record a dtype" + # (→ stay native). Restore from kwargs_orig to keep that signal. + if "torch_dtype" in kwargs_orig: + kwargs["torch_dtype"] = kwargs_orig["torch_dtype"] + if "dtype" in kwargs_orig: + kwargs["dtype"] = kwargs_orig["dtype"] if kwargs_orig.get("quantization_config", None) is not None: kwargs["quantization_config"] = kwargs_orig["quantization_config"] diff --git a/src/transformers/quantizers/__init__.py b/src/transformers/quantizers/__init__.py index 7117bc2b5d80..6c70ccd986fc 100755 --- a/src/transformers/quantizers/__init__.py +++ b/src/transformers/quantizers/__init__.py @@ -13,4 +13,5 @@ # limitations under the License. from .auto import AutoHfQuantizer, AutoQuantizationConfig, register_quantization_config, register_quantizer from .base import HfQuantizer +from .quantizer_gguf import GGUFQuantizer from .quantizers_utils import get_module_from_name diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index a972c6f637bb..25ddba7aecb6 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -28,6 +28,7 @@ FineGrainedFP8Config, FourOverSixConfig, FPQuantConfig, + GgufQuantizeConfig, GPTQConfig, HiggsConfig, HqqConfig, @@ -55,6 +56,7 @@ from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer from .quantizer_fouroversix import FourOverSixHfQuantizer from .quantizer_fp_quant import FPQuantHfQuantizer +from .quantizer_gguf import GGUFQuantizer from .quantizer_gptq import GptqHfQuantizer from .quantizer_higgs import HiggsHfQuantizer from .quantizer_hqq import HqqHfQuantizer @@ -92,6 +94,7 @@ "mxfp4": Mxfp4HfQuantizer, "metal": MetalHfQuantizer, "sinq": SinqHfQuantizer, + "gguf": GGUFQuantizer, } AUTO_QUANTIZATION_CONFIG_MAPPING = { @@ -118,6 +121,7 @@ "mxfp4": Mxfp4Config, "metal": MetalConfig, "sinq": SinqConfig, + "gguf": GgufQuantizeConfig, } LOADING_ATTRIBUTES_CONFIG_TYPES = ( diff --git a/src/transformers/quantizers/quantizer_gguf.py b/src/transformers/quantizers/quantizer_gguf.py new file mode 100644 index 000000000000..9bad7c58bcc1 --- /dev/null +++ b/src/transformers/quantizers/quantizer_gguf.py @@ -0,0 +1,419 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""GGUF quantizer — mirrors the FP8 quantizer pattern. + +Three load paths: + +* `from_pretrained(repo, gguf_file=...)` on MPS: modules swap to + :class:`GgufLinear` / :class:`GgufExperts` at meta time. Bytes flow through + the standard rename pipeline via the target-aware :class:`GGUFDequantize` + (which passes uint8 bytes straight to the swapped buffers) and the experts + merge converter is rewritten in-place to per-projection renames so the + whole load is one machinery — no post-load fix-up. +* `from_pretrained(repo, gguf_file=..., dtype=...)` or non-MPS: + :class:`GGUFDequantize` dequantizes and the model loads as a standard + `nn.Linear` chain in the requested dtype. +* `from_pretrained(repo, quantization_config=GgufQuantizeConfig(...))`: + load fp16/bf16 normally; :class:`GGUFQuantize` (the + :class:`GGUFDequantize` reverse op) packs matching Linears on the fly and + the meta-time swap drops in :class:`GgufLinear` modules. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from .base import HfQuantizer + + +# `quantizers/__init__.py` is imported even from torch-less installs (PIL-only +# CI does this). Keep the top-level import set torch-free; the methods that +# actually need torch import it lazily. +if TYPE_CHECKING: + import torch # noqa: F401 + + +logger = logging.getLogger(__name__) + + +class GGUFQuantizer(HfQuantizer): + """GGUF quantizer.""" + + requires_calibration = False + + def __init__(self, quantization_config, **kwargs): + """Distinguish the three sources a GGUF-quantized model can come from and + set ``pre_quantized`` accordingly. The source is inferred from the config, + not passed explicitly: + + * **safetensors reload** — the config carries ``module_quant_types``, the + per-module swap plan recorded by a prior ``save_pretrained``. The + quantized bytes already live in the model's own safetensors shards, so + no ``.gguf`` file is involved; we only need to replay the swap. This is a + pre-quantized load (``pre_quantized=True``). + * **external gguf file** — the config has ``gguf_file=``. The source bytes + come from that ``.gguf`` and are read by :meth:`load_checkpoint_state`. + Also pre-quantized (``pre_quantized=True``). + * **on-the-fly quantization** — neither of the above. The checkpoint is a + regular fp16/bf16 model and the :class:`GGUFQuantize` conversion op packs + each matching weight into GGUF bytes at load time, so the model is *not* + pre-quantized (``pre_quantized=False``). + + The first two are mutually exclusive with the third; ``on_the_fly`` is + simply "neither a saved swap plan nor a gguf file was given". + """ + from ..utils.quantization_config import GgufQuantizeConfig + + if quantization_config is None: + quantization_config = GgufQuantizeConfig() + has_module_map = bool(getattr(quantization_config, "module_quant_types", None)) + has_gguf_file = getattr(quantization_config, "gguf_file", None) is not None + on_the_fly = not (has_module_map or has_gguf_file) + kwargs.setdefault("pre_quantized", not on_the_fly) + super().__init__(quantization_config=quantization_config, **kwargs) + # Re-annotate with the concrete type so attribute accesses type-check. + self.quantization_config: GgufQuantizeConfig = quantization_config + self.on_the_fly = on_the_fly + # `keep_quantized`: True ⇒ the model keeps GGUF byte modules (GgufLinear / + # GgufExperts, the Metal-kernel path); False ⇒ weights are dequantized to + # plain float `nn.Linear`. Resolved for real in + # `_process_model_before_weight_loading` via `_resolve_keep_quantized`; this + # is just the default for the on-the-fly case. + self.keep_quantized = on_the_fly + # GGUF file state — populated by :meth:`load_checkpoint_state` when the + # gguf_file path is active. Empty for the on-the-fly path. + self.weight_mapping: list = [] + self.gguf_tensors: dict = {} + # `{gguf_name: ggml_quant_type}` metadata read off the GGUF header (no + # tensor data). Drives the coarse all-supported decision below. + self.gguf_tensor_types: dict = {} + # Whether every quant type in the source file has a metal kernel (set in + # `load_checkpoint_state`). Gates the byte/swap path vs full dequant. + self._all_supported: bool = True + + # ---- HfQuantizer hooks ---------------------------------------------------- + + def validate_environment(self, *args, **kwargs): + """Disk offload writes dequantized fp shards to disk and reloads them as + plain tensors — incompatible with the GGUF byte-passthrough / kernel path. + Reject a `device_map` that sends any module to ``"disk"``.""" + device_map = kwargs.get("device_map") + if device_map is not None and ( + (isinstance(device_map, dict) and "disk" in device_map.values()) or "disk" in device_map + ): + raise RuntimeError( + "One or more modules is configured to be mapped to disk. Disk offload is not supported for models " + "loaded from GGUF files." + ) + + def update_weight_conversions(self, weight_conversions): + """Prepend the GGUF rename / merge / reverse-permute converters + produced from the source `.gguf` (`load_checkpoint_state` populates + `self.weight_mapping`) onto the model's existing conversion list. + + FP8's `update_weight_conversions` exists to *fuse per-block scales* + into the weight tensors before any merge/concat ops. GGUF has no + separate scale tensor — each block carries its own scale inline — so + there's nothing to fuse here. The dequant decision is binary and + made up-front in `load_checkpoint_state` (every quantized source is + either kept as bytes for the GgufLinear/GgufExperts swap, or already + dequantized at load time when `dequantize=True`). By the time this + runs, the value side is settled; we just splice the GGUF rename map + in front of the model-side conversions. + """ + return list(self.weight_mapping) + list(weight_conversions) + + def get_quantize_ops(self): + """On-the-fly path: the loader calls this when `param_needs_quantization` + returns True. Returns the :class:`GGUFQuantize` op that packs a live + fp16/bf16 weight into GGUF block bytes.""" + from ..gguf_conversion_ops import GGUFQuantize + + return GGUFQuantize(self) + + def param_needs_quantization(self, model, param_name, **kwargs) -> bool: + """Only the on-the-fly path needs per-param quantization (re-packing the + incoming fp16/bf16 weight into GGUF bytes). The pre-quantized path leans + on the rename pipeline + the post-load byte-copy.""" + if not self.on_the_fly: + return False + import torch.nn as nn + + from ..integrations.gguf_linear import GgufLinear + + target_path = param_name.rsplit(".", 1)[0] + try: + module = model.get_submodule(target_path) + except AttributeError: + return False + # After `_process_model_before_weight_loading` swapped the matching + # Linears, those modules are GgufLinear and need quantization. Anything + # else (norms, lm_head, embeddings) is left alone. + return isinstance(module, GgufLinear) and not param_name.endswith(".bias") and isinstance(module, nn.Module) + + def load_checkpoint_state(self, gguf_path: str): + """Pre-load the `.gguf` file's tensors + rename map + KV metadata. + + Called by :func:`PreTrainedModel.from_pretrained` for the gguf_file path + after :func:`_get_resolved_checkpoint_files` resolves the file. Returns + the state-dict that the standard loader needs — GGUF isn't a + safetensors format, so this hook gives the quantizer ownership of the + on-disk → in-memory state load. + + Dequant decision lives here, not in a conversion op: + + * `dequantize=True` (or non-MPS): dequant every quantized tensor + up-front. Modules stay as plain `nn.Linear` / `nn.Embedding` + with fp weights. No op in the conversion chain. + * Default (MPS swap path): keep quantized tensors as + `GGUFQuantizedTensor` raw bytes when their quant_type has a + Metal kernel (`_QUANT_INFO`) — they flow into `GgufLinear` / + `GgufExperts` uint8 buffers. Anything else still has to be + pre-dequanted: embedding-bound tensors (their HF target is + `nn.Embedding`, not a swappable Linear), and any quant_type + without a Metal kernel (IQ1_S, IQ2_*, IQ3_*, Q2_K, Q3_K, Q4_1) + — for those the swap step skips the layer and the bytes would + otherwise be assigned raw into a plain `nn.Linear.weight`. + """ + import torch + + from ..integrations.gguf_dequant import GGUFQuantizedTensor, dequantize_gguf_tensor + from ..integrations.gguf_linear import _quant_type_name, gguf_linear_supports + from ..modeling_gguf_pytorch_utils import load_gguf_checkpoint + + parsed = load_gguf_checkpoint(gguf_path, return_tensors=True) + self.weight_mapping = list(parsed.get("weight_mapping", []) or []) + self.gguf_tensor_types = dict(parsed.get("tensor_quant_types", {}) or {}) + tensors = parsed["tensors"] + # GGUF tensor names whose HF target is an `nn.Embedding` need fp values at + # assignment time. `token_embd` is universal; extend as more archs land. + embedding_prefixes = ("token_embd",) + _FLOAT_TYPES = {"F32", "F16", "BF16", "F64"} + # Coarse, metadata-only decision: can the whole model take the byte/kernel + # path? Only if every quantized (non-float, non-embedding) tensor has a + # metal kernel — checked off the header types, no rename, no tensor load. + # If any unsupported quant is present (e.g. Q2_K, IQ1_S), the whole model + # dequantizes; otherwise we keep bytes and swap uniformly. (Per-module + # support is unnecessary: a uniform swap reverts float-source Linears and + # the kernels cover everything else.) + self._all_supported = all( + gguf_linear_supports(qt) + for name, qt in self.gguf_tensor_types.items() + if name.split(".")[0] not in embedding_prefixes and _quant_type_name(qt) not in _FLOAT_TYPES + ) + # The Gemma/Nemotron `w + 1` norm de-offset is handled by the `SubtractOne` + # converter in fp32 — `_mark_offset_norms_fp32` keeps those norms fp32. + dequant_all = ( + self.quantization_config.dequantize or not torch.backends.mps.is_available() or not self._all_supported + ) + for name, t in list(tensors.items()): + if not isinstance(t, GGUFQuantizedTensor): + continue + # Already-float sources (F32 / F16 / BF16): drop the `GGUFQuantizedTensor` + # subclass so downstream ops see a plain tensor (its `__torch_function__` + # re-wraps on `.to`, which on MPS corrupts mmap-backed storage during a + # device cast). These flow into the swap as-is and get reverted to a + # plain `nn.Linear` post-load if their target was swapped. + if t.is_floating_point(): + tensors[name] = t.as_subclass(torch.Tensor) + continue + # Quantized: dequantize embeddings (never swapped) and everything on the + # dequant path; otherwise keep the raw bytes for the GgufLinear swap. + if dequant_all or name.split(".")[0] in embedding_prefixes: + tensors[name] = dequantize_gguf_tensor(t, t.quant_type, device=t.device) + self.gguf_tensors = tensors + return tensors + + def _resolve_keep_quantized(self) -> bool: + """Decide whether the model keeps GGUF byte modules (GgufLinear / + GgufExperts, the Metal-kernel path) or is dequantized to plain float + `nn.Linear`. Single source of truth for the three load paths: + + * safetensors reload (`module_quant_types` present) → always keep: + the saved uint8 bytes only fit the swapped buffers. + * on-the-fly (`GgufQuantizeConfig` without a file) → always keep: + the config explicitly asked to quantize. + * gguf_file= → keep only on MPS, without an explicit `dequantize`, and + only if every quant type in the file has a kernel (`_all_supported`). + """ + import torch + + if bool(getattr(self.quantization_config, "module_quant_types", None)): + return True + if self.on_the_fly: + return True + return ( + torch.backends.mps.is_available() and not self.quantization_config.dequantize and self._all_supported + ) + + def _mark_offset_norms_fp32(self, model) -> None: + """Gemma/Nemotron store RMSNorm weights as `w + 1`; the `SubtractOne` + converter de-offsets them, but only correctly in fp32 — the loader casts + to the target dtype *before* the converter chain runs. Force exactly those + norms to stay fp32 for the load through `_keep_in_fp32_modules_strict` + (which covers both fp16 and bf16). The set of norms is the targets of the + `SubtractOne` converters in `weight_mapping` — no hardcoded arch list, so a + new arch works as soon as its converter graph declares the offset. Runs on + every load path (the dequant path needs it too). Must run before + `_get_dtype_plan` reads the flag — `preprocess_model` is called first. + """ + from ..core_model_loading import WeightConverter + from ..gguf_conversion_ops import SubtractOne + + offset_targets = { + pat + for entry in self.weight_mapping + if isinstance(entry, WeightConverter) and any(isinstance(op, SubtractOne) for op in entry.operations) + for pat in entry.target_patterns + } + if offset_targets: + model._keep_in_fp32_modules_strict = set(model._keep_in_fp32_modules_strict or []) | offset_targets + + def _process_model_before_weight_loading(self, model, **kwargs): + """Swap nn.Linear / fused-expert modules to GgufLinear / GgufExperts at + meta time, when `_resolve_keep_quantized` selects the byte path. + + The swap is **uniform** — every Linear becomes a placeholder, with no + source→target rename and no per-module quant type (those are bound in + `postprocess_model` from the bytes that arrive). The on-the-fly path is the + one exception: its quant type is a single config value, known up front. + """ + self._mark_offset_norms_fp32(model) + self.keep_quantized = self._resolve_keep_quantized() + if not self.keep_quantized: + return + from ..integrations.gguf_linear import replace_with_gguf_linear + + cfg = self.quantization_config + skip = set(cfg.modules_to_not_convert or []) + # Tied output embeddings (e.g. lm_head on tie_word_embeddings models) have + # no standalone quantized tensor in the GGUF — they reuse `token_embd`, which + # is dequantized. Keep them as plain modules so the tie lands a float weight, + # rather than swapping them to an unfilled GgufLinear. + tied = getattr(model, "all_tied_weights_keys", None) or {} + skip |= {k.removesuffix(".weight") for k in (set(tied) | set(tied.values()))} + if self.on_the_fly: + # Live fp model: uniform config quant type, filtered to modules_to_convert; + # `GGUFQuantize` packs the weights at load, so bind kernels at meta. + replace_with_gguf_linear( + model, modules_to_not_convert=skip, modules_to_convert=cfg.modules_to_convert, quant_type=cfg.quant_type + ) + else: + # gguf_file / safetensors reload: load the raw bytes/floats without dtype + # coercion (so a uint8 buffer doesn't truncate a float router), then bind + # or revert per module in `postprocess_model`. + self.pre_quantized = True + replace_with_gguf_linear(model, modules_to_not_convert=skip) + + def postprocess_model(self, model, **kwargs): + """Run the base post-process, then — on the swap path — reconcile each + swapped module against the bytes that loaded (bind kernels / revert float + sources) and apply the generation defaults.""" + super().postprocess_model(model, **kwargs) + if self.keep_quantized: + self._reconcile_swapped_modules(model) + self._apply_generation_defaults(model) + return model + + @property + def is_trainable(self): + return True + + def is_serializable(self, safe_serialization: bool = True) -> bool: + # save_pretrained → safetensors works: GgufLinear / GgufExperts buffers + # are uint8, biases / norms are fp32, all native safetensors dtypes. + # The swap layout is replayed on reload via + # `GgufQuantizeConfig.module_quant_types`, recorded in `postprocess_model`. + return True + + # ---- Internals ------------------------------------------------------------ + + def _reconcile_swapped_modules(self, model) -> None: + """Post-load reconciliation for the uniform swap. For every placeholder: + + * bind the metal kernels from the quant type that arrived with the bytes + (a `GGUFQuantizedTensor` carries it; the gate/up merge keeps it via + `torch.cat`), or from the saved `module_quant_types` on a safetensors + reload (where the bytes are plain uint8 with no carried type); + * revert any `GgufLinear` whose GGUF source was actually float (e.g. MoE + routers, stored F32) back to a plain `nn.Linear`. + + Records the resolved per-module quant types on the config so a + `save_pretrained` round-trip can rebind on reload. This replaces the old + rename-based `_build_quant_info` — the swap plan now falls out of the + bytes, not a source→target name map. + """ + from ..integrations.gguf_linear import GgufExperts, GgufLinear, _quant_type_name, gguf_linear_supports + + saved = dict(getattr(self.quantization_config, "module_quant_types", None) or {}) + recorded: dict = {} + for name, mod in list(model.named_modules()): + if isinstance(mod, GgufLinear): + quant_type = mod.quant_type or saved.get(name) or _quant_type_name(getattr(mod.weight, "quant_type", None)) + if gguf_linear_supports(quant_type): + if mod._mv_op is None: + mod.bind_after_load(quant_type) + recorded[name] = mod.quant_type + else: + # Float source (router) or unsupported → plain nn.Linear. + self._revert_to_linear(model, name, mod) + elif isinstance(mod, GgufExperts): + if mod._id_op_gate_up is None: + info = saved.get(name) or {} + mod.bind_after_load(info.get("gate_up_quant"), info.get("down_quant")) + recorded[name] = {"gate_up_quant": mod.gate_up_quant, "down_quant": mod.down_quant} + self.quantization_config.module_quant_types = recorded + # The swapped model is now a standard HF module tree holding uint8 buffers; + # its "original format" is GGUF, which can't round-trip through safetensors + # (a plain reload has no rename map). Drop the GGUF load-renames so + # `save_pretrained` writes HF names — reload re-swaps uniformly and rebinds + # from `module_quant_types`. + model._weight_conversions = None + + @staticmethod + def _revert_to_linear(model, name: str, mod) -> None: + """Swap a uniformly-created `GgufLinear` whose source turned out to be float + back to a plain `nn.Linear` holding the loaded float weight.""" + import torch + import torch.nn as nn + + weight = mod.weight + new = nn.Linear( + mod.in_features, mod.out_features, bias=mod.bias is not None, device=weight.device, dtype=weight.dtype + ) + with torch.no_grad(): + new.weight.copy_(weight.reshape(new.weight.shape)) + if mod.bias is not None: + new.bias.copy_(mod.bias) + new.weight.requires_grad_(False) + if new.bias is not None: + new.bias.requires_grad_(False) + model.set_submodule(name, new) + + def _apply_generation_defaults(self, model) -> None: + """Set `cache_implementation` + `compile_config` defaults on + `model.generation_config` to leverage the best perf when running GGUF Models""" + from ..generation.configuration_utils import CompileConfig + + gc = getattr(model, "generation_config", None) + if gc is None: + return + changed = [] + if getattr(gc, "cache_implementation", None) is None: + gc.cache_implementation = "static" + changed.append('cache_implementation="static"') + if getattr(gc, "compile_config", None) is None: + gc.compile_config = CompileConfig(mode="reduce-overhead", dynamic=False) + changed.append('compile_config=CompileConfig(mode="reduce-overhead", dynamic=False)') + if changed: + logger.info( + "GGUFQuantizer: defaulting generation_config to %s for the GgufLinear kernel path. " + "Pass values explicitly on `model.generation_config` to override.", + ", ".join(changed), + ) diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index a6c1f5334516..c7813ff0ce69 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -63,6 +63,7 @@ class QuantizationMethod(str, Enum): METAL = "metal" FOUR_OVER_SIX = "fouroversix" SINQ = "sinq" + GGUF = "gguf" class AwqFormat(str, Enum): @@ -1091,6 +1092,131 @@ def post_init(self): raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights}") +class GgufQuantizeConfig(QuantizationConfigMixin): + """Quantize specific modules of a model into GGUF block bytes. Used either + implicitly when `from_pretrained` is called with a `gguf_file=` + argument, or explicitly to convert a fp16/bf16 model on the fly: + + ```python + AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen1.5-MoE-A2.7B", + quantization_config=GgufQuantizeConfig( + quant_type="Q4_0", + modules_to_convert=["model.layers.*.self_attn.q_proj", ...], + ), + ) + ``` + + Args: + quant_type (`str`, *optional*, defaults to `"Q4_0"`): + GGUF quantization name. The loader path (`gguf_file=` or safetensors + reload via `module_quant_types`) supports the full kernel suite — + Q4_0 / Q5_0 / Q5_1 / Q8_0 / Q4_K / Q5_K / Q6_K / IQ4_NL / IQ4_XS. + The on-the-fly quantize path is gated by gguf-py's own quantizer, + which only ships Python implementations for Q4_0 and Q8_0 today + (K-quants and IQ4_* are read-only upstream). + modules_to_convert (`list[str]`, *optional*): + Glob-style names of `nn.Linear` modules to quantize. `None` means every + quantizable `nn.Linear` is converted. + modules_to_not_convert (`list[str]`, *optional*): + Explicit skip list (substring match against module names). + gguf_file (`str`, *optional*): + Path or repo-relative filename of the source `.gguf` to load. Set + automatically by `from_pretrained` when the caller passes + `gguf_file=`. When set, the quantizer pre-loads the file from + `_get_resolved_checkpoint_files` instead of running the on-the-fly + quantize path. + dequantize (`bool`, *optional*, defaults to `False`): + Force the dequantize-on-load path even with a GGUF source. Set + automatically when `from_pretrained` receives an explicit `dtype=`, + so the caller-requested precision wins over native GGUF kernels. + module_quant_types (`dict[str, str | dict[str, str]]`, *optional*): + Per-module quant info, populated by the quantizer at swap time and + replayed on `from_pretrained` so a `save_pretrained` → + `from_pretrained` round-trip reconstructs the same GGUF modules + from safetensors (no .gguf file needed). For `nn.Linear`: + `{name: "Q4_K"}`. For fused-expert modules: `{name: {"gate_up_quant": + "Q4_K", "down_quant": "Q8_0"}}`. Mixed-quant GGUFs (Q4_K_M etc.) need + this to round-trip — the on-the-fly path can also rely on it to + avoid re-walking the model on reload. + """ + + def __init__( + self, + quant_type: str = "Q4_0", + modules_to_convert: list[str] | None = None, + modules_to_not_convert: list[str] | None = None, + gguf_file: str | None = None, + dequantize: bool = False, + module_quant_types: dict | None = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.GGUF + self.quant_type = quant_type + self.modules_to_convert = modules_to_convert + self.modules_to_not_convert = modules_to_not_convert + self.gguf_file = gguf_file + self.dequantize = dequantize + self.module_quant_types = dict(module_quant_types) if module_quant_types else None + self.post_init() + + # All quant types the GGUF metal kernels handle. The on-the-fly path is + # narrower (gguf-py limit) and validated lazily inside `GGUFQuantize.convert`. + _SUPPORTED_QUANT_TYPES = ( + "Q4_0", + "Q5_0", + "Q5_1", + "Q8_0", + "Q4_K", + "Q5_K", + "Q6_K", + "IQ4_NL", + "IQ4_XS", + ) + + def post_init(self): + if self.quant_type not in self._SUPPORTED_QUANT_TYPES: + raise ValueError( + f"GgufQuantizeConfig.quant_type must be one of {self._SUPPORTED_QUANT_TYPES}, got {self.quant_type!r}" + ) + + def to_dict(self): + """Serialise for `config.json`. `gguf_file` is a load-time pointer + to an external file and shouldn't survive into the saved config — once + the model is saved as safetensors the bytes live there. Reload picks + up the swap layout via :attr:`module_quant_types`.""" + out = super().to_dict() + out.pop("gguf_file", None) + return out + + @classmethod + def from_gguf_file(cls, gguf_file: str, quantization_config=None, *, dequantize: bool = False): + """Resolve the config backing a ``from_pretrained(..., gguf_file=...)`` load. + + ``gguf_file=`` is just shorthand for + ``quantization_config=GgufQuantizeConfig(gguf_file=...)``; centralising the + construction here keeps `from_pretrained` free of GGUF special-cases. Three + inputs are accepted: + + * ``quantization_config is None`` → build a fresh config pointing at ``gguf_file``. + * an existing :class:`GgufQuantizeConfig` → fill in ``gguf_file`` / ``dequantize`` + where the caller left them unset (an explicit config wins otherwise). + * any other quant config → error; the two can't be combined. + + ``dequantize`` forces the dequantize-on-load path. `from_pretrained` passes + ``dtype is not None`` here so an explicit ``dtype=`` overrides the native kernels. + """ + if quantization_config is None: + return cls(gguf_file=gguf_file, dequantize=dequantize) + if isinstance(quantization_config, cls): + if quantization_config.gguf_file is None: + quantization_config.gguf_file = gguf_file + if dequantize: + quantization_config.dequantize = True + return quantization_config + raise ValueError("Cannot combine `gguf_file=` with a non-GGUF `quantization_config=`. Drop one of them.") + + class CompressedTensorsConfig(QuantizationConfigMixin): """ This is a wrapper class that handles compressed-tensors quantization config options. diff --git a/tests/quantization/ggml/test_ggml.py b/tests/quantization/ggml/test_ggml.py index aa5cdbc7adc6..139bdfe34c4f 100644 --- a/tests/quantization/ggml/test_ggml.py +++ b/tests/quantization/ggml/test_ggml.py @@ -841,7 +841,7 @@ def test_nemotron_weights_conversion_fp16(self): for layer_name, original_params in original_state_dict.items(): if layer_name in converted_state_dict: self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape) - torch.testing.assert_close(original_params, converted_state_dict[layer_name]) + torch.testing.assert_close(original_params, converted_state_dict[layer_name].to(original_params.dtype)) else: raise ValueError(f"Layer {layer_name} is not presented in GGUF model") @@ -919,7 +919,7 @@ def test_gemma2_weights_conversion_fp32(self): for layer_name, original_params in original_state_dict.items(): if layer_name in converted_state_dict: self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape) - torch.testing.assert_close(original_params, converted_state_dict[layer_name]) + torch.testing.assert_close(original_params, converted_state_dict[layer_name].to(original_params.dtype)) else: raise ValueError(f"Layer {layer_name} is not presented in GGUF model") @@ -958,7 +958,7 @@ def test_gemma3_text_weights_conversion_bf16(self): for layer_name, original_params in original_state_dict.items(): if layer_name in converted_state_dict: self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape) - torch.testing.assert_close(original_params, converted_state_dict[layer_name]) + torch.testing.assert_close(original_params, converted_state_dict[layer_name].to(original_params.dtype)) else: raise ValueError(f"Layer {layer_name} is not presented in GGUF model") @@ -983,7 +983,7 @@ def test_gemma3_vision_weights_conversion_bf16(self): for layer_name, original_params in original_state_dict.items(): if layer_name in converted_state_dict: self.assertTrue(original_params.shape == converted_state_dict[layer_name].shape) - torch.testing.assert_close(original_params, converted_state_dict[layer_name]) + torch.testing.assert_close(original_params, converted_state_dict[layer_name].to(original_params.dtype)) else: raise ValueError(f"Layer {layer_name} is not presented in GGUF model") diff --git a/tests/quantization/ggml/test_gguf_arch_coverage.py b/tests/quantization/ggml/test_gguf_arch_coverage.py new file mode 100644 index 000000000000..bef5f658a317 --- /dev/null +++ b/tests/quantization/ggml/test_gguf_arch_coverage.py @@ -0,0 +1,401 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Fast (non-slow) tests pinning the GGUF integration's *coverage matrix* — +every model_type the loader claims to support, every quant_type the metal +kernels claim to handle, every MoE arch routed through the experts swap. + +End-to-end per-arch generation tests live in `test_ggml.py` (slow CI). +The tests here only inspect the registries + the quantizer's contract with +the loading pipeline; they download nothing and run on every PR. +""" + +from __future__ import annotations + +import unittest + +from parameterized import parameterized + +from transformers.testing_utils import require_gguf, require_torch +from transformers.utils import is_torch_available + + +# HF model_types the public GGUF integration claims to load (mirrors the per-arch +# coverage in `test_ggml.py`). If you add a new entry to `_GGUF_ARCH_CONVERTERS`, +# add it here too so the coverage assertion catches accidental removal later. +EXPECTED_MODEL_TYPES = sorted( + { + # Llama / RoPE family + "llama", + "mistral", + "phi3", + "cohere", + "qwen2", + "qwen3", + "deci", + "stablelm", + "starcoder2", + # Norm-subtract-one variants + "nemotron", + "gemma2", + "gemma3", + "gemma3_text", + # MoE families + "qwen2_moe", + "qwen3_moe", + "minimax_m2", + "gpt_oss", + # Misc encoder/decoder & misc archs + "bloom", + "gpt2", + "mamba", + "lfm2", + "falcon", + "t5", + "t5encoder", + "umt5", + } +) + +# Quant types `GgufLinear` / `GgufExperts` accept on the metal fast path. The +# canonical source is `_QUANT_INFO` in `integrations.gguf_linear`; we re-list +# here so a drift between the kernel-side table and what tests expect is loud. +EXPECTED_KERNEL_QUANT_TYPES = ( + "Q4_0", + "Q5_0", + "Q5_1", + "Q8_0", + "Q4_K", + "Q5_K", + "Q6_K", + "IQ4_NL", + "IQ4_XS", +) + + +@require_torch +@require_gguf +class GgufArchCoverageTests(unittest.TestCase): + """Registry / well-formedness checks — no downloads.""" + + def test_every_expected_model_type_is_registered(self): + from transformers.modeling_gguf_pytorch_utils import _GGUF_ARCH_CONVERTERS + + missing = sorted(set(EXPECTED_MODEL_TYPES) - set(_GGUF_ARCH_CONVERTERS)) + self.assertFalse( + missing, + f"Model types previously supported by the GGUF loader are no longer registered " + f"in `_GGUF_ARCH_CONVERTERS`: {missing}. Re-add the entry (or alias it onto an " + f"existing converter list) so the loader keeps the same coverage.", + ) + + def test_no_unknown_model_types_silently_added(self): + """Catches the reverse: a new entry sneaking into `_GGUF_ARCH_CONVERTERS` + without being added to `EXPECTED_MODEL_TYPES`. Forces the test author of + the new arch to acknowledge the coverage matrix.""" + from transformers.modeling_gguf_pytorch_utils import _GGUF_ARCH_CONVERTERS + + unexpected = sorted(set(_GGUF_ARCH_CONVERTERS) - set(EXPECTED_MODEL_TYPES)) + self.assertFalse( + unexpected, + f"New entries in `_GGUF_ARCH_CONVERTERS` that aren't in EXPECTED_MODEL_TYPES: " + f"{unexpected}. Add them to EXPECTED_MODEL_TYPES (and to `test_ggml.py` if " + f"they're not already covered there).", + ) + + @parameterized.expand([(m,) for m in EXPECTED_MODEL_TYPES]) + def test_arch_entry_is_well_formed(self, model_type: str): + from transformers.core_model_loading import WeightConverter, WeightRenaming + from transformers.modeling_gguf_pytorch_utils import get_gguf_converters + + rules = get_gguf_converters(model_type) + self.assertGreater(len(rules), 0, f"{model_type}: empty converter list") + for rule in rules: + self.assertIsInstance( + rule, + (WeightRenaming, WeightConverter), + f"{model_type}: every entry must be a WeightRenaming or WeightConverter, got {type(rule).__name__}", + ) + + +@require_torch +@require_gguf +class GgufQuantizerWiringTests(unittest.TestCase): + """The quantizer is FP8-shaped: it owns the dequant decision in + `load_checkpoint_state` and does NOT inject a `GGUFDequantize` op into the + conversion chain. `update_weight_conversions` is a pure splice that puts + the GGUF rename map in front of the model-side conversions.""" + + def test_update_weight_conversions_injects_nothing(self): + """Mirror of `Fp8Quantizer.update_weight_conversions` semantics: no + scale-fusion / dequant op gets prepended on the GGUF side (GGUF blocks + carry their own scale; no fusion needed).""" + from transformers import GgufQuantizeConfig + from transformers.core_model_loading import WeightConverter, WeightRenaming + from transformers.gguf_conversion_ops import GGUFDequantize, Unsqueeze + from transformers.quantizers.quantizer_gguf import GGUFQuantizer + + # Fake the rename map (no .gguf file needed for this hook). A + # WeightConverter must carry at least one op, so use a harmless Unsqueeze. + fake_mapping = [ + WeightRenaming(source_patterns=r"^foo$", target_patterns="bar"), + WeightConverter(source_patterns="x", target_patterns="y", operations=[Unsqueeze(0)]), + ] + q = GGUFQuantizer(GgufQuantizeConfig()) + q.weight_mapping = fake_mapping + out = q.update_weight_conversions([]) + # Output is the rename map untouched (identity splice). + self.assertEqual(len(out), 2) + for orig, new in zip(fake_mapping, out): + self.assertIs(orig, new) + # No conversion op was prepended into either rule. + for rule in out: + if isinstance(rule, WeightConverter): + for op in rule.operations: + self.assertNotIsInstance( + op, + GGUFDequantize, + f"{rule}: GGUFDequantize should never be auto-prepended — " + f"the dequant decision lives in load_checkpoint_state.", + ) + + +@require_torch +@require_gguf +class GgufExpertsRegistryTests(unittest.TestCase): + """MoE archs with GGUF rules either ship a `GgufExperts` subclass in + `MODEL_TYPE_TO_GGUF_EXPERTS` (byte-passthrough path) or are explicitly + excluded with a reason in the registry source.""" + + # Archs that are MoE but intentionally NOT in MODEL_TYPE_TO_GGUF_EXPERTS, + # with the reason captured here so the exclusion is reviewable. + _MOE_ARCHS_EXCLUDED = { + "gpt_oss": ( + "Transposed `(in, out)` gate_up layout + per-expert biases. Ships no " + "transposed-aware mul_mat_id kernel today; falls back to bf16 via " + "load_checkpoint_state(dequantize=True) and the stock `GptOssExperts` " + "+ `batched_mm_experts_forward(is_transposed=True, has_bias=True)`." + ), + } + + def test_moe_archs_have_experts_class_or_documented_exclusion(self): + if not is_torch_available(): + self.skipTest("requires torch") + from transformers.integrations.gguf_linear import MODEL_TYPE_TO_GGUF_EXPERTS + + moe_archs = sorted(m for m in EXPECTED_MODEL_TYPES if "moe" in m or m in ("minimax_m2", "gpt_oss")) + for arch in moe_archs: + with self.subTest(arch=arch): + if arch in self._MOE_ARCHS_EXCLUDED: + self.assertNotIn(arch, MODEL_TYPE_TO_GGUF_EXPERTS) + continue + self.assertIn( + arch, + MODEL_TYPE_TO_GGUF_EXPERTS, + f"{arch}: MoE arch has GGUF rename rules but no entry in " + f"MODEL_TYPE_TO_GGUF_EXPERTS. Add a class (or document the exclusion).", + ) + + +@require_torch +@require_gguf +class GgufQuantTypeCoverageTests(unittest.TestCase): + """Every quant_type the kernels claim to support must have both a + `_KERNEL_FMT` mapping (so module __init__ can resolve the metal op name) + AND a working torch dequant kernel (so the dequant path can fall back + cleanly on CPU / CUDA).""" + + def test_quant_info_covers_expected_types(self): + from transformers.integrations.gguf_linear import _QUANT_INFO + + missing = sorted(set(EXPECTED_KERNEL_QUANT_TYPES) - set(_QUANT_INFO)) + self.assertFalse(missing, f"`_QUANT_INFO` is missing: {missing}") + + @parameterized.expand([(qt,) for qt in EXPECTED_KERNEL_QUANT_TYPES]) + def test_kernel_fmt_entry_exists(self, quant_type: str): + from transformers.integrations.gguf_kernels import _KERNEL_FMT + + self.assertIn( + quant_type, + _KERNEL_FMT, + f"{quant_type}: missing from `_KERNEL_FMT`; module __init__ will fail to resolve `mul_mat__f32`.", + ) + + @parameterized.expand([(qt,) for qt in EXPECTED_KERNEL_QUANT_TYPES]) + def test_torch_dequant_kernel_exists(self, quant_type: str): + """The CPU dequant kernel must exist for the same set of types, so + non-MPS hosts (and `dequantize=True`) have a path.""" + import gguf + + from transformers.integrations.gguf_dequant import supported_quant_types + + ggml_type = getattr(gguf.GGMLQuantizationType, quant_type) + self.assertIn( + ggml_type, + supported_quant_types(), + f"{quant_type}: no torch dequant kernel registered in `_DISPATCH`.", + ) + + +@require_torch +@require_gguf +class GgufLoadCheckpointStateDequantDecisionTests(unittest.TestCase): + """The dequant decision in `load_checkpoint_state` is what makes the + no-op-in-update_weight_conversions design work. Pin its branches: + + * `dequantize=True` → every quantized tensor is dequanted at load time. + * Default + MPS available → only embedding-bound tensors are pre-dequanted; + the rest stay as `GGUFQuantizedTensor` for the swap path. + * Default + no MPS → falls back to full dequant. + """ + + def _run(self, monkeypatch, dequantize: bool, mps_available: bool): + import torch + + from transformers import GgufQuantizeConfig + from transformers.integrations.gguf_dequant import GGUFQuantizedTensor + from transformers.quantizers.quantizer_gguf import GGUFQuantizer + + # Fake a `load_gguf_checkpoint` return: a couple of quantized tensors + # under realistic names + one already-fp norm. + bytes32 = torch.zeros(32, dtype=torch.uint8) + + def fake_quant_tensor(): + import gguf + + t = GGUFQuantizedTensor(bytes32.clone(), quant_type=gguf.GGMLQuantizationType.Q4_0) + return t + + fake_tensors = { + "token_embd.weight": fake_quant_tensor(), # embedding → must dequant + "blk.0.attn_q.weight": fake_quant_tensor(), # linear → swap-path target + "blk.0.attn_norm.weight": torch.ones(8, dtype=torch.float32), # already fp + } + + def fake_load_gguf_checkpoint(_path, return_tensors=True): + return {"tensors": dict(fake_tensors), "weight_mapping": []} + + # Patch the loader + the dequant kernel so we don't need a real GGUF file + # or torch ops. + import transformers.integrations.gguf_dequant as _dq + import transformers.modeling_gguf_pytorch_utils as _mod + + monkeypatch(_mod, "load_gguf_checkpoint", fake_load_gguf_checkpoint) + called = {"n": 0} + + def fake_dequantize(t, quant_type, device=None): + called["n"] += 1 + return torch.zeros(8, 16, dtype=torch.float32) + + monkeypatch(_dq, "dequantize_gguf_tensor", fake_dequantize) + # `load_checkpoint_state` gates the byte-passthrough path purely on MPS + # availability (the Metal kernels are resolved later, when GgufLinear / + # GgufExperts are built), so simulating MPS is enough here. + monkeypatch(torch.backends.mps, "is_available", lambda: mps_available) + + q = GGUFQuantizer(GgufQuantizeConfig(dequantize=dequantize)) + out = q.load_checkpoint_state("/does/not/exist") + return out, called["n"] + + def _monkeypatcher(self): + """Pure-stdlib monkeypatch: save the original, restore via addCleanup.""" + + def patch(obj, name, value): + sentinel = object() + original = getattr(obj, name, sentinel) + setattr(obj, name, value) + self.addCleanup(setattr, obj, name, original) if original is not sentinel else self.addCleanup( + delattr, obj, name + ) + + return patch + + def test_dequant_true_dequants_everything(self): + out, n = self._run(self._monkeypatcher(), dequantize=True, mps_available=True) + self.assertEqual(n, 2, "both quantized tensors should have been dequanted up-front") + for key in ("token_embd.weight", "blk.0.attn_q.weight"): + from transformers.integrations.gguf_dequant import GGUFQuantizedTensor + + self.assertNotIsInstance(out[key], GGUFQuantizedTensor, key) + + def test_default_mps_keeps_linear_targets_as_bytes(self): + out, n = self._run(self._monkeypatcher(), dequantize=False, mps_available=True) + self.assertEqual(n, 1, "only token_embd should be pre-dequanted on the MPS swap path") + from transformers.integrations.gguf_dequant import GGUFQuantizedTensor + + self.assertNotIsInstance(out["token_embd.weight"], GGUFQuantizedTensor) + self.assertIsInstance(out["blk.0.attn_q.weight"], GGUFQuantizedTensor) + + def test_default_no_mps_falls_back_to_full_dequant(self): + out, n = self._run(self._monkeypatcher(), dequantize=False, mps_available=False) + self.assertEqual(n, 2, "no MPS → dequant everything so CPU / CUDA can run as plain fp") + from transformers.integrations.gguf_dequant import GGUFQuantizedTensor + + for key in ("token_embd.weight", "blk.0.attn_q.weight"): + self.assertNotIsInstance(out[key], GGUFQuantizedTensor, key) + + +@require_torch +@require_gguf +class GgufUniformSwapTests(unittest.TestCase): + """The swap is uniform (no rename, no per-module plan): placeholders are sized + from the Q4_K_M average, the kernel path is gated by a coarse all-supported + check, and quant types are bound (or float sources reverted) post-load.""" + + def test_bytes_per_row_exact_and_average(self): + from transformers.integrations.gguf_linear import _AVG_BYTES_PER_ELEM, _bytes_per_row + + # Known quant type → exact width (Q4_0 = 18 bytes per 32 elems). + self.assertEqual(_bytes_per_row(64, "Q4_0"), (64 // 32) * 18) + # Unknown (uniform-swap placeholder) → Q4_K_M average. + self.assertEqual(_bytes_per_row(256, None), max(1, round(256 * _AVG_BYTES_PER_ELEM))) + + def test_unsupported_quant_forces_full_dequant(self): + """A file containing a quant type with no metal kernel (Q2_K) must clear + `_all_supported`, so the whole model dequantizes even on (simulated) MPS.""" + import gguf + import torch + + import transformers.integrations.gguf_dequant as _dq + import transformers.modeling_gguf_pytorch_utils as _mod + from transformers import GgufQuantizeConfig + from transformers.integrations.gguf_dequant import GGUFQuantizedTensor + from transformers.quantizers.quantizer_gguf import GGUFQuantizer + + def patch(obj, name, value): + sentinel = object() + orig = getattr(obj, name, sentinel) + setattr(obj, name, value) + self.addCleanup(setattr, obj, name, orig) if orig is not sentinel else self.addCleanup( + delattr, obj, name + ) + + types = {"blk.0.attn_q.weight": gguf.GGMLQuantizationType.Q2_K} # unsupported + + def fake_load(_path, return_tensors=True): + tensors = {n: GGUFQuantizedTensor(torch.zeros(32, dtype=torch.uint8), quant_type=qt) for n, qt in types.items()} + return {"tensors": tensors, "weight_mapping": [], "tensor_quant_types": dict(types)} + + patch(_mod, "load_gguf_checkpoint", fake_load) + called = {"n": 0} + + def fake_dequant(t, quant_type, device=None): + called["n"] += 1 + return torch.zeros(8, 16, dtype=torch.float32) + + patch(_dq, "dequantize_gguf_tensor", fake_dequant) + patch(torch.backends.mps, "is_available", lambda: True) + + q = GGUFQuantizer(GgufQuantizeConfig()) + out = q.load_checkpoint_state("/does/not/exist") + self.assertFalse(q._all_supported) + self.assertEqual(called["n"], 1, "unsupported quant must be dequantized despite MPS") + self.assertNotIsInstance(out["blk.0.attn_q.weight"], GGUFQuantizedTensor) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/quantization/ggml/test_gguf_linear.py b/tests/quantization/ggml/test_gguf_linear.py new file mode 100644 index 000000000000..66bbb47a6223 --- /dev/null +++ b/tests/quantization/ggml/test_gguf_linear.py @@ -0,0 +1,247 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Tests for the slim GGUF integration. + +Covers: +* :class:`GgufLinear` construction + ``state_dict`` / ``load_state_dict`` round-trip +* :func:`replace_with_gguf_linear` meta-time swap mechanics +* ``MODEL_TYPE_TO_GGUF_EXPERTS`` registry lookup +* the public load APIs: + - ``from_pretrained(..., gguf_file=..., dtype=torch.bfloat16)`` → dequant path + - ``from_pretrained(..., quantization_config=GgufQuantizeConfig(...))`` → on-the-fly swap +* generation-config defaults set by ``GGUFQuantizer.postprocess_model`` + +Forward-pass correctness lives behind MPS + kernels-available — the strict +no-fallback policy means we can't validate forward on CI hosts. The structural +tests below run anywhere. +""" + +from __future__ import annotations + +import tempfile +import unittest + +import numpy as np + +from transformers.testing_utils import require_gguf, require_torch +from transformers.utils import is_torch_available + + +if is_torch_available(): + import torch + import torch.nn as nn + + from transformers.integrations.gguf_linear import ( + MODEL_TYPE_TO_GGUF_EXPERTS, + GgufExperts, + GgufLinear, + gguf_linear_supports, + replace_with_gguf_linear, + ) + + +def _random_q4_0_bytes(M: int, K: int, seed: int = 0) -> bytes: + nblocks = M * (K // 32) + rng = np.random.default_rng(seed) + d = rng.uniform(0.01, 0.5, size=nblocks).astype(np.float16) + qs = rng.integers(0, 256, size=(nblocks, 16), dtype=np.uint8) + out = np.empty((nblocks, 18), dtype=np.uint8) + out[:, 0:2] = d.view(np.uint8).reshape(nblocks, 2) + out[:, 2:18] = qs + return out.tobytes() + + +def _random_q4_K_bytes(M: int, K: int, seed: int = 0) -> bytes: + nblocks = M * (K // 256) + rng = np.random.default_rng(seed) + d = rng.uniform(0.01, 0.5, size=nblocks).astype(np.float16) + dmin = rng.uniform(0.0, 0.5, size=nblocks).astype(np.float16) + body = rng.integers(0, 256, size=(nblocks, 140), dtype=np.uint8) + out = np.empty((nblocks, 144), dtype=np.uint8) + out[:, 0:2] = d.view(np.uint8).reshape(nblocks, 2) + out[:, 2:4] = dmin.view(np.uint8).reshape(nblocks, 2) + out[:, 4:144] = body + return out.tobytes() + + +@require_torch +@require_gguf +class GgufLinearConstructionTest(unittest.TestCase): + def test_supports_lookup(self): + import gguf + + for name in ("Q4_0", "Q4_K", "Q5_0", "Q5_1", "Q8_0", "IQ4_NL", "IQ4_XS"): + self.assertTrue(gguf_linear_supports(getattr(gguf.GGMLQuantizationType, name)), name) + self.assertFalse(gguf_linear_supports(gguf.GGMLQuantizationType.F32)) + + def test_weight_buffer_size(self): + # Q4_K: 144 bytes / 256 elems → buffer shape (out_features, bytes_per_row). + layer = GgufLinear(in_features=256, out_features=32, quant_type="Q4_K", bias=False) + self.assertEqual(tuple(layer.weight.shape), (32, (256 // 256) * 144)) + self.assertEqual(layer.weight.dtype, torch.uint8) + + def test_rejects_unsupported_dim(self): + with self.assertRaises(ValueError): + GgufLinear(in_features=130, out_features=32, quant_type="Q4_0") # 130 not multiple of 32 + + +@require_torch +@require_gguf +class GgufLinearStateDictRoundtripTest(unittest.TestCase): + """``state_dict`` → ``load_state_dict`` preserves bytes + extra-state metadata.""" + + def test_q4_0_roundtrip(self): + K, M = 64, 32 + qb = _random_q4_0_bytes(M, K) + layer = GgufLinear(in_features=K, out_features=M, quant_type="Q4_0", bias=False) + # Buffer is `(out_features, bytes_per_row)` 2D uint8. + layer.weight.copy_(torch.frombuffer(bytearray(qb), dtype=torch.uint8).view_as(layer.weight)) + clone = GgufLinear(in_features=K, out_features=M, quant_type="Q4_0", bias=False) + clone.load_state_dict(layer.state_dict()) + self.assertTrue(torch.equal(layer.weight, clone.weight)) + + +@require_torch +@require_gguf +class ReplaceWithGgufLinearTest(unittest.TestCase): + """``replace_with_gguf_linear`` does a uniform meta-time swap — every Linear + becomes a quant-type-agnostic placeholder (kernels bound post-load), except + those in `modules_to_not_convert`.""" + + def test_uniform_swap_with_skip(self): + class M(nn.Module): + def __init__(self): + super().__init__() + self.config = type("C", (), {"model_type": "llama"})() + self.a = nn.Linear(64, 32, bias=False) + self.b = nn.Linear(64, 32, bias=False) + + m = M() + n = replace_with_gguf_linear(m, modules_to_not_convert={"b"}) # swap everything but `.b` + self.assertEqual(n, 1) + self.assertIsInstance(m.a, GgufLinear) + self.assertIsInstance(m.b, nn.Linear) + # Placeholder: quant type unknown at swap, no kernels bound yet, uint8 buffer + # sized from the Q4_K_M average until the real bytes load. + self.assertIsNone(m.a.quant_type) + self.assertIsNone(m.a._mv_op) + self.assertEqual(m.a.weight.dtype, torch.uint8) + + def test_moe_registry_lookup(self): + # Every MoE arch we promise byte-passthrough support for must land on + # the GgufExperts base. gpt_oss is intentionally excluded — see + # `test_gguf_arch_coverage.py::GgufExpertsRegistryTests` for the why. + for arch in ("qwen2_moe", "qwen3_moe", "minimax_m2", "mixtral", "deepseek_v3"): + self.assertIn(arch, MODEL_TYPE_TO_GGUF_EXPERTS, arch) + self.assertIs(MODEL_TYPE_TO_GGUF_EXPERTS[arch], GgufExperts, arch) + + +@require_torch +@require_gguf +class GgufLinearForwardOnMpsTest(unittest.TestCase): + """Forward path only runs on MPS with kernels available — no slow fallback.""" + + def test_forward_errors_off_mps(self): + layer = GgufLinear(in_features=32, out_features=32, quant_type="Q4_0", bias=False) + x = torch.zeros(1, 32, dtype=torch.float32, device="cpu") + with self.assertRaisesRegex(RuntimeError, "GgufLinear runs only on MPS"): + layer(x) + + +@require_torch +@require_gguf +class FromPretrainedDequantPathTest(unittest.TestCase): + """`gguf_file=` + explicit `dtype=` → dequant path; no GgufLinear in the model.""" + + def test_dequant_on_load(self): + # Pull a tiny GGUF off the Hub. The dequant path runs on CPU. + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained( + "unsloth/Llama-3.2-1B-Instruct-GGUF", + gguf_file="Llama-3.2-1B-Instruct-Q4_K_M.gguf", + dtype=torch.bfloat16, + device_map="cpu", + ) + n_gguf = sum(1 for _, mod in model.named_modules() if isinstance(mod, GgufLinear)) + self.assertEqual(n_gguf, 0) + # First projection weight came through dequant. + q = next(p for name, p in model.named_parameters() if name.endswith("q_proj.weight")) + self.assertEqual(q.dtype, torch.bfloat16) + + +@require_torch +@require_gguf +class FromPretrainedOnTheFlyTest(unittest.TestCase): + """`quantization_config=GgufQuantizeConfig(...)` → meta-time swap, GGUFQuantize op packs bytes.""" + + def test_on_the_fly_swap(self): + from transformers import AutoModelForCausalLM, GgufQuantizeConfig + + model = AutoModelForCausalLM.from_pretrained( + "HuggingFaceTB/SmolLM2-135M", + dtype=torch.float32, + quantization_config=GgufQuantizeConfig(quant_type="Q4_0", modules_to_convert=["model.layers.*.mlp.*"]), + ) + gguf_linears = [name for name, mod in model.named_modules() if isinstance(mod, GgufLinear)] + # 30 layers × 3 mlp linears = 90. + self.assertEqual(len(gguf_linears), 90) + # The buffer is uint8 and non-empty after the load. + sample = model.get_submodule(gguf_linears[0]).weight + self.assertEqual(sample.dtype, torch.uint8) + self.assertGreater(sample.numel(), 0) + self.assertTrue(bool((sample != 0).any()), "weight should have been populated by GGUFQuantize") + + +@require_torch +@require_gguf +class GenerationDefaultsTest(unittest.TestCase): + """`postprocess_model` should default `cache_implementation` + `compile_config` + when the GgufLinear swap is in effect.""" + + def test_defaults_applied(self): + if not torch.backends.mps.is_available(): + self.skipTest("GgufLinear swap only triggers on MPS") + from transformers.integrations.gguf_kernels import ensure_metal_kernels + + try: + ensure_metal_kernels() + except Exception: + self.skipTest("GgufLinear swap also requires the `ArthurZ/gguf-kernels` package") + from transformers import AutoModelForCausalLM + from transformers.generation.configuration_utils import CompileConfig + + model = AutoModelForCausalLM.from_pretrained( + "unsloth/Llama-3.2-1B-Instruct-GGUF", + gguf_file="Llama-3.2-1B-Instruct-Q4_K_M.gguf", + ) + self.assertEqual(model.generation_config.cache_implementation, "static") + self.assertIsInstance(model.generation_config.compile_config, CompileConfig) + self.assertEqual(model.generation_config.compile_config.mode, "reduce-overhead") + + +@require_torch +@require_gguf +class SaveRoundtripTest(unittest.TestCase): + """`state_dict` → safetensors round-trip via ``save_pretrained`` / ``from_pretrained``.""" + + def test_uint8_buffers_survive_safetensors(self): + # Quick check that uint8 buffer values survive a CPU round-trip via state_dict. + K, M = 32, 32 + layer = GgufLinear(in_features=K, out_features=M, quant_type="Q4_0", bias=False) + src = torch.frombuffer(bytearray(_random_q4_0_bytes(M, K, seed=11)), dtype=torch.uint8) + layer.weight.copy_(src.view_as(layer.weight)) + with tempfile.TemporaryDirectory(): + sd = layer.state_dict() + clone = GgufLinear(in_features=K, out_features=M, quant_type="Q4_0", bias=False) + clone.load_state_dict(sd) + self.assertTrue(torch.equal(layer.weight, clone.weight)) + + +if __name__ == "__main__": + unittest.main()