|
|
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()) |
|
|
|
