wehub-resource-sync
45 lines
1.8 KiB
Python
45 lines
1.8 KiB
Python
import os
|
|
import shutil
|
|
import tempfile
|
|
import torch
|
|
import unittest
|
|
from modelscope import Model
|
|
from peft.utils import WEIGHTS_NAME
|
|
|
|
from swift.tuners 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())
|