Unverified Commit c852d4fb authored by Younes Belkada's avatar Younes Belkada Committed by GitHub
Browse files

FIX [`bnb`] Make `unexpected_keys` optional (#29420)



* make `unexpected_keys` optional

* push

* Apply suggestions from code review
Co-authored-by: default avataramyeroberts <22614925+amyeroberts@users.noreply.github.com>

---------
Co-authored-by: default avataramyeroberts <22614925+amyeroberts@users.noreply.github.com>
parent 87e2ea33
...@@ -3769,7 +3769,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix ...@@ -3769,7 +3769,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
): ):
set_module_tensor_to_device(model, key, "cpu", value) set_module_tensor_to_device(model, key, "cpu", value)
else: else:
hf_quantizer.create_quantized_param(model, value, key, "cpu", state_dict) hf_quantizer.create_quantized_param(model, value, key, "cpu", state_dict, unexpected_keys)
# retrieve uninitialized modules and initialize before maybe overriding that with the pretrained weights. # retrieve uninitialized modules and initialize before maybe overriding that with the pretrained weights.
if _fast_init: if _fast_init:
......
...@@ -12,7 +12,7 @@ ...@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import importlib import importlib
from typing import TYPE_CHECKING, Any, Dict, List, Union from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from packaging import version from packaging import version
...@@ -143,7 +143,7 @@ class Bnb4BitHfQuantizer(HfQuantizer): ...@@ -143,7 +143,7 @@ class Bnb4BitHfQuantizer(HfQuantizer):
param_name: str, param_name: str,
target_device: "torch.device", target_device: "torch.device",
state_dict: Dict[str, Any], state_dict: Dict[str, Any],
unexpected_keys: List[str], unexpected_keys: Optional[List[str]] = None,
): ):
""" """
combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device() combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device()
...@@ -198,7 +198,8 @@ class Bnb4BitHfQuantizer(HfQuantizer): ...@@ -198,7 +198,8 @@ class Bnb4BitHfQuantizer(HfQuantizer):
for k, v in state_dict.items(): for k, v in state_dict.items():
if param_name + "." in k: if param_name + "." in k:
quantized_stats[k] = v quantized_stats[k] = v
unexpected_keys.remove(k) if unexpected_keys is not None and k in unexpected_keys:
unexpected_keys.remove(k)
new_value = bnb.nn.Params4bit.from_prequantized( new_value = bnb.nn.Params4bit.from_prequantized(
data=param_value, data=param_value,
......
...@@ -12,7 +12,7 @@ ...@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import importlib import importlib
from typing import TYPE_CHECKING, Any, Dict, List, Union from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from packaging import version from packaging import version
...@@ -162,7 +162,7 @@ class Bnb8BitHfQuantizer(HfQuantizer): ...@@ -162,7 +162,7 @@ class Bnb8BitHfQuantizer(HfQuantizer):
param_name: str, param_name: str,
target_device: "torch.device", target_device: "torch.device",
state_dict: Dict[str, Any], state_dict: Dict[str, Any],
unexpected_keys: List[str], unexpected_keys: Optional[List[str]] = None,
): ):
""" """
combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device() combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device()
...@@ -207,7 +207,8 @@ class Bnb8BitHfQuantizer(HfQuantizer): ...@@ -207,7 +207,8 @@ class Bnb8BitHfQuantizer(HfQuantizer):
module._parameters[tensor_name] = new_value module._parameters[tensor_name] = new_value
if fp16_statistics is not None: if fp16_statistics is not None:
setattr(module.weight, "SCB", fp16_statistics.to(target_device)) setattr(module.weight, "SCB", fp16_statistics.to(target_device))
unexpected_keys.remove(fp16_statistics_key) if unexpected_keys is not None:
unexpected_keys.remove(fp16_statistics_key)
def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
model.is_loaded_in_8bit = True model.is_loaded_in_8bit = True
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment