Refacto GGUF weight conversion

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 authoritative list is the `_GGUF_ARCH_CONVERTERS` mapping 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,80 @@ 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. **Pick the HF `model_type`.** The loader looks up the rename table by the HF `model_type` (e.g. `"llama"`, `"qwen2_moe"`, `"gemma3_text"`). This is the same string `config.model_type` produces.
+2. **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`.
+3. **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
+
+The conventions are the same as in `convert_mistral4_weight_to_hf.py` and `Fp8Quantizer`:
+
+```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 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` — same pattern as `Fp8Dequantize` for FP8. You don't need to add it manually.
+
+### Existing transform ops
+
+These live in [`src/transformers/gguf_conversion_ops.py`](https://github.com/huggingface/transformers/blob/main/src/transformers/gguf_conversion_ops.py) and are reusable:
+
+| 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.
+
+### Worked 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 unchanged — the loader moves the raw uint8 bytes to the target device first, then the kernel produces the float tensor on-device.
+
+If you add a new quant type, add the corresponding `_dq_<TYPE>` kernel to `gguf_dequant.py` and register it in `_build_dispatch()`.

setup.py

@@ -112,6 +112,7 @@
     "pydantic>=2",
     "pytest>=7.2.0,<9.0.0",
     "pytest-asyncio>=1.2.0",
+    "pytest-benchmark",
     "pytest-random-order",
     "pytest-rerunfailures<16.0",
     "pytest-timeout",
@@ -218,6 +219,7 @@ def deps_list(*pkgs):
     deps_list(
         "pytest",
         "pytest-asyncio",
+        "pytest-benchmark",
         "pytest-random-order",
         "pytest-rich",
         "pytest-xdist",

src/transformers/core_model_loading.py

@@ -166,7 +166,12 @@ def convert(
                 all_tensors.extend(tensors)
             else:
                 all_tensors.append(tensors)
-        return {target_pattern: torch.cat(all_tensors, dim=self.dim)}
+        result = torch.cat(all_tensors, dim=self.dim)
+        # Release source refs eagerly so accelerator caching allocators (MPS/CUDA) can pool/reclaim
+        # buffers immediately instead of waiting for the caller to drop the input_dict.
+        all_tensors.clear()
+        input_dict.clear()
+        return {target_pattern: result}
 
     def get_target_pattern(self, target_patterns: list[str]) -> str:
         # Here we always return the target pattern
@@ -198,9 +203,22 @@ def convert(
         **kwargs,
     ) -> dict[str, torch.Tensor]:
         merged: dict[str, torch.Tensor] = {}
-        for source_pattern, tensors in input_dict.items():
-            target_pattern = self.get_target_pattern(input_dict, source_pattern, target_patterns)
+        # Snapshot keys so we can pop entries from `input_dict` as we consume them — this releases
+        # per-expert source tensors before the next stack runs, halving peak memory on MoE fuses
+        # (critical on MPS where MTLBuffers stay live until refcounts drop).
+        original_len = len(input_dict)
+        for source_pattern in list(input_dict.keys()):
+            tensors = input_dict.pop(source_pattern)
+            if original_len == 1:
+                if len(target_patterns) == 1:
+                    target_pattern = target_patterns[0]
+                else:
+                    raise ValueError("Undefined Operation encountered!")
+            else:
+                target_pattern = source_pattern
             merged[target_pattern] = torch.stack(tensors, dim=self.dim)
+            tensors.clear()
+            del tensors
         return merged
 
     def get_target_pattern(self, input_dict: dict, source_pattern: str, target_patterns: list[str]) -> str:
@@ -461,6 +479,8 @@ def convert(
 
         for k, v in split_and_fused.items():
             split_and_fused[k] = torch.cat(v, dim=self.concat_dim)
+        # Eager release of per-source tensor lists once the fused outputs are built.
+        input_dict.clear()
 
         return split_and_fused
 
@@ -1365,6 +1385,14 @@ def convert_and_load_state_dict_in_model(
             else:
                 mapping = param_name_to_load.setdefault(renamed_key, WeightRenaming(original_key, renamed_key))
                 source_pattern = original_key
+                # Pre-quantized loaders (e.g. GGUF) need a dequant op on the rename path:
+                # ``WeightRenaming.convert`` returns the raw quantized tensor directly under
+                # the renamed key, and ``set_param_for_module`` skips the shape check when a
+                # quantizer is active. Without an explicit dequant op the layer would land in
+                # the model as raw bytes cast to fp32 with the wrong (byte-) shape.
+                renaming_dequant = getattr(hf_quantizer, "renaming_dequantize_op", None)
+                if renaming_dequant is not None and mapping.quantization_operation is None:
+                    mapping.quantization_operation = renaming_dequant
 
             # 3. Handle dtype casting
             needs_quantization = (
@@ -1435,6 +1463,10 @@ def convert_and_load_state_dict_in_model(
         else:
             loading_info.unexpected_keys.add(renamed_key)
 
+    # When loading onto MPS, periodically drain the MPS allocator pool back to the system so
+    # buffers from completed converter stages don't accumulate and push us into swap. CUDA's pool
+    # is already pressure-aware; MPS's is not.
+    mps_target = torch.backends.mps.is_available() and any(str(d).startswith("mps") for d in device_map.values())
     try:
         for first_param_name, mapping in tqdm(param_name_to_load.items(), desc="Loading weights"):
             try:
@@ -1465,6 +1497,8 @@ def convert_and_load_state_dict_in_model(
 
                 # Cleanup all the tensors that were gathered before next iteration
                 del realized_value
+                if mps_target:
+                    torch.mps.empty_cache()
 
             except SkipParameters:
                 continue

src/transformers/dependency_versions_table.py

@@ -44,6 +44,7 @@
     "pydantic": "pydantic>=2",
     "pytest": "pytest>=7.2.0,<9.0.0",
     "pytest-asyncio": "pytest-asyncio>=1.2.0",
+    "pytest-benchmark": "pytest-benchmark",
     "pytest-random-order": "pytest-random-order",
     "pytest-rerunfailures": "pytest-rerunfailures<16.0",
     "pytest-timeout": "pytest-timeout",