ESM models quantization broken in current BNB

[ddofer]

System Info

WSL2 on Windows 11. Latest HF, pytorch, bnb, PEFT, anaconda as of 8.2.2024.

Reproduction

nf4_config = BitsAndBytesConfig( load_in_4bit=True, )
model = AutoModelForSequenceClassification.from_pretrained("facebook/esm2_t12_35M_UR50D",    problem_type="single_label_classification",   quantization_config=nf4_config )

Results in an error: KeyError: 'inv_freq'

Full error message trace:

File ~/anaconda3/envs/hf/lib/python3.10/site-packages/transformers/modeling_utils.py:3926, in PreTrainedModel._load_pretrained_model(cls, model, state_dict, loaded_keys, resolved_archive_file, pretrained_model_name_or_path, ignore_mismatched_sizes, sharded_metadata, _fast_init, low_cpu_mem_usage, device_map, offload_folder, offload_state_dict, dtype, hf_quantizer, keep_in_fp32_modules)
   3924                     hf_quantizer.create_quantized_param(model, param, key, "cpu", state_dict)
   3925     else:
-> 3926         new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
   3927             model_to_load,
   3928             state_dict,
   3929             loaded_keys,
   3930             start_prefix,
   3931             expected_keys,
   3932             device_map=device_map,
   3933             offload_folder=offload_folder,
   3934             offload_index=offload_index,
   3935             state_dict_folder=state_dict_folder,
   3936             state_dict_index=state_dict_index,
   3937             dtype=dtype,
   3938             hf_quantizer=hf_quantizer,
   3939             is_safetensors=is_safetensors,
   3940             keep_in_fp32_modules=keep_in_fp32_modules,
   3941             unexpected_keys=unexpected_keys,
   3942         )
   3943         error_msgs += new_error_msgs
   3944 else:

File ~/anaconda3/envs/hf/lib/python3.10/site-packages/transformers/modeling_utils.py:802, in _load_state_dict_into_meta_model(model, state_dict, loaded_state_dict_keys, start_prefix, expected_keys, device_map, offload_folder, offload_index, state_dict_folder, state_dict_index, dtype, hf_quantizer, is_safetensors, keep_in_fp32_modules, unexpected_keys)
    797 elif param_device == "cpu" and state_dict_index is not None:
    798     state_dict_index = offload_weight(param, param_name, model, state_dict_folder, state_dict_index)
    799 elif (
    800     hf_quantizer is None
    801     or (not hf_quantizer.requires_parameters_quantization)
--> 802     or (not hf_quantizer.check_quantized_param(model, param, param_name, state_dict))
    803 ):
    804     # For backward compatibility with older versions of `accelerate` and for non-quantized params
    805     set_module_tensor_to_device(model, param_name, param_device, **set_module_kwargs)
    806 else:

File ~/anaconda3/envs/hf/lib/python3.10/site-packages/transformers/quantizers/quantizer_bnb_4bit.py:124, in Bnb4BitHfQuantizer.check_quantized_param(self, model, param_value, param_name, state_dict)
    121 import bitsandbytes as bnb
    123 module, tensor_name = get_module_from_name(model, param_name)
--> 124 if isinstance(module._parameters[tensor_name], bnb.nn.Params4bit):
    125     # Add here check for loaded components' dtypes once serialization is implemented
    126     return True
    127 elif isinstance(module, bnb.nn.Linear4bit) and tensor_name == "bias":
    128     # bias could be loaded by regular set_module_tensor_to_device() from accelerate,
    129     # but it would wrongly use uninitialized weight there.

KeyError: 'inv_freq'

Expected behavior

Succesful import of the model and quantization. Previously this had worked, something broke in the latest bnb or peft . There are issues in general with the ESM models, although they are officially meant to be supported.

[matthewdouglas]

This is fixed in transformers. See: huggingface/transformers#29329