"docs/source/en/model_doc/bart.mdx" did not exist on "b256f3518d470ba53be519992c3b9d97d174db48"
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
):
set_module_tensor_to_device(model, key, "cpu", value)
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.
if _fast_init:
......
......@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
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
......@@ -143,7 +143,7 @@ class Bnb4BitHfQuantizer(HfQuantizer):
param_name: str,
target_device: "torch.device",
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()
......@@ -198,6 +198,7 @@ class Bnb4BitHfQuantizer(HfQuantizer):
for k, v in state_dict.items():
if param_name + "." in k:
quantized_stats[k] = v
if unexpected_keys is not None and k in unexpected_keys:
unexpected_keys.remove(k)
new_value = bnb.nn.Params4bit.from_prequantized(
......
......@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
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
......@@ -162,7 +162,7 @@ class Bnb8BitHfQuantizer(HfQuantizer):
param_name: str,
target_device: "torch.device",
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()
......@@ -207,6 +207,7 @@ class Bnb8BitHfQuantizer(HfQuantizer):
module._parameters[tensor_name] = new_value
if fp16_statistics is not None:
setattr(module.weight, "SCB", fp16_statistics.to(target_device))
if unexpected_keys is not None:
unexpected_keys.remove(fp16_statistics_key)
def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
......
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