44 lines
1.8 KiB
Python
44 lines
1.8 KiB
Python
import math
|
|
import torch
|
|
import unittest
|
|
from modelscope import Model, Preprocessor
|
|
from torch import nn
|
|
|
|
from swift.tuners import LoRAConfig, Swift
|
|
|
|
|
|
class TestMergedLinear(unittest.TestCase):
|
|
|
|
def test_swift_lora_forward(self):
|
|
|
|
from swift.tuners.lora import MergedLinear
|
|
|
|
def reset_parameters(self):
|
|
nn.Linear.reset_parameters(self)
|
|
if hasattr(self, 'lora_A'):
|
|
# initialize A the same way as the default for nn.Linear and B to zero
|
|
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
|
|
nn.init.ones_(self.lora_B)
|
|
|
|
MergedLinear.reset_parameters = reset_parameters
|
|
|
|
model = Model.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
|
|
preprocessor = Preprocessor.from_pretrained('damo/nlp_structbert_sentence-similarity_chinese-base')
|
|
inputs = preprocessor('how are you')
|
|
lora_config = LoRAConfig(
|
|
target_modules=['query', 'key', 'value'], use_merged_linear=True, enable_lora=[True, True, True])
|
|
outputs = model(**inputs)
|
|
model = Swift.prepare_model(model, config=lora_config)
|
|
model.eval()
|
|
outputs_lora = model(**inputs)
|
|
model.deactivate_adapter('default')
|
|
outputs_deactivate = model(**inputs)
|
|
model.activate_adapter('default')
|
|
outputs_reactivate = model(**inputs)
|
|
Swift.merge_and_unload(model)
|
|
outputs_merged = model(**inputs)
|
|
self.assertTrue(torch.allclose(outputs.logits, outputs_deactivate.logits))
|
|
self.assertTrue(not torch.allclose(outputs.logits, outputs_lora.logits))
|
|
self.assertTrue(torch.allclose(outputs_lora.logits, outputs_reactivate.logits))
|
|
self.assertTrue(torch.allclose(outputs_lora.logits, outputs_merged.logits, atol=1e-4))
|