Spaces:
Sleeping
Sleeping
| # coding=utf-8 | |
| # Copyright 2023-present the HuggingFace Inc. team. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import os | |
| import tempfile | |
| import unittest | |
| import torch | |
| from parameterized import parameterized | |
| from transformers import AutoModelForCausalLM | |
| from peft import ( | |
| PeftModel, | |
| get_peft_model, | |
| get_peft_model_state_dict, | |
| prepare_model_for_int8_training, | |
| ) | |
| from .testing_common import PeftTestConfigManager | |
| # This has to be in the order: model_id, lora_kwargs, prefix_tuning_kwargs, prompt_encoder_kwargs, prompt_tuning_kwargs | |
| PEFT_MODELS_TO_TEST = [ | |
| ("hf-internal-testing/tiny-random-OPTForCausalLM", {"target_modules": ["q_proj", "v_proj"]}, {}, {}, {}), | |
| ] | |
| class PeftTestMixin: | |
| torch_device = "cuda" if torch.cuda.is_available() else "cpu" | |
| class PeftModelTester(unittest.TestCase, PeftTestMixin): | |
| r""" | |
| Test if the PeftModel behaves as expected. This includes: | |
| - test if the model has the expected methods | |
| We use parametrized.expand for debugging purposes to test each model individually. | |
| """ | |
| def test_attributes_parametrized(self, test_name, model_id, config_cls, config_kwargs): | |
| self._test_model_attr(model_id, config_cls, config_kwargs) | |
| def _test_model_attr(self, model_id, config_cls, config_kwargs): | |
| model = AutoModelForCausalLM.from_pretrained(model_id) | |
| config = config_cls( | |
| base_model_name_or_path=model_id, | |
| **config_kwargs, | |
| ) | |
| model = get_peft_model(model, config) | |
| self.assertTrue(hasattr(model, "save_pretrained")) | |
| self.assertTrue(hasattr(model, "from_pretrained")) | |
| self.assertTrue(hasattr(model, "push_to_hub")) | |
| def _test_prepare_for_training(self, model_id, config_cls, config_kwargs): | |
| model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device) | |
| config = config_cls( | |
| base_model_name_or_path=model_id, | |
| **config_kwargs, | |
| ) | |
| model = get_peft_model(model, config) | |
| dummy_input = torch.LongTensor([[1, 1, 1]]).to(self.torch_device) | |
| dummy_output = model.get_input_embeddings()(dummy_input) | |
| self.assertTrue(not dummy_output.requires_grad) | |
| # load with `prepare_model_for_int8_training` | |
| model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device) | |
| model = prepare_model_for_int8_training(model) | |
| for param in model.parameters(): | |
| self.assertTrue(not param.requires_grad) | |
| config = config_cls( | |
| base_model_name_or_path=model_id, | |
| **config_kwargs, | |
| ) | |
| model = get_peft_model(model, config) | |
| # For backward compatibility | |
| if hasattr(model, "enable_input_require_grads"): | |
| model.enable_input_require_grads() | |
| else: | |
| def make_inputs_require_grad(module, input, output): | |
| output.requires_grad_(True) | |
| model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) | |
| dummy_input = torch.LongTensor([[1, 1, 1]]).to(self.torch_device) | |
| dummy_output = model.get_input_embeddings()(dummy_input) | |
| self.assertTrue(dummy_output.requires_grad) | |
| def test_prepare_for_training_parametrized(self, test_name, model_id, config_cls, config_kwargs): | |
| self._test_prepare_for_training(model_id, config_cls, config_kwargs) | |
| def _test_save_pretrained(self, model_id, config_cls, config_kwargs): | |
| model = AutoModelForCausalLM.from_pretrained(model_id) | |
| config = config_cls( | |
| base_model_name_or_path=model_id, | |
| **config_kwargs, | |
| ) | |
| model = get_peft_model(model, config) | |
| model = model.to(self.torch_device) | |
| with tempfile.TemporaryDirectory() as tmp_dirname: | |
| model.save_pretrained(tmp_dirname) | |
| model_from_pretrained = AutoModelForCausalLM.from_pretrained(model_id) | |
| model_from_pretrained = PeftModel.from_pretrained(model_from_pretrained, tmp_dirname) | |
| # check if the state dicts are equal | |
| state_dict = get_peft_model_state_dict(model) | |
| state_dict_from_pretrained = get_peft_model_state_dict(model_from_pretrained) | |
| # check if same keys | |
| self.assertEqual(state_dict.keys(), state_dict_from_pretrained.keys()) | |
| # check if tensors equal | |
| for key in state_dict.keys(): | |
| self.assertTrue( | |
| torch.allclose( | |
| state_dict[key].to(self.torch_device), state_dict_from_pretrained[key].to(self.torch_device) | |
| ) | |
| ) | |
| # check if `adapter_model.bin` is present | |
| self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_model.bin"))) | |
| # check if `adapter_config.json` is present | |
| self.assertTrue(os.path.exists(os.path.join(tmp_dirname, "adapter_config.json"))) | |
| # check if `pytorch_model.bin` is not present | |
| self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "pytorch_model.bin"))) | |
| # check if `config.json` is not present | |
| self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json"))) | |
| def test_save_pretrained(self, test_name, model_id, config_cls, config_kwargs): | |
| self._test_save_pretrained(model_id, config_cls, config_kwargs) | |