lora bugfix (#16)

- fix: lora bug (4c504d33aca884a998533b089cb905b597a82467)


Co-authored-by: Jack Min Ong <[email protected]>

modeling_lora.py

@@ -11,7 +11,7 @@ from torch import nn
 from torch.nn import Parameter
 from transformers import PretrainedConfig
 
-from .modeling_xlm_roberta import XLMRobertaFlashConfig, XLMRobertaModel
+from .modeling_xlm_roberta import XLMRobertaFlashConfig, XLMRobertaModel, XLMRobertaPreTrainedModel
 
 
 LORA_NO_UPDATE = '__lora_no_update__'
@@ -210,13 +210,19 @@ class LoRAParametrization(nn.Module):
             layer.current_task = task_idx
 
 
-class XLMRobertaLoRA(XLMRobertaModel):
+class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     def __init__(
         self,
         config: XLMRobertaFlashConfig,
+        roberta: Optional[XLMRobertaModel] = None
     ):
         super().__init__(config)
 
+        if roberta is None:
+            self.roberta = XLMRobertaModel(config)
+        else:
+            self.roberta = roberta
+
         self._lora_adaptations = config.lora_adaptations
         if (
             not isinstance(self._lora_adaptations, list)
@@ -231,7 +237,6 @@ class XLMRobertaLoRA(XLMRobertaModel):
         self._rank = config.lora_rank
         self._dropout_p = config.lora_dropout_p
         self._alpha = config.lora_alpha
-
         self._register_lora(
             num_adaptations=len(self._lora_adaptations),
             rank=self._rank,
@@ -284,9 +289,8 @@ class XLMRobertaLoRA(XLMRobertaModel):
                 pretrained_model_name_or_path, *model_args, **kwargs
             )
         else:
-            dtype = config.torch_dtype if config.torch_dtype else torch.bfloat16
-            torch.set_default_dtype(dtype)
-            return cls(config)
+            roberta = XLMRobertaModel.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+            return cls(config, roberta=roberta)
 
     def _register_lora(self, num_adaptations, rank, dropout_p, alpha):
         self.apply(
@@ -331,7 +335,8 @@ class XLMRobertaLoRA(XLMRobertaModel):
     def forward(self, *args, task: Union[str, None] = LORA_NO_UPDATE, **kwargs):
         if task != LORA_NO_UPDATE:
             self.current_task = task
-        return super().forward(*args, **kwargs)
+
+        return self.roberta(*args, **kwargs)
 
     def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
         for _, param in self.named_parameters(recurse=recurse):
@@ -373,4 +378,4 @@ class XLMRobertaLoRA(XLMRobertaModel):
                 )
             self.current_task = task
 
-        return super().encode(*args, **kwargs)
+        return self.roberta.encode(*args, **kwargs)