File size: 1,793 Bytes
7feac49 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
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())
|