Repository: huggingface/transformers
Pull request: [loading] `tie_weights` is finally symmetric

# What does this PR do?

As per the title. This makes `tie_weights` work fully symmetrically, even in cases where one weight is tied between more than 2 keys.
This will hopefully end the pain of all models checkpoints having the saved weight in the "wrong" key.

This fixes several opened issues: https://github.com/huggingface/transformers/issues/42313 (though we still have an issue with the tokenizer here apparently), https://github.com/huggingface/transformers/issues/42460 and I think others that I lost track of.

Note: for `facebook/bart-large` bart checkpoint, all 3 tied weight keys are explicitly present in the safetensors checkpoint. So according to this PR, they are not tied anymore (with a warning), whereas they would be anyway before. IMO this is fine as if they are in the checkpoints, they were not tied when saved. We could introduce a value check if both source and targets are present and still tie if everything is equal, but I think this is too expensive and too niche to add for now. We may consider doing it if we notice this is more frequent so those `bart` checkpoints

Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -2283,37 +2283,83 @@ def tie_weights(self, missing_keys: Optional[set[str]] = None, recompute_mapping
         If `recompute_mapping=True`, it will re-check all internal submodels and their config to determine the params
         that need to be tied. This is the default when `model.tie_weights()` is called on its own, outside of
         `__init__`, and `from_pretrained`, in case the config values were changed somewhere.
+
+        Note that during `from_pretrained`, tying is *symmetric*: if the mapping says "tie target -> source" but
+        `source` is missing in the checkpoint while `target` exists, we *swap* source and target so we can still
+        tie everything to the parameter that actually exists.
         """
         # In this case, the keys stored in `all_tied_weights_keys` are already correct
         if not recompute_mapping:
             tied_keys = self.all_tied_weights_keys
         else:
             tied_keys = self.get_expanded_tied_weights_keys(all_submodels=True)
 
-        for target_param_name, source_param_name in tied_keys.items():
-            source_param = self.get_parameter_or_buffer(source_param_name)
-            if "." in target_param_name:
-                parent_name, name = target_param_name.rsplit(".", 1)
-                parent = self.get_submodule(parent_name)
-            else:
-                name = target_param_name
-                parent = self
-            setattr(parent, name, source_param)
-            self._adjust_bias(parent, source_param)
+        tied_keys = list(tied_keys.items())
+        for i, (target_param_name, source_param_name) in enumerate(tied_keys):
+            # Usually we tie a single target to a single source, but when both are missing we may later tie
+            # both the source and target to a third "backup" parameter that is present in the checkpoint, so we use
+            # a list here
+            target_param_names = [target_param_name]
+
+            # This is `from_pretrained` -> let's check symmetrically in case the source key is not present
             if missing_keys is not None:
+                remove_from_missing = True
                 source_is_there = source_param_name not in missing_keys
                 target_is_there = target_param_name not in missing_keys
-                # If we tied correctly, remove the target from the missing keys
-                if source_is_there:
-                    missing_keys.discard(target_param_name)
-                # If the source is not present, but the target is, the checkpoint is corrupted
-                # TODO: maybe we could simply tie in the opposite direction here instead of error?
-                elif target_is_there:
-                    raise ValueError(
-                        f"This checkpoint seem corrupted. The tied weights mapping for this model specifies to tie "
-                        f"{source_param_name} (which should be present and is not), to {target_param_name} (which is "
-                        f"present)."
+                # Both are already present -> it means the config is wrong and do not reflect the actual
+                # checkpoint -> let's raise a warning and do nothing
+                if source_is_there and target_is_there:
+                    logger.warning(
+                        f"The tied weights mapping and config for this model specifies to tie {source_param_name} to "
+                        f"{target_param_name}, but both are present in the checkpoints, so we will NOT tie them. "
+                        "You should update the config with `tie_word_embeddings=False` to silence this warning"
                     )
+                    # Skip to next iteration
+                    continue
+                # We're missing the source but we have the target -> we swap them, tying the parameter that exists
+                elif not source_is_there and target_is_there:
+                    target_param_name, source_param_name = source_param_name, target_param_name
+                    target_param_names = [target_param_name]
+                # Both are missing -> check other keys in case more than 2 keys are tied to the same weight
+                elif not source_is_there and not target_is_there:
+                    for target_backup, source_backup in tied_keys[i + 1 :]:
+                        # In case of more than 2 keys tied to the same weight, they are guaranteed to all have
+                        # the same source thanks to `get_expanded_tied_weights_keys` so this check is enough
+                        if source_backup == source_param_name:
+                            target_backup_is_there = target_backup not in missing_keys
+                            # If the target is present, we found the correct weight to tie into (we know the source is missing)
+                            if target_backup_is_there:
+                                source_param_name = target_backup
+                                # Append the source as well, since both are missing we'll tie both
+                                target_param_names.append(source_param_name)
+                                break
+                    # If we did not break from the loop, it was impossible to find a source key -> let's raise
+                    else:
+                        # TODO Cyril: here ideally we want to raise instead of warning, but will break our CI as we have
+                        # tests loading model from empty dicts to perform init checks - since we don't raise, add a flag
+                        # to NOT remove from missing keys as it's actually still missing
+                        remove_from_missing = False
+                        logger.warning(
+                            f"This checkpoint seem corrupted. The tied weights mapping for this model specifies to tie "
+                            f"{source_param_name} to {target_param_name}, but both are absent from the checkpoint, "
+                            "and we could not find another related tied weight for those keys"
+                        )
+
+            # Perform the actual tying
+            source_param = self.get_parameter_or_buffer(source_param_name)
+            for target_param_name in target_param_names:
+                if "." in target_param_name:
+                    parent_name, name = target_param_name.rsplit(".", 1)
+                    parent = self.get_submodule(parent_name)
+                else:
+                    name = target_param_name
+                    parent = self
+                # Tie the weights
+                setattr(parent, name, source_param)
+                self._adjust_bias(parent, source_param)
+                # Remove from missing if necesary
+                if missing_keys is not None and remove_from_missing:
+                    missing_keys.discard(target_param_name)
 
     def _adjust_bias(self, output_embeddings, input_embeddings):
         if getattr(output_embeddings, "bias", None) is not None and hasattr(output_embeddings, "weight"):

filename: src/transformers/models/bart/modeling_bart.py
patch: @@ -897,21 +897,6 @@ def __init__(self, config: BartConfig):
         # Initialize weights and apply final processing
         self.post_init()
 
-    def tie_weights(self, missing_keys: Optional[set[str]] = None, recompute_mapping: bool = True):
-        """We need to overload here to handle the wrong key saved in some main checkpoints."""
-        if self.config.tie_word_embeddings:
-            # Some model checkpoints like "facebook/bart-large-cnn"'s embedding weight is in decoder.embed_tokens,
-            # need check here, see issue #36247
-            if missing_keys is not None:
-                if "shared.weight" in missing_keys and "decoder.embed_tokens.weight" not in missing_keys:
-                    self.encoder.embed_tokens.weight = self.decoder.embed_tokens.weight
-                    self.shared.weight = self.decoder.embed_tokens.weight
-                    missing_keys.discard("encoder.embed_token.weight")
-                    missing_keys.discard("shared.weight")
-
-        # needs to be done after, otherwise it raises an Error because the correct weights are not present
-        super().tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping)
-
     def get_input_embeddings(self):
         return self.shared
 
@@ -1049,21 +1034,6 @@ def __init__(self, config: BartConfig):
         # Initialize weights and apply final processing
         self.post_init()
 
-    def tie_weights(self, missing_keys: Optional[set[str]] = None, recompute_mapping: bool = True):
-        """We need to overload here to handle the wrong key saved in some main checkpoints."""
-        if self.config.tie_word_embeddings:
-            # Some model checkpoints like "facebook/bart-large-cnn"'s embedding weight is in decoder.embed_tokens,
-            # need check here, see issue #36247
-            if missing_keys is not None:
-                if "model.shared.weight" in missing_keys and "model.decoder.embed_tokens.weight" not in missing_keys:
-                    self.model.encoder.embed_tokens.weight = self.model.decoder.embed_tokens.weight
-                    self.model.shared.weight = self.model.decoder.embed_tokens.weight
-                    missing_keys.discard("model.encoder.embed_token.weight")
-                    missing_keys.discard("model.shared.weight")
-
-        # needs to be done after, otherwise it raises an Error because the correct weights are not present
-        super().tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping)
-
     def resize_token_embeddings(
         self, new_num_tokens: int, pad_to_multiple_of: Optional[int] = None, mean_resizing: bool = True
     ) -> nn.Embedding:
@@ -1242,21 +1212,6 @@ def __init__(self, config: BartConfig, **kwargs):
         # Initialize weights and apply final processing
         self.post_init()
 
-    def tie_weights(self, missing_keys: Optional[set[str]] = None, recompute_mapping: bool = True):
-        """We need to overload here to handle the wrong key saved in some main checkpoints."""
-        if self.config.tie_word_embeddings:
-            # Some model checkpoints like "facebook/bart-large-cnn"'s embedding weight is in decoder.embed_tokens,
-            # need check here, see issue #36247
-            if missing_keys is not None:
-                if "model.shared.weight" in missing_keys and "model.decoder.embed_tokens.weight" not in missing_keys:
-                    self.model.encoder.embed_tokens.weight = self.model.decoder.embed_tokens.weight
-                    self.model.shared.weight = self.model.decoder.embed_tokens.weight
-                    missing_keys.discard("model.encoder.embed_token.weight")
-                    missing_keys.discard("model.shared.weight")
-
-        # needs to be done after, otherwise it raises an Error because the correct weights are not present
-        super().tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping)
-
     @auto_docstring
     def forward(
         self,
@@ -1388,21 +1343,6 @@ def __init__(self, config):
         # Initialize weights and apply final processing
         self.post_init()
 
-    def tie_weights(self, missing_keys: Optional[set[str]] = None, recompute_mapping: bool = True):
-        """We need to overload here to handle the wrong key saved in some main checkpoints."""
-        if self.config.tie_word_embeddings:
-            # Some model checkpoints like "facebook/bart-large-cnn"'s embedding weight is in decoder.embed_tokens,
-            # need check here, see issue #36247
-            if missing_keys is not None:
-                if "model.shared.weight" in missing_keys and "model.decoder.embed_tokens.weight" not in missing_keys:
-                    self.model.encoder.embed_tokens.weight = self.model.decoder.embed_tokens.weight
-                    self.model.shared.weight = self.model.decoder.embed_tokens.weight
-                    missing_keys.discard("model.encoder.embed_token.weight")
-                    missing_keys.discard("model.shared.weight")
-
-        # needs to be done after, otherwise it raises an Error because the correct weights are not present
-        super().tie_weights(missing_keys=missing_keys, recompute_mapping=recompute_mapping)
-
     @auto_docstring
     def forward(
         self,

filename: tests/models/openai/test_modeling_openai.py
patch: @@ -269,13 +269,6 @@ def test_model_from_pretrained(self):
         model = OpenAIGPTModel.from_pretrained(model_name)
         self.assertIsNotNone(model)
 
-    @unittest.skip("Tied weights mapping is reversed, so this is supposed to error out")
-    def test_correct_missing_keys(self):
-        # openai defines `_tied_weights_keys = {"transformer.tokens_embed.weight": "lm_head.weight"}` instead
-        # of the usual `_tied_weights_keys = {"lm_head.weight": "transformer.tokens_embed.weight"}`, so removing
-        # the head parameters actually removes the source weight, so this test is supposed to fail
-        pass
-
 
 @require_torch
 class OPENAIGPTModelLanguageGenerationTest(unittest.TestCase):

filename: tests/utils/test_modeling_utils.py
patch: @@ -92,6 +92,8 @@
     is_torch_npu_available,
 )
 
+from ..test_modeling_common import compare_state_dicts
+
 
 sys.path.append(str(Path(__file__).parent.parent.parent / "utils"))
 
@@ -184,10 +186,35 @@ def __init__(self, config):
         def forward(self, x):
             return self.linear_2(self.linear(x))
 
-        def tie_weights(self, missing_keys=None, **kwargs):
-            self.linear_2.weight = self.linear.weight
-            if missing_keys is not None:
-                missing_keys.discard("linear_2.weight")
+    class BaseModelWithMultipleTiedWeights(PreTrainedModel):
+        config_class = PreTrainedConfig
+        _tied_weights_keys = {"linear_2.weight": "linear.weight", "linear_3.weight": "linear.weight"}
+
+        def __init__(self, config):
+            super().__init__(config)
+            self.linear = nn.Linear(5, 5)
+            self.linear_2 = nn.Linear(5, 5)
+            self.linear_3 = nn.Linear(5, 5)
+            self.post_init()
+
+        def forward(self, x):
+            return self.linear_2(self.linear(x))
+
+    class BaseModelWithMultipleMixedTiedWeights(PreTrainedModel):
+        config_class = PreTrainedConfig
+        # Here the tied keys both refer to `linear.weight`, but they are inconsistent in the mapping, i.e. they
+        # are provided as a "circular" dependency
+        _tied_weights_keys = {"linear_2.weight": "linear.weight", "linear_3.weight": "linear_2.weight"}
+
+        def __init__(self, config):
+            super().__init__(config)
+            self.linear = nn.Linear(5, 5)
+            self.linear_2 = nn.Linear(5, 5)
+            self.linear_3 = nn.Linear(5, 5)
+            self.post_init()
+
+        def forward(self, x):
+            return self.linear_2(self.linear(x))
 
     class ModelWithHead(PreTrainedModel):
         base_model_prefix = "base"
@@ -258,11 +285,6 @@ def __init__(self, config):
         def forward(self, x):
             return self.decoder(self.base(x))
 
-        def tie_weights(self, missing_keys=None, **kwargs):
-            self.decoder.weight = self.base.linear.weight
-            if missing_keys is not None:
-                missing_keys.discard("decoder.weight")
-
     class Prepare4dCausalAttentionMaskModel(nn.Module):
         def forward(self, inputs_embeds):
             batch_size, seq_length, _ = inputs_embeds.shape
@@ -1372,6 +1394,148 @@ def test_tied_weights_reload(self):
             # Should only complain about the missing bias
             self.assertSetEqual(load_info["missing_keys"], {"decoder.bias"})
 
+    def test_tied_weights_can_load_symmetrically(self):
+        """Test that we can correctly load and tie weights even though the wrong key was saved."""
+        model = BaseModelWithTiedWeights(PreTrainedConfig())
+        # Just to be sure it's actually tied
+        self.assertIs(model.linear.weight, model.linear_2.weight, msg="Weights are not tied!")
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Save the config
+            with open(os.path.join(tmp_dir, "config.json"), "w") as f:
+                f.write(json.dumps(model.config.to_dict()))
+
+            state_dict = model.state_dict()
+            # Save using the wrong key
+            state_dict.pop("linear.weight")
+            safe_save_file(state_dict, os.path.join(tmp_dir, "model.safetensors"))
+
+            new_model, load_info = BaseModelWithTiedWeights.from_pretrained(tmp_dir, output_loading_info=True)
+            # Assert no missing keys
+            self.assertSetEqual(load_info["missing_keys"], set(), msg=f"{load_info['missing_keys']} are missing!")
+            # It's still the same weight
+            self.assertIs(new_model.linear.weight, new_model.linear_2.weight, msg="Weights are not tied!")
+
+            # Make sure both state dict are the same
+            compare_state_dicts(model.state_dict(), new_model.state_dict())
+
+    def test_tied_weights_can_load_symmetrically_multiple_keys(self):
+        """Test that we can correctly load and tie weights even though the wrong key was saved, when we
+        have more than 1 target to the same source."""
+        # First class is consistent in how they provide the source, second is not -> make sure it works in both cases
+        for model_class in [BaseModelWithMultipleTiedWeights, BaseModelWithMultipleMixedTiedWeights]:
+            with self.subTest(model_class.__name__):
+                model = model_class(PreTrainedConfig())
+                # Just to be sure it's actually tied
+                self.assertIs(model.linear.weight, model.linear_2.weight, msg="Weights are not tied!")
+                self.assertIs(model.linear.weight, model.linear_3.weight, msg="Weights are not tied!")
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    # Save the config
+                    with open(os.path.join(tmp_dir, "config.json"), "w") as f:
+                        f.write(json.dumps(model.config.to_dict()))
+
+                    state_dict = model.state_dict()
+                    # Keep only 1 of the 3 tied keys, but not the source (which is `linear.weight`)
+                    state_dict.pop("linear.weight")
+                    state_dict.pop("linear_3.weight")
+                    safe_save_file(state_dict, os.path.join(tmp_dir, "model.safetensors"))
+
+                    new_model, load_info = BaseModelWithMultipleTiedWeights.from_pretrained(
+                        tmp_dir, output_loading_info=True
+                    )
+                    # Assert no missing keys
+                    self.assertSetEqual(
+                        load_info["missing_keys"], set(), msg=f"{load_info['missing_keys']} are missing!"
+                    )
+                    # It's still the same weight
+                    self.assertIs(new_model.linear.weight, new_model.linear_2.weight, msg="Weights are not tied!")
+                    self.assertIs(new_model.linear.weight, new_model.linear_3.weight, msg="Weights are not tied!")
+
+                    # Make sure both state dict are the same
+                    compare_state_dicts(model.state_dict(), new_model.state_dict())
+
+                    # Now, do the same but try to keep `linear_2.weight` in the saved key instead of `linear_3.weight`
+                    # to make sure it does not matter
+                    state_dict = model.state_dict()
+                    # Keep only 1 of the 3 tied keys, but not the source (which is `linear.weight`)
+                    state_dict.pop("linear.weight")
+                    state_dict.pop("linear_2.weight")
+                    safe_save_file(state_dict, os.path.join(tmp_dir, "model.safetensors"))
+
+                    new_model, load_info = BaseModelWithMultipleTiedWeights.from_pretrained(
+                        tmp_dir, output_loading_info=True
+                    )
+                    # Assert no missing keys
+                    self.assertSetEqual(
+                        load_info["missing_keys"], set(), msg=f"{load_info['missing_keys']} are missing!"
+                    )
+                    # It's still the same weight
+                    self.assertIs(new_model.linear.weight, new_model.linear_2.weight, msg="Weights are not tied!")
+                    self.assertIs(new_model.linear.weight, new_model.linear_3.weight, msg="Weights are not tied!")
+
+                    # Make sure both state dict are the same
+                    compare_state_dicts(model.state_dict(), new_model.state_dict())
+
+    def test_tied_weights_are_not_tied_if_both_present(self):
+        """Test that if both the source and target of tied weights are present, we do NOT tie them, and instead
+        raise a warning"""
+        model = BaseModelWithTiedWeights(PreTrainedConfig())
+        # Just to be sure it's actually tied
+        self.assertIs(model.linear.weight, model.linear_2.weight, msg="Weights are not tied!")
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Save the config
+            with open(os.path.join(tmp_dir, "config.json"), "w") as f:
+                f.write(json.dumps(model.config.to_dict()))
+
+            state_dict = model.state_dict()
+            # Clone every param to make sure nothing is tied -> we save everything
+            state_dict = {k: v.clone() for k, v in state_dict.items()}
+            safe_save_file(state_dict, os.path.join(tmp_dir, "model.safetensors"))
+
+            logger = logging.get_logger("transformers.modeling_utils")
+            with CaptureLogger(logger) as cl:
+                new_model, load_info = BaseModelWithTiedWeights.from_pretrained(tmp_dir, output_loading_info=True)
+
+            # We should have raised a warning here saying that we will NOT tie the weights
+            self.assertIn("both are present in the checkpoints, so we will NOT tie them.", cl.out)
+            # Assert no missing keys
+            self.assertSetEqual(load_info["missing_keys"], set(), msg=f"{load_info['missing_keys']} are missing!")
+            # It should not be the same weight anymore
+            self.assertIsNot(
+                new_model.linear.weight, new_model.linear_2.weight, msg="Weights are tied but they should not!"
+            )
+
+            # Make sure both state dict are the same (the values are still the same, it's just not tied)
+            compare_state_dicts(model.state_dict(), new_model.state_dict())
+
+    def test_tied_weights_are_missing_if_both_absent(self):
+        """Test that if both the source and target of tied weights are absent, we do tie them, but they are missing"""
+        model = BaseModelWithTiedWeights(PreTrainedConfig())
+        # Just to be sure it's actually tied
+        self.assertIs(model.linear.weight, model.linear_2.weight, msg="Weights are not tied!")
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Save the config
+            with open(os.path.join(tmp_dir, "config.json"), "w") as f:
+                f.write(json.dumps(model.config.to_dict()))
+
+            state_dict = model.state_dict()
+            # Remove both from the state dict
+            state_dict.pop("linear.weight")
+            state_dict.pop("linear_2.weight")
+            safe_save_file(state_dict, os.path.join(tmp_dir, "model.safetensors"))
+
+            logger = logging.get_logger("transformers.modeling_utils")
+            with CaptureLogger(logger) as cl:
+                new_model, load_info = BaseModelWithTiedWeights.from_pretrained(tmp_dir, output_loading_info=True)
+
+            # We should have raised a warning here saying that we will NOT tie the weights
+            self.assertIn(
+                "This checkpoint seem corrupted. The tied weights mapping for this model specifies to tie", cl.out
+            )
+            # Assert both are in the missing keys
+            self.assertSetEqual(load_info["missing_keys"], {"linear.weight", "linear_2.weight"})
+            # They should still be tied though
+            self.assertIs(new_model.linear.weight, new_model.linear_2.weight, msg="Weights are not tied!")
+
     def test_unexpected_keys_warnings(self):
         model = ModelWithHead(PreTrainedConfig())
         logger = logging.get_logger("transformers.modeling_utils")