| import os |
| import shutil |
| import tempfile |
| import unittest |
|
|
| import torch |
| from modelscope import Model |
| from peft.utils import WEIGHTS_NAME |
|
|
| from swift import LoRAConfig, SwiftModel |
|
|
|
|
| @unittest.skip |
| class TestSwift(unittest.TestCase): |
|
|
| def setUp(self): |
| print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) |
| self.tmp_dir = tempfile.TemporaryDirectory().name |
| if not os.path.exists(self.tmp_dir): |
| os.makedirs(self.tmp_dir) |
|
|
| def tearDown(self): |
| shutil.rmtree(self.tmp_dir) |
| super().tearDown() |
|
|
| def test_swift_multiple_adapters(self): |
| model = Model.from_pretrained('modelscope/Llama-2-7b-ms', device_map='auto') |
| lora_config = LoRAConfig(target_modules=['q_proj', 'k_proj', 'v_proj']) |
| model: SwiftModel = SwiftModel(model, config={'lora': lora_config}) |
| self.assertTrue(isinstance(model, SwiftModel)) |
| model.save_pretrained(self.tmp_dir, adapter_name=['lora']) |
| state_dict = model.state_dict() |
| with open(os.path.join(self.tmp_dir, 'configuration.json'), 'w') as f: |
| f.write('{}') |
| self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'lora'))) |
| self.assertTrue(os.path.exists(os.path.join(self.tmp_dir, 'lora', WEIGHTS_NAME))) |
| model = Model.from_pretrained('modelscope/Llama-2-7b-ms', device_map='auto') |
| model = SwiftModel.from_pretrained(model, self.tmp_dir, adapter_name=['lora'], device_map='auto') |
|
|
| state_dict2 = model.state_dict() |
| for key in state_dict: |
| self.assertTrue(key in state_dict2) |
| self.assertTrue(all(torch.isclose(state_dict[key], state_dict2[key]).flatten().detach().cpu())) |
|
|
| self.assertTrue(len(set(model.hf_device_map.values())) == torch.cuda.device_count()) |
|
|
