Repository: huggingface/transformers
Pull request: Clean bnb integration using weight converter

# What does this PR do?

This PR cleans bnb integration, we don't need to deal with missing keys.

Changed files:
filename: src/transformers/core_model_loading.py
patch: @@ -674,9 +674,6 @@ def convert_and_load_state_dict_in_model(
 
     renamings = [entry for entry in weight_mapping if isinstance(entry, WeightRenaming)]
     converters = [entry for entry in weight_mapping if isinstance(entry, WeightConverter)]
-    if hf_quantizer:
-        # We will add the quantizer's deserialization WeightConverter here.
-        pass
 
     param_name_to_load: dict[str, Union[WeightRenaming | WeightConverter]] = {}
 
@@ -724,11 +721,19 @@ def convert_and_load_state_dict_in_model(
                 source_pattern = renamed_key
 
             # 5. Handle dtype casting
-            if hf_quantizer is not None and hf_quantizer.param_needs_quantization(model, renamed_key):
+            if (
+                hf_quantizer
+                and not hf_quantizer.pre_quantized
+                and hf_quantizer.param_needs_quantization(model, renamed_key)
+            ):
                 mapping.quantization_operation = hf_quantizer.get_quantize_ops()
 
             _dtype = dtype
-            if dtype_plan != {} and dtype_policy_alt.search(renamed_key):
+            if hf_quantizer and hf_quantizer.pre_quantized and original_key != renamed_key:
+                # if the key was renamed as it is not available in the state dict otherwise, it means that we are deserializing it,
+                # so we need to make sure to load the tensor with the same dtype from the checkpoint
+                _dtype = None
+            elif dtype_plan != {} and dtype_policy_alt.search(renamed_key):
                 matched_dtype_pattern = dtype_policy_alt.search(renamed_key)
                 if matched_dtype_pattern is not None:
                     _dtype = dtype_plan[matched_dtype_pattern.group()]

filename: src/transformers/integrations/bitsandbytes.py
patch: @@ -1,5 +1,4 @@
 import inspect
-from collections import defaultdict
 from inspect import signature
 
 from ..core_model_loading import ConversionOps
@@ -37,54 +36,64 @@ def __init__(self, hf_quantizer):
     def convert(
         self,
         input_dict: dict[str, list[torch.Tensor]],
+        full_layer_name: str | None = None,
         model: torch.nn.Module | None = None,
-        missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
         """
         we need to store some parameters to create the quantized weight. For example, bnb requires 6 values that are stored in the checkpoint to recover the quantized weight. So we store them in a dict that it stored in hf_quantizer for now as we can't save it in the op since we create an op per tensor.
         """
-        target_key, value = tuple(input_dict.items())[0]
+        value = list(input_dict.values())[0]
         value = value[0] if isinstance(value, list) else value
 
-        full_name = target_key
         # update param name to get the weights instead of the quantized stats
-        target_key = self.hf_quantizer.get_param_name(target_key)
-        module, _ = get_module_from_name(model, target_key)
-
-        if not self.hf_quantizer.pre_quantized:
-            # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
-            # Since weights are saved in the correct "orientation", we skip transposing when loading.
-            if issubclass(module.source_cls, Conv1D):
-                value = value.T
-            old_value = model.get_parameter_or_buffer(target_key)
-            new_value = bnb.nn.Params4bit(value, requires_grad=False, **old_value.__dict__).to(value.device)
-            # remove missing keys that were create when initializing Params4bit
-            for key in new_value.quant_state.as_dict(packed=True).keys():
-                missing_keys.discard(f"{full_name}.{key}")
-            module._is_hf_initialized = True
-            return {target_key: new_value}
-        else:
-            module_name = target_key.rsplit(".", 1)[0]
-            # Save the states for later quantization when they are all gathered
-            if not hasattr(self.hf_quantizer, "param_quant_stats"):
-                self.hf_quantizer.param_quant_stats = defaultdict(dict)
-            self.hf_quantizer.param_quant_stats[module_name].update({full_name: value})
-            missing_keys.discard(full_name)
-            # We are ready for quantization in this case (note, the +1 is for the weight itself)
-            if len(self.hf_quantizer.param_quant_stats[module_name]) == len(self.hf_quantizer.bnb_keys) + 1:
-                weight = self.hf_quantizer.param_quant_stats[module_name].pop(f"{module_name}.weight")
-                new_value = bnb.nn.Params4bit.from_prequantized(
-                    data=weight,
-                    quantized_stats=self.hf_quantizer.param_quant_stats[module_name],
-                    requires_grad=False,
-                    device=value.device,
-                    module=module,
-                )
-                module._is_hf_initialized = True
-                del self.hf_quantizer.param_quant_stats[module_name]
-                return {target_key: new_value}
-            return {}
+        module, _ = get_module_from_name(model, full_layer_name)
+
+        # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
+        # Since weights are saved in the correct "orientation", we skip transposing when loading.
+        if issubclass(module.source_cls, Conv1D):
+            value = value.T
+
+        old_value = model.get_parameter_or_buffer(full_layer_name)
+        new_value = bnb.nn.Params4bit(value, requires_grad=False, **old_value.__dict__).to(value.device)
+        module._is_hf_initialized = True
+        return {full_layer_name: new_value}
+
+
+class Bnb4bitDeserialize(ConversionOps):
+    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]:
+        """
+        Deserialization of bnb keys. We need 6 keys to recreate the quantized weights
+        """
+        if len(input_dict) == 1:
+            # special case when we only fetched the weight
+            # since we collected keys, we need to return it like that
+            return {full_layer_name: input_dict["weight"]}
+
+        for key, value in input_dict.items():
+            if isinstance(value, list):
+                input_dict[key] = value[0]
+
+        weight = input_dict.pop("weight")
+        module, _ = get_module_from_name(model, full_layer_name)
+        new_value = bnb.nn.Params4bit.from_prequantized(
+            data=weight,
+            quantized_stats=input_dict,
+            requires_grad=False,
+            device=weight.device,
+            module=module,
+        )
+        module._is_hf_initialized = True
+        return {full_layer_name: new_value}
 
 
 class Bnb8bitQuantize(ConversionOps):
@@ -95,48 +104,56 @@ def convert(
         self,
         input_dict: dict[str, list[torch.Tensor]],
         model: torch.nn.Module | None = None,
-        missing_keys=None,
+        full_layer_name: str | None = None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
-        target_key, value = tuple(input_dict.items())[0]
+        value = list(input_dict.values())[0]
         value = value[0] if isinstance(value, list) else value
 
-        module, tensor_name = get_module_from_name(model, target_key)
-        module_name = target_key.rsplit(".", 1)[0]
-
-        if not self.hf_quantizer.pre_quantized:
-            # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
-            # Since weights are saved in the correct "orientation", we skip transposing when loading.
-            if issubclass(module.source_cls, Conv1D):
-                value = value.T
-            value_device = value.device
-            kwargs = getattr(module, tensor_name).__dict__
-            kwargs.pop("SCB", None)
-            new_value = bnb.nn.Int8Params(value.to("cpu"), requires_grad=False, **kwargs).to(value_device)
-            missing_keys.discard(f"{module_name}.weight_format")
-            missing_keys.discard(f"{module_name}.SCB")
-            return {target_key: new_value}
-        else:
-            missing_keys.discard(target_key)
-            # useless key that gets saved for no reason
-            if tensor_name.endswith("weight_format"):
-                return {}
-            # Save the states for later quantization when they are all gathered
-            if not hasattr(self.hf_quantizer, "param_quant_stats"):
-                self.hf_quantizer.param_quant_stats = defaultdict(dict)
-            self.hf_quantizer.param_quant_stats[module_name].update({target_key: value})
-            # We are ready for quantization in this case (SCB and the weight)
-            if len(self.hf_quantizer.param_quant_stats[module_name]) == 2:
-                weight = self.hf_quantizer.param_quant_stats[module_name].pop(f"{module_name}.weight")
-                kwargs = getattr(module, "weight").__dict__
-                weight_device = weight.device
-                new_value = bnb.nn.Int8Params(weight.to("cpu"), requires_grad=False, **kwargs).to(weight_device)
-                setattr(new_value, "SCB", self.hf_quantizer.param_quant_stats[module_name][f"{module_name}.SCB"])
-                del self.hf_quantizer.param_quant_stats[module_name]
-                # sometimes, weight_format is not saved so we need to remove it from missing keys ...
-                missing_keys.discard(f"{module_name}.weight_format")
-                return {f"{module_name}.weight": new_value}
-            return {}
+        module, _ = get_module_from_name(model, full_layer_name)
+
+        # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
+        # Since weights are saved in the correct "orientation", we skip transposing when loading.
+        if issubclass(module.source_cls, Conv1D):
+            value = value.T
+        value_device = value.device
+        kwargs = model.get_parameter_or_buffer(full_layer_name).__dict__
+        kwargs.pop("SCB", None)
+        new_value = bnb.nn.Int8Params(value.to("cpu"), requires_grad=False, **kwargs).to(value_device)
+        return {full_layer_name: new_value}
+
+
+class Bnb8bitDeserialize(ConversionOps):
+    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]:
+        """
+        Deserialization of bnb keys.
+        """
+        if len(input_dict) == 1:
+            # special case when we only fetched the weight
+            # since we collected keys, we need to return it like that
+            return {full_layer_name: input_dict["weight"]}
+
+        for key, value in input_dict.items():
+            if isinstance(value, list):
+                input_dict[key] = value[0]
+
+        module, _ = get_module_from_name(model, full_layer_name)
+
+        weight = input_dict["weight"]
+        kwargs = model.get_parameter_or_buffer(full_layer_name).__dict__
+        kwargs["SCB"] = input_dict["SCB"]
+        new_value = bnb.nn.Int8Params(weight, requires_grad=False, **kwargs).to(weight.device)
+        module._is_hf_initialized = True
+        return {full_layer_name: new_value}
 
 
 def _replace_with_bnb_linear(
@@ -178,8 +195,6 @@ def _replace_with_bnb_linear(
                             has_fp16_weights=quantization_config.llm_int8_has_fp16_weight,
                             threshold=quantization_config.llm_int8_threshold,
                         )
-                        # hack to create the correct keys in the state dict with the right dtype
-                        new_module.weight.SCB = torch.empty(1, dtype=torch.float32)
                         if pre_quantized:
                             new_module.weight.data = new_module.weight.data.to(dtype=torch.int8)
                         model._modules[name] = new_module
@@ -205,25 +220,6 @@ def _replace_with_bnb_linear(
                                 quant_type=quantization_config.bnb_4bit_quant_type,
                                 **extra_kwargs,
                             )
-                            from bitsandbytes.functional import QuantState
-
-                            # hack to create the correct keys in the state dict with the right dtype
-                            absmax_dtype = (
-                                torch.uint8 if quantization_config.bnb_4bit_use_double_quant else torch.float32
-                            )
-                            new_module.weight.quant_state = QuantState(
-                                absmax=torch.empty(1, dtype=absmax_dtype),
-                                code=torch.empty(1, dtype=torch.float32),
-                                shape=(1,),
-                                offset=torch.empty(1),
-                                quant_type=quantization_config.bnb_4bit_quant_type,
-                                state2=QuantState(
-                                    absmax=torch.empty(1, dtype=torch.float32),
-                                    code=torch.empty(1, dtype=torch.float32),
-                                )
-                                if quantization_config.bnb_4bit_use_double_quant
-                                else None,
-                            )
                             if pre_quantized:
                                 # this is kind of an edge case when supporting both loading and quantization ...
                                 # we need to set the right dtype as we cast the checkpoint with the dtype of the meta model

filename: src/transformers/quantizers/quantizer_bnb_4bit.py
patch: @@ -37,6 +37,7 @@
 if is_torch_available():
     import torch
 
+    from ..core_model_loading import WeightConverter
     from ..pytorch_utils import Conv1D
 
 logger = logging.get_logger(__name__)
@@ -298,3 +299,24 @@ def get_quantize_ops(self):
         from ..integrations.bitsandbytes import Bnb4bitQuantize
 
         return Bnb4bitQuantize(self)
+
+    def get_weight_conversions(self):
+        from ..integrations.bitsandbytes import Bnb4bitDeserialize
+
+        if self.pre_quantized:
+            return [
+                WeightConverter(
+                    source_keys=[
+                        "weight.nested_absmax",
+                        "weight.nested_quant_map",
+                        "weight.quant_map",
+                        "weight.absmax",
+                        "weight.quant_state.bitsandbytes__nf4",
+                        "weight.quant_state.bitsandbytes__fp4",
+                        "weight",
+                    ],
+                    target_keys="weight",
+                    operations=[Bnb4bitDeserialize(self)],
+                )
+            ]
+        return []

filename: src/transformers/quantizers/quantizer_bnb_8bit.py
patch: @@ -34,6 +34,7 @@
 if is_torch_available():
     import torch
 
+    from ..core_model_loading import WeightConverter
     from ..pytorch_utils import Conv1D
 
 logger = logging.get_logger(__name__)
@@ -246,3 +247,16 @@ def get_quantize_ops(self):
         from ..integrations.bitsandbytes import Bnb8bitQuantize
 
         return Bnb8bitQuantize(self)
+
+    def get_weight_conversions(self):
+        from ..integrations.bitsandbytes import Bnb8bitDeserialize
+
+        if self.pre_quantized:
+            return [
+                WeightConverter(
+                    source_keys=["SCB", "weight_format", "weight"],
+                    target_keys="weight",
+                    operations=[Bnb8bitDeserialize(self)],
+                )
+            ]
+        return []