Update esmfold model not to use param_buffer_assignment

[jiminha]

What does this PR do?

Transformer 4.43 update includes this PR: huggingface/transformers#31771
This PR introduces utilizing _assign_to_params_buffers as a way to speed up weight loading if the dtypes of models are the same. This PR only checks though the very first key's dtype of the model parameter against state_dict and determine if this feature can be used or not. This particular model "facebook/esmfold_v1" weights/bias of the encoder layers are float16 and rest of them are float32, and the first key happen to be float32, so it determines this model can use this feature.

For gaudi2/gaudi3 - weights/biases are loaded as float16.
For gaudi1 - runtime error seen due to float16 not being supported.

Issues Fixed

Run time error on Gaudi1 due to model weights are initialized with float16.

python run_esmfold.py

Traceback (most recent call last):
  File "/root/optimum-habana-fork/examples/protein-folding/run_esmfold.py", line 86, in <module>
    model = model.to(device)
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 2871, in to
    return super().to(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/habana_frameworks/torch/core/weight_sharing.py", line 179, in wrapped_to
    result = self.original_to(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1176, in to
    return self._apply(convert)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 779, in _apply
    module._apply(fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 779, in _apply
    module._apply(fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 779, in _apply
    module._apply(fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 804, in _apply
    param_applied = fn(param)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1162, in convert
    return t.to(
  File "/usr/local/lib/python3.10/dist-packages/habana_frameworks/torch/core/weight_sharing.py", line 57, in __torch_function__
    return super().__torch_function__(func, types, new_args, kwargs)
RuntimeError: float16/half is not supported on Gaudi.  

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

The code quality check failed, please run make style.

[regisss]

Wow I didn't see that, thanks for the fix!