Repository: huggingface/transformers (Python)
Date: 2024-06-27 11:06:08+00:00
Task type: Bug Fix

Issue:
Can't load models with a gamma or beta parameter

It seems that you cannot create parameters with the string `gamma` or `beta` in any modules you write if you intend to save/load them with the transformers library. There is a small function called `_fix_keys`  implemented in the model loading ([link](https://github.com/huggingface/transformers/blob/0290ec19c901adc0f1230ebdccad11c40af026f5/src/transformers/modeling_utils.py#L3637-L3642)). It renames __all__  instances of  `beta` or `gamma` in any substring of the sate_dict keys to be `bias` and  `weight`.  This means if your modules actually have a parameter with these names, they won't be loaded when using a pretrained model.

As far as I can tell, it's completely undocumented that people shouldn't create any parameters with the string `gamma` or `beta` in them. 

Here is a minimal reproducible example:

```
import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig

class Model(PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.gamma = nn.Parameter(torch.zeros(4))

    def forward(self):
        return self.gamma.sum()


if __name__ == '__main__':
    config = PretrainedConfig()

    # 1) First run this
    #model = Model(config)
    #print(model())

    #model.save_pretrained('test_out')

    # 2) Then try this
    model = Model.from_pretrained('test_out', config=config)
    print(model())
```

When you run this code, you get the following error:
```
Some weights of Model were not initialized from the model checkpoint at test_out and are newly initialized: ['gamma']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```




Issue discussion:
Yes that's correct, it's a bug I pointed out in my [video series](https://www.youtube.com/watch?v=wElORCdXHTU&t=1s&ab_channel=NielsRogge) on contributing to Transformers.

This is due to these lines: https://github.com/huggingface/transformers/blob/0290ec19c901adc0f1230ebdccad11c40af026f5/src/transformers/modeling_utils.py#L579-L582.

I assume they are there for backwards compatibility reasons. If we would know which models require this exception, we could fix this.
I assumed the same, but it's a pretty annoying bug to have to find on your own. Would it be worth adding a warning to the init method of the `PreTrainedModel ` class to let users know if their parameters have the string "gamma" or "beta" in them and encourage them to change it? At least while this block of code still exists in the codebase.

It's further complicated by the fact that `accelerate` uses torch.load_state, which doesn't do this renaming. So there is an incompatibility between two highly coupled libraries.
Hi @malik-ali, thanks for raising this issue! Indeed, this isn't a desired behaviour. 

> If we would know which models require this exception, we could fix this.

I think this would be very hard to do. There are many saved checkpoints both on and off the hub, as well as all sorts of custom models which might rely on this behaviour.

> Would it be worth adding a warning to the init method of the PreTrainedModel  class to let users know if their parameters have the string "gamma" or "beta" in them and encourage them to change it? At least while this block of code still exists in the codebase.

Yes, I think a warning for a few cycle releases is the best way to go. I would put this in the `_load_state_dict_into_model` function and trigger if "gamma" or "beta" are in the key. 

It won't be possible to tell if the parameter is from an "old" state or a new model, but we can warn that the renaming is happening, that the behaviour will be removed in a future release and they should update the weights in their state dict to use "weight" or "bias" to be loaded properly. 

@malik-ali Would you like to open a PR to add this? This way you get the github contribution for your suggested solution
@amyeroberts I'd be happy to! Just one question: if we add this to the `_load_state_dict_into_model`, is it correct that users would only see this warning when loading their pretrained model? 

I ask because I ran into this issue after training a model for several days and later loading it. It would have been nice to see the warning before doing all the training, so that I could rename the parameters on the spot. Do you think a warning like that would be feasible?

(My fix was to manually rename the keys of the saved state_dict and then rename the parameters in my model)
Good point! In this case, we'll need to add a warning in two places to make sure we catch both new model creations and old state dicts being loaded in. 
+1 Find this problem today...
@amyeroberts I might not have a chance to push a fix for this for at least a few weeks so please feel free to make any changes as you (or anyone) wishes!
@malik-ali OK - thanks for letting us know. I've added a 'Good difficult Issue' to flag for anyone in the community that might want to tackle this in the meantime

Solution patch:
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -104,6 +104,8 @@
 
 XLA_USE_BF16 = os.environ.get("XLA_USE_BF16", "0").upper()
 XLA_DOWNCAST_BF16 = os.environ.get("XLA_DOWNCAST_BF16", "0").upper()
+PARAM_RENAME_WARNING = "A parameter name that contains `{}` will be renamed internally to `{}`. Please use a different name to suppress this warning."
+
 
 if is_accelerate_available():
     from accelerate import dispatch_model, infer_auto_device_map, init_empty_weights
@@ -662,8 +664,10 @@ def _load_state_dict_into_model(model_to_load, state_dict, start_prefix):
     for key in state_dict.keys():
         new_key = None
         if "gamma" in key:
+            logger.warning(PARAM_RENAME_WARNING.format("gamma", "weight"))
             new_key = key.replace("gamma", "weight")
         if "beta" in key:
+            logger.warning(PARAM_RENAME_WARNING.format("beta", "bias"))
             new_key = key.replace("beta", "bias")
         if new_key:
             old_keys.append(key)
@@ -807,8 +811,10 @@ def _load_state_dict_into_meta_model(
     for key in state_dict.keys():
         new_key = None
         if "gamma" in key:
+            logger.warning(PARAM_RENAME_WARNING.format("gamma", "weight"))
             new_key = key.replace("gamma", "weight")
         if "beta" in key:
+            logger.warning(PARAM_RENAME_WARNING.format("beta", "bias"))
             new_key = key.replace("beta", "bias")
         if new_key:
             old_keys.append(key)


Test patch:
diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py
--- a/tests/utils/test_modeling_utils.py
+++ b/tests/utils/test_modeling_utils.py
@@ -1499,6 +1499,57 @@ def test_model_from_pretrained_from_mlx(self):
             outputs_from_saved = new_model(input_ids)
             self.assertTrue(torch.allclose(outputs_from_saved["logits"], outputs["logits"]))
 
+    def test_warning_for_beta_gamma_parameters(self):
+        class TestModelGamma(PreTrainedModel):
+            def __init__(self, config):
+                super().__init__(config)
+                self.gamma_param = nn.Parameter(torch.ones(10))
+                self.post_init()
+
+            def forward(self):
+                return self.gamma_param.sum()
+
+        logger = logging.get_logger("transformers.modeling_utils")
+        config = PretrainedConfig()
+        warning_msg_gamma = "A parameter name that contains `gamma` will be renamed internally"
+        model = TestModelGamma(config)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir)
+            with LoggingLevel(logging.WARNING):
+                with CaptureLogger(logger) as cl1:
+                    _, loading_info = TestModelGamma.from_pretrained(tmp_dir, config=config, output_loading_info=True)
+
+        missing_keys = loading_info["missing_keys"]
+        unexpected_keys = loading_info["unexpected_keys"]
+        self.assertIn(warning_msg_gamma, cl1.out)
+        self.assertIn("gamma_param", missing_keys)
+        self.assertIn("weight_param", unexpected_keys)
+
+        class TestModelBeta(PreTrainedModel):
+            def __init__(self, config):
+                super().__init__(config)
+                self.beta_param = nn.Parameter(torch.ones(10))
+                self.post_init()
+
+            def forward(self):
+                return self.beta_param.sum()
+
+        warning_msg_beta = "A parameter name that contains `beta` will be renamed internally"
+        model = TestModelBeta(config)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model.save_pretrained(tmp_dir)
+            with LoggingLevel(logging.WARNING):
+                with CaptureLogger(logger) as cl2:
+                    _, loading_info = TestModelBeta.from_pretrained(tmp_dir, config=config, output_loading_info=True)
+
+        missing_keys = loading_info["missing_keys"]
+        unexpected_keys = loading_info["unexpected_keys"]
+        self.assertIn(warning_msg_beta, cl2.out)
+        self.assertIn("beta_param", missing_keys)
+        self.assertIn("bias_param", unexpected_keys)
+
 
 @slow
 @require_torch