Repository: huggingface/transformers
Pull request: fix(models): Unpack BitNet packed weights to fix CI failure

### What does this PR do?

→ Fixes `tests/models/bitnet/test_modeling_bitnet.py::BitNetIntegrationTest::test_model_generation && tests/models/bitnet/test_modeling_bitnet.py::BitNetIntegrationTest::test_model_logits`.

→ The reasoning; picked a couple of canonical repo patterns and applied them to BitNet to fix the issue:
1. [Bnb4bitDeserialize](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/bitsandbytes.py#L61-L93) for `bitsandbytes.py::BitNetDeserialize`, with a small caveat; `bitsandbytes.py` checks `if len(input_dict) == 1: return input_dict` because it has multiple state dict keys (weight, absmax, quant_map, etc.); but `BitNet` only has weight, so this check would cause an early return, hence we omit it.
2. [get_weight_conversions()](https://github.com/huggingface/transformers/blob/main/src/transformers/quantizers/quantizer_bnb_4bit.py#L166-L185) in `quantizer_bnb_4bit.py`, which integrates into accelerate’s loading pipeline before the weights reach the model.
3. Packing implementation: [pack_weights(quantized_weights: torch.Tensor)](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/bitnet.py#L17-L52), [unpack_weights(packed: torch.Tensor, dtype: torch.dtype)](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/bitnet.py#L55-L121)

For more details on reproducing the bug, please visit the linked issue!

Fixes #43720.

### Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you fix any necessary existing tests?

Changed files:
filename: src/transformers/integrations/bitnet.py
patch: @@ -368,3 +368,38 @@ def replace_with_bitnet_linear(model, modules_to_not_convert: list[str] | None =
         )
 
     return model
+
+
+class BitNetDeserialize:
+    def __init__(self, hf_quantizer):
+        self.hf_quantizer = hf_quantizer
+
+    def convert(
+        self,
+        input_dict: dict[str, list[torch.Tensor]],
+        model: torch.nn.Module | None = None,
+        full_layer_name: str | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        for key, value in input_dict.items():
+            if isinstance(value, list):
+                input_dict[key] = value[0]
+        key_weight = "weight"
+        weight = input_dict.pop(key_weight)
+        from ..quantizers.quantizers_utils import get_module_from_name
+
+        needs_unpacking = False
+        target_dtype = weight.dtype
+        if model is not None and full_layer_name is not None:
+            module, _ = get_module_from_name(model, full_layer_name)
+            if hasattr(module, "out_features") and hasattr(module, "in_features"):
+                # Packed: shape[0] * VALUES_PER_ITEM == out_features
+                # Unpacked: shape[0] == out_features
+                expected_out = module.out_features
+                actual_out = weight.shape[0]
+                if actual_out * VALUES_PER_ITEM == expected_out:
+                    needs_unpacking = True
+        if needs_unpacking:
+            weight_uint8 = weight.to(torch.uint8)
+            weight = unpack_weights(weight_uint8, dtype=target_dtype)
+        return {key_weight: weight}

filename: src/transformers/quantizers/quantizer_bitnet.py
patch: @@ -103,3 +103,20 @@ def is_qat_trainable(self) -> bool:
             self.quantization_config.linear_class == "autobitlinear"
             and self.quantization_config.quantization_mode == "online"
         )
+
+    def get_weight_conversions(self):
+        from ..core_model_loading import WeightConverter
+        from ..integrations.bitnet import BitNetDeserialize
+
+        if (
+            self.quantization_config.linear_class == "autobitlinear"
+            and self.quantization_config.quantization_mode == "offline"
+        ):
+            return [
+                WeightConverter(
+                    source_patterns=["weight"],
+                    target_patterns=["weight"],
+                    operations=[BitNetDeserialize(self)],
+                )
+            ]
+        return []

filename: tests/quantization/bitnet_integration/test_bitnet.py
patch: @@ -36,6 +36,42 @@
     import torch
 
 
+class BitNetPackedWeightsTest(unittest.TestCase):
+    def test_offline_autobitlinear_weight_conversion(self):
+        """get_weight_conversions() must return a WeightConverter for autobitlinear+offline"""
+        from transformers.quantizers.quantizer_bitnet import BitNetHfQuantizer
+
+        config = BitNetQuantConfig(linear_class="autobitlinear", quantization_mode="offline")
+        quantizer = BitNetHfQuantizer(config)
+        conversions = quantizer.get_weight_conversions()
+        self.assertEqual(len(conversions), 1)
+        self.assertEqual(conversions[0].source_patterns, ["weight"])
+        self.assertEqual(conversions[0].target_patterns, ["weight"])
+
+    def test_unpack_packed_weights(self):
+        """BitNetDeserialize.convert() must unpack packed weights to the original ternary values"""
+        from transformers.integrations.bitnet import AutoBitLinear, BitNetDeserialize, pack_weights
+
+        out_features = 128
+        in_features = 64
+
+        class SimpleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = AutoBitLinear(in_features=in_features, out_features=out_features, bias=False)
+
+        model = SimpleModel()
+        # same as the ckpt loading with safetensors: ternary weights {-1, 0, 1} packed into uint8, then cast to bfloat16
+        original = torch.randint(-1, 2, (out_features, in_features)).to(torch.bfloat16)
+        packed = pack_weights(original.clone().float()).to(torch.bfloat16)
+        # packed shape is [out_features // 4, in_features]
+        self.assertEqual(packed.shape[0], out_features // 4)
+        deserializer = BitNetDeserialize(hf_quantizer=None)
+        result = deserializer.convert({"weight": packed}, model=model, full_layer_name="linear.weight")
+        self.assertEqual(result["weight"].shape, (out_features, in_features))
+        self.assertTrue(torch.equal(result["weight"], original))
+
+
 @require_torch_accelerator
 class BitNetQuantConfigTest(unittest.TestCase):
     def test_to_dict(self):