Fix load_adapter OOM caused by full-model warmup sizing (#46145)

* Fix load_adapter OOM caused by full-model warmup sizing

load_adapter passed every named parameter on the model, including the
base model, as expected_keys to _load_pretrained_model. Downstream,
caching_allocator_warmup summed those into a full base-model byte count
and issued a single same-size allocation on top of the already-resident
base model, OOMing whenever the base model occupies more than ~half of
GPU memory.

The file already defined an is_adapter_key helper for identifying
parameters belonging to the freshly-injected adapter, but it was declared
after the _load_pretrained_model call. Hoist the helper above the call
and apply it to expected_keys.

Adds a regression test that captures the device map passed to
caching_allocator_warmup during load_adapter and asserts it contains only
adapter-owned parameter names, not base-model names.

* Address review: use unittest.mock.patch and expand test docstring

Author: Yooniel

src/transformers/integrations/peft.py

@@ -583,6 +583,13 @@ def load_adapter(
             # Create and add fresh new adapters into the model, unless the weights are hotswapped
             inject_adapter_in_model(peft_config, self, adapter_name)
 
+        adapter_key_markers = {adapter_name}
+        if peft_config is not None and getattr(peft_config, "peft_type", None) is not None:
+            adapter_key_markers.add(peft_config.peft_type.value.lower())
+
+        def is_adapter_key(key: str) -> bool:
+            return any(marker in key for marker in adapter_key_markers)
+
         if not self._hf_peft_config_loaded:
             self._hf_peft_config_loaded = True
 
@@ -670,9 +677,9 @@ def load_adapter(
             state_dict=adapter_state_dict,
             checkpoint_files=checkpoint_files,
             load_config=load_config,
-            # pass expected keys explicitly, otherwise they are determined from the state_dict, which can contain
-            # unexpected entries, like "layer.SCB" from a bnb layer.
-            expected_keys=[n for n, _ in self.named_parameters()],
+            # Pass expected keys explicitly while excluding non-adapter parameters.
+            # Otherwise `caching_allocator_warmup` sizes for the full base model.
+            expected_keys=[n for n, _ in self.named_parameters() if is_adapter_key(n)],
         )
 
         if peft_config.inference_mode:
@@ -683,13 +690,6 @@ def load_adapter(
                 if isinstance(module, BaseTunerLayer):
                     module.requires_grad_(False)
 
-        adapter_key_markers = {adapter_name}
-        if peft_config is not None and getattr(peft_config, "peft_type", None) is not None:
-            adapter_key_markers.add(peft_config.peft_type.value.lower())
-
-        def is_adapter_key(key: str) -> bool:
-            return any(marker in key for marker in adapter_key_markers)
-
         loading_info.missing_keys = {k for k in loading_info.missing_keys if is_adapter_key(k)}
 
         log_state_dict_report(

tests/peft_integration/test_peft_integration.py

@@ -19,6 +19,7 @@
 import tempfile
 import unittest
 from pathlib import Path
+from unittest.mock import patch
 
 from datasets import Dataset, DatasetDict
 from huggingface_hub import hf_hub_download
@@ -694,6 +695,58 @@ def test_peft_add_adapter_with_state_dict_low_cpu_mem_usage(self):
                 # after loading, no meta device should be remaining
                 self.assertFalse(any((p.device.type == "meta") for p in model.parameters()))
 
+    def test_peft_load_adapter_warmup_uses_adapter_expected_keys(self):
+        """
+        Check that adapter loading only warms up memory for adapter parameters by capturing the device map passed to
+        `caching_allocator_warmup`.
+
+        Note: this test depends on `_load_pretrained_model` calling `caching_allocator_warmup`; if that internal
+        contract changes, update the test to keep checking the keys used for warmup sizing.
+        """
+        from peft import LoraConfig
+
+        import transformers.modeling_utils as modeling_utils
+
+        adapter_name = "warmup_test_adapter"
+        adapter_key_markers = (adapter_name, "lora")
+
+        for model_id in self.transformers_test_model_ids:
+            for transformers_class in self.transformers_test_model_classes:
+                model = transformers_class.from_pretrained(model_id).to(torch_device)
+
+                peft_config = LoraConfig()
+                template_model = transformers_class.from_pretrained(model_id)
+                template_model.add_adapter(LoraConfig(), adapter_name=adapter_name)
+                dummy_state_dict = {
+                    name: torch.zeros_like(param)
+                    for name, param in template_model.named_parameters()
+                    if any(marker in name for marker in adapter_key_markers)
+                }
+                del template_model
+                self.assertTrue(dummy_state_dict)
+
+                captured_device_maps = []
+
+                def capture_warmup(model, expanded_device_map, hf_quantizer):
+                    captured_device_maps.append(dict(expanded_device_map))
+
+                with patch.object(modeling_utils, "caching_allocator_warmup", side_effect=capture_warmup):
+                    with CaptureLogger(logging.get_logger("transformers.integrations.peft")):
+                        model.load_adapter(
+                            adapter_state_dict=dummy_state_dict,
+                            adapter_name=adapter_name,
+                            peft_config=peft_config,
+                        )
+
+                self.assertTrue(captured_device_maps)
+                warmed_keys = set().union(*(device_map.keys() for device_map in captured_device_maps))
+                self.assertTrue(warmed_keys)
+
+                unexpected_base_keys = [
+                    key for key in warmed_keys if not any(marker in key for marker in adapter_key_markers)
+                ]
+                self.assertEqual(unexpected_base_keys, [])
+
     def test_peft_from_pretrained_hub_kwargs(self):
         """
         Tests different combinations of PEFT model + from_pretrained + hub kwargs