draft

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

embedding.py

@@ -47,7 +47,8 @@ class XLMRobertaEmbeddings(nn.Module):
         token_type_ids: (batch, seqlen)
         """
         batch_size, seqlen = input_ids.shape
-        embeddings = self.word_embeddings(input_ids)
+        print('input shape', input_ids.shape)
+        embeddings = self.word_embeddings(input_ids, task='sts')
         if self.max_position_embeddings > 0:
             if position_ids is None:
                 position_ids = create_position_ids_from_input_ids(input_ids, padding_idx=self.word_embeddings.padding_idx).to(input_ids.device)

mha.py

@@ -340,8 +340,8 @@ class CrossAttention(nn.Module):
 class LinearResidual(nn.Linear):
     """Wrap nn.Linear to return the residual as well. For compatibility with FusedDense."""
 
-    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        return super().forward(input), input
+    def forward(self, input: torch.Tensor, task=None) -> torch.Tensor:
+        return super().forward(input, task=task), input
 
 
 def _update_kv_cache(kv, inference_params, layer_idx):
@@ -450,6 +450,7 @@ class MHA(nn.Module):
 
         if fused_bias_fc and FusedDense is None:
             raise ImportError("fused_dense is not installed")
+        print('is this true', fused_bias_fc)
         linear_cls = nn.Linear if not fused_bias_fc else FusedDense
         linear_resid_cls = (
             LinearResidual if not fused_bias_fc else partial(FusedDense, return_residual=True)
@@ -646,7 +647,7 @@ class MHA(nn.Module):
             if not self.return_residual:
                 qkv = self.Wqkv(x)
             else:
-                qkv, x = self.Wqkv(x)
+                qkv, x = self.Wqkv(x, task='sts')
             if self.dwconv:
                 qkv = rearrange(
                     self.dwconv_qkv(rearrange(qkv, "b s d -> b d s"))[..., :-2], "b d s -> b s d"

modeling_lora.py

@@ -98,15 +98,15 @@ class LoRAParametrization(nn.Module):
         # to mimic the original implementation: A @ dropout(x), we do (A * dropout(ones)) @ x
         return A * self.lora_dropout(self.lora_dropout_mask)
 
-    def lora_forward(self, X):
-        assert self.current_task is not None
+    def lora_forward(self, X, current_task=None):
+        print('lora input shape', X.shape)
         return (
             X
             + torch.matmul(
                 *self.swap(
                     (
-                        self.lora_B[self.current_task],
-                        self.dropout_fn(self.lora_A[self.current_task]),
+                        self.lora_B[current_task],
+                        self.dropout_fn(self.lora_A[current_task]),
                     )
                 )
             ).view(X.shape)
@@ -114,7 +114,10 @@ class LoRAParametrization(nn.Module):
         )
 
     def forward(self, X):
-        return self.forward_fn(X)
+        print('forward input shape', X.shape, X)
+        out = self.forward_fn(X)
+        print(out.shape)
+        return out
 
     @property
     def current_task(self):
@@ -178,6 +181,7 @@ class LoRAParametrization(nn.Module):
         rank: int,
         dropout_p: float,
         alpha: float,
+        adaptation_map: dict,
     ):
         if isinstance(layer, nn.Linear):
             parametrize.register_parametrization(
@@ -191,6 +195,16 @@ class LoRAParametrization(nn.Module):
                     alpha=alpha,
                 ),
             )
+            original_forward = layer.forward
+
+            def new_forward(self, input, task):
+                print('an aq mitxari aba')
+                output = original_forward(input, task=task)
+                weight = self.parametrizations.weight(self.weight, task)
+                return nn.functional.linear(input, weight, self.bias)
+
+            layer.forward = new_forward.__get__(layer, layer.__class__)
+
         elif isinstance(layer, nn.Embedding):
             parametrize.register_parametrization(
                 layer,
@@ -203,6 +217,23 @@ class LoRAParametrization(nn.Module):
                     alpha=alpha,
                 ),
             )
+            original_forward = layer.forward
+
+            def new_forward(self, input, task):
+                print('input here', input, input.shape)
+                print('func', original_forward)
+                # original_forward['parametrizations'] = None
+                # print('funcc', original_forward.__dict__)
+                output = original_forward(input)
+                print(output.shape, 'output shape')
+                task_idx = adaptation_map[task] if task else None
+                if task_idx:
+                    output = self.parametrizations.weight[0].lora_forward(output, current_task=task_idx)
+                    print('thats it')
+                return output
+
+            layer.forward = new_forward.__get__(layer, layer.__class__)
+
 
     @staticmethod
     def select_task_for_layer(layer: nn.Module, task_idx: Optional[int] = None):
@@ -247,6 +278,13 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         self._task_idx = None
         # By default, disable LoRA until it's specified which adapter/task to use
         self.current_task = None
+        for name, param in super().named_parameters():
+            if name == 'roberta.encoder.layers.22.mixer.Wqkv.parametrizations.weight.0.lora_A':
+                print('A0', param[0])
+                print('A1', param[1])
+            if name == 'roberta.encoder.layers.22.mixer.Wqkv.parametrizations.weight.0.lora_B':
+                print('B0', param[0])
+                print('B1', param[1])
 
     @property
     def main_params_trainable(self):
@@ -300,6 +338,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 rank=rank,
                 dropout_p=dropout_p,
                 alpha=alpha,
+                adaptation_map=self._adaptation_map,
             )
         )
 

modeling_xlm_roberta.py

@@ -204,7 +204,7 @@ class XLMRobertaEncoder(nn.Module):
     def gradient_checkpointing(self, value):
         self._grad_checkpointing = value
 
-    def forward(self, hidden_states, key_padding_mask=None, subset_mask=None):
+    def forward(self, hidden_states, key_padding_mask=None, subset_mask=None, task=None):
         """If subset_mask is not None, we only want output for the subset of the sequence.
         This means that we only compute the last layer output for these tokens.
         subset_mask: (batch, seqlen), dtype=torch.bool