Fix tie_weights skipping logic is not tied to model thread scope

src/transformers/initialization.py

@@ -15,6 +15,7 @@
 import sys
 from collections import defaultdict
 from contextlib import contextmanager
+from contextvars import ContextVar
 
 import torch
 
@@ -38,6 +39,19 @@
     "sparse_": torch.nn.init.sparse_,
 }
 
+# Track the current no-tie scope per execution context so concurrent model loads
+# do not leak tie_weights suppression across threads.
+_SKIP_TIE_WEIGHTS_SCOPE: ContextVar[object | None] = ContextVar("_SKIP_TIE_WEIGHTS_SCOPE", default=None)
+
+
+def should_skip_tie_weights(model) -> bool:
+    scope = _SKIP_TIE_WEIGHTS_SCOPE.get()
+    if scope is None:
+        return False
+
+    # Only skip tying for the model instance created inside the active scope.
+    return getattr(model, "_skip_tie_weights_scope", None) is scope
+
 
 def uniform_(
     tensor: torch.Tensor, a: float = 0.0, b: float = 1.0, generator: torch.Generator | None = None
@@ -287,16 +301,10 @@ def no_tie_weights():
     weights in the state_dict during `from_pretrained`, and otherwise tying them would remove them from it, as it's
     called in `post_init` when instantiating.
     """
-    from .modeling_utils import PreTrainedModel
-
-    def empty_func(*args, **kwargs):
-        pass
-
+    # Use an opaque scope token so nested or concurrent loads can identify only
+    # the models instantiated under this context manager.
+    state_token = _SKIP_TIE_WEIGHTS_SCOPE.set(object())
     try:
-        original_tie_weights = PreTrainedModel.tie_weights
-        PreTrainedModel.tie_weights = empty_func
-
         yield
     finally:
-        # Set back the original
-        PreTrainedModel.tie_weights = original_tie_weights
+        _SKIP_TIE_WEIGHTS_SCOPE.reset(state_token)

src/transformers/modeling_utils.py

@@ -1324,6 +1324,10 @@ def post_init(self):
             if no_split := getattr(module, "_no_split_modules", None):
                 self._no_split_modules.update(no_split)
 
+        # Preserve the current no-tie scope on this instance so only the model
+        # being initialized in that scope skips tie_weights().
+        self._skip_tie_weights_scope = init._SKIP_TIE_WEIGHTS_SCOPE.get()
+
         # Maybe initialize the weights and tie the keys
         self.init_weights()
         self._backward_compatibility_gradient_checkpointing()
@@ -2517,6 +2521,9 @@ def tie_weights(self, missing_keys: set[str] | None = None, recompute_mapping: b
         `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.
         """
+        if init.should_skip_tie_weights(self):
+            return
+
         # 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

tests/utils/test_modeling_utils.py

@@ -218,6 +218,31 @@ def __init__(self, config):
         def forward(self, x):
             return self.linear_2(self.linear(x))
 
+    class DummyModelWithTiedEmbeddings(PreTrainedModel):
+        config_class = PreTrainedConfig
+        _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}
+
+        def __init__(self, config):
+            super().__init__(config)
+            self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+            self.post_init()
+
+        def get_input_embeddings(self):
+            return self.embed_tokens
+
+        def set_input_embeddings(self, value):
+            self.embed_tokens = value
+
+        def get_output_embeddings(self):
+            return self.lm_head
+
+        def set_output_embeddings(self, value):
+            self.lm_head = value
+
+        def forward(self, input_ids):
+            return self.lm_head(self.embed_tokens(input_ids))
+
     class ModelWithHead(PreTrainedModel):
         base_model_prefix = "base"
         config_class = PreTrainedConfig
@@ -414,6 +439,23 @@ def tearDown(self):
         torch.set_default_dtype(self.old_dtype)
         super().tearDown()
 
+    def _build_missing_tied_embeddings_checkpoint(self, tmp_dir):
+        reference_model = DummyModelWithTiedEmbeddings(
+            PreTrainedConfig(vocab_size=11, hidden_size=7, tie_word_embeddings=True)
+        )
+        reference_model.config.save_pretrained(tmp_dir)
+
+        state_dict = reference_model.state_dict()
+        del state_dict["lm_head.weight"]
+        safe_save_file(state_dict, os.path.join(tmp_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"})
+        return reference_model
+
+    def _assert_tied_embeddings_load_succeeded(self, model, reference_model):
+        self.assertIs(model.lm_head.weight, model.embed_tokens.weight, msg="Weights are not tied!")
+        for name, value in model.state_dict().items():
+            self.assertNotEqual(value.device.type, "meta", msg=f"{name} is still on meta!")
+        compare_state_dicts(reference_model.state_dict(), model.state_dict())
+
     @require_torch
     def test_get_total_byte_count_does_not_require_process_group(self):
         model = BaseModel(PreTrainedConfig())
@@ -1602,6 +1644,77 @@ def test_tied_weights_are_always_tied_from_config(self):
             model = LlamaForCausalLM._from_config(copy.deepcopy(config))
             self.assertTrue(model.lm_head.weight is not model.model.embed_tokens.weight)
 
+    def test_no_tie_weights_is_thread_local_during_concurrent_from_pretrained(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            reference_model = self._build_missing_tied_embeddings_checkpoint(tmp_dir)
+            first_loader_initialized = threading.Event()
+            release_first_loader = threading.Event()
+            first_loader_lock = threading.Lock()
+            results = []
+            errors = []
+            first_loader_claimed = False
+            original_init = DummyModelWithTiedEmbeddings.__init__
+
+            def instrumented_init(model_self, config):
+                original_init(model_self, config)
+
+                nonlocal first_loader_claimed
+                with first_loader_lock:
+                    should_block = not first_loader_claimed
+                    if should_block:
+                        first_loader_claimed = True
+
+                if should_block:
+                    first_loader_initialized.set()
+                    if not release_first_loader.wait(timeout=10):
+                        raise TimeoutError("Timed out waiting for the first loader to resume.")
+
+            def worker():
+                try:
+                    model, loading_info = DummyModelWithTiedEmbeddings.from_pretrained(
+                        tmp_dir, output_loading_info=True
+                    )
+                    results.append((model, loading_info))
+                except Exception as error:
+                    errors.append(error)
+
+            first_thread = threading.Thread(target=worker)
+            second_thread = threading.Thread(target=worker)
+
+            try:
+                with patch.object(DummyModelWithTiedEmbeddings, "__init__", new=instrumented_init):
+                    first_thread.start()
+                    self.assertTrue(first_loader_initialized.wait(timeout=10))
+
+                    second_thread.start()
+                    second_thread.join(timeout=20)
+                    self.assertFalse(second_thread.is_alive())
+            finally:
+                release_first_loader.set()
+                first_thread.join(timeout=20)
+                second_thread.join(timeout=20)
+
+            self.assertFalse(first_thread.is_alive())
+            self.assertFalse(second_thread.is_alive())
+            self.assertEqual(errors, [])
+            self.assertEqual(len(results), 2)
+
+            for model, loading_info in results:
+                self.assertSetEqual(loading_info["missing_keys"], set())
+                self._assert_tied_embeddings_load_succeeded(model, reference_model)
+
+    def test_no_tie_weights_is_model_specific_during_nested_from_pretrained(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            reference_model = self._build_missing_tied_embeddings_checkpoint(tmp_dir)
+
+            # `from_pretrained` uses its own no-tie scope while instantiating. An
+            # outer active scope must not suppress the final tie_weights() call.
+            with init.no_tie_weights():
+                model, load_info = DummyModelWithTiedEmbeddings.from_pretrained(tmp_dir, output_loading_info=True)
+
+                self.assertSetEqual(load_info["missing_keys"], set())
+                self._assert_tied_embeddings_load_succeeded(model, reference_model)
+
     def test_unexpected_keys_warnings(self):
         model = ModelWithHead(PreTrainedConfig(tie_word_embeddings=True))
         logger = logging.get_logger("transformers.modeling_utils")