diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 3cc3d3dd69..089cda9e4c 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -1421,7 +1421,7 @@ def __init__( # to make sure classifier layer is trainable; this may add a new ModulesToSaveWrapper _set_trainable(self, adapter_name) - def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: + def add_adapter(self, adapter_name: str, peft_config: PeftConfig, low_cpu_mem_usage: bool = False) -> None: """ Add an adapter to the model based on the passed configuration. @@ -1437,6 +1437,10 @@ def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: The name of the adapter to be added. peft_config ([`PeftConfig`]): The configuration of the adapter to be added. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the process when loading saved + adapters. Don't use this option when creating a new PEFT adapter for training. + """ # ensure that additional adapters also add the classifier layer to modules_to_save if hasattr(peft_config, "modules_to_save"): @@ -1446,7 +1450,7 @@ def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: else: peft_config.modules_to_save.extend(classifier_module_names) - return super().add_adapter(adapter_name, peft_config) + return super().add_adapter(adapter_name, peft_config, low_cpu_mem_usage=low_cpu_mem_usage) def forward( self, @@ -2140,7 +2144,7 @@ def __init__( # to make sure classifier layer is trainable; this may add a new ModulesToSaveWrapper _set_trainable(self, adapter_name) - def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: + def add_adapter(self, adapter_name: str, peft_config: PeftConfig, low_cpu_mem_usage: bool = False) -> None: """ Add an adapter to the model based on the passed configuration. @@ -2156,6 +2160,10 @@ def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: The name of the adapter to be added. peft_config ([`PeftConfig`]): The configuration of the adapter to be added. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the process when loading saved + adapters. Don't use this option when creating a new PEFT adapter for training. + """ # ensure that additional adapters also add the classifier layer to modules_to_save if hasattr(peft_config, "modules_to_save"): @@ -2165,7 +2173,7 @@ def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: else: peft_config.modules_to_save.extend(classifier_module_names) - return super().add_adapter(adapter_name, peft_config) + return super().add_adapter(adapter_name, peft_config, low_cpu_mem_usage=low_cpu_mem_usage) def forward( self, @@ -2357,7 +2365,7 @@ def __init__( # to make sure classifier layer is trainable; this may add a new ModulesToSaveWrapper _set_trainable(self, adapter_name) - def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: + def add_adapter(self, adapter_name: str, peft_config: PeftConfig, low_cpu_mem_usage: bool = False) -> None: """ Add an adapter to the model based on the passed configuration. @@ -2373,6 +2381,10 @@ def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: The name of the adapter to be added. peft_config ([`PeftConfig`]): The configuration of the adapter to be added. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the process when loading saved + adapters. Don't use this option when creating a new PEFT adapter for training. + """ # ensure that additional adapters also add the classifier layer to modules_to_save if hasattr(peft_config, "modules_to_save"): @@ -2382,7 +2394,7 @@ def add_adapter(self, adapter_name: str, peft_config: PeftConfig) -> None: else: peft_config.modules_to_save.extend(qa_module_names) - return super().add_adapter(adapter_name, peft_config) + return super().add_adapter(adapter_name, peft_config, low_cpu_mem_usage=low_cpu_mem_usage) def forward( self, diff --git a/tests/test_decoder_models.py b/tests/test_decoder_models.py index 6204db93f6..ad10805baf 100644 --- a/tests/test_decoder_models.py +++ b/tests/test_decoder_models.py @@ -200,6 +200,9 @@ def test_save_pretrained_selected_adapters(self, test_name, model_id, config_cls def test_save_pretrained_selected_adapters_pickle(self, test_name, model_id, config_cls, config_kwargs): self._test_save_pretrained_selected_adapters(model_id, config_cls, config_kwargs, safe_serialization=False) + def test_load_model_low_cpu_mem_usage(self): + self._test_load_model_low_cpu_mem_usage(PEFT_DECODER_MODELS_TO_TEST[0], LoraConfig, {}) + @parameterized.expand( PeftTestConfigManager.get_grid_parameters(FULL_GRID, filter_params_func=skip_oft_or_hra_and_gpt2) ) diff --git a/tests/test_encoder_decoder_models.py b/tests/test_encoder_decoder_models.py index 2b9f68fc21..d013bbad99 100644 --- a/tests/test_encoder_decoder_models.py +++ b/tests/test_encoder_decoder_models.py @@ -82,6 +82,9 @@ def test_save_pretrained_selected_adapters(self, test_name, model_id, config_cls def test_save_pretrained_selected_adapters_pickle(self, test_name, model_id, config_cls, config_kwargs): self._test_save_pretrained_selected_adapters(model_id, config_cls, config_kwargs, safe_serialization=False) + def test_load_model_low_cpu_mem_usage(self): + self._test_load_model_low_cpu_mem_usage(PEFT_ENCODER_DECODER_MODELS_TO_TEST[0], LoraConfig, {}) + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) def test_from_pretrained_config_construction(self, test_name, model_id, config_cls, config_kwargs): self._test_from_pretrained_config_construction(model_id, config_cls, config_kwargs) diff --git a/tests/test_feature_extraction_models.py b/tests/test_feature_extraction_models.py index 05cbeb73d4..884c8cf7e0 100644 --- a/tests/test_feature_extraction_models.py +++ b/tests/test_feature_extraction_models.py @@ -17,7 +17,7 @@ from parameterized import parameterized from transformers import AutoModel -from peft import PrefixTuningConfig, PromptLearningConfig +from peft import LoraConfig, PrefixTuningConfig, PromptLearningConfig from .testing_common import PeftCommonTester, PeftTestConfigManager @@ -99,6 +99,9 @@ def test_save_pretrained(self, test_name, model_id, config_cls, config_kwargs): def test_save_pretrained_selected_adapters(self, test_name, model_id, config_cls, config_kwargs): self._test_save_pretrained_selected_adapters(model_id, config_cls, config_kwargs) + def test_load_model_low_cpu_mem_usage(self): + self._test_load_model_low_cpu_mem_usage(PEFT_FEATURE_EXTRACTION_MODELS_TO_TEST[0], LoraConfig, {}) + @parameterized.expand(PeftTestConfigManager.get_grid_parameters(FULL_GRID)) def test_from_pretrained_config_construction(self, test_name, model_id, config_cls, config_kwargs): self._test_from_pretrained_config_construction(model_id, config_cls, config_kwargs)