2-adapter-tuning

block.py

@@ -233,7 +233,7 @@ class Block(nn.Module):
                     is_rms_norm=isinstance(self.norm1, RMSNorm),
                 )
             if not isinstance(self.mlp, nn.Identity):
-                mlp_out = self.mlp(hidden_states, task_type=mixer_kwargs.get('task_type'))
+                mlp_out = self.mlp(hidden_states, cu_adapter_mask=mixer_kwargs.get('cu_adapter_mask'))
                 if self.return_residual:  # mlp out is actually a pair here
                     mlp_out, hidden_states = mlp_out
                 if not self.fused_dropout_add_ln:

embedding.py

@@ -40,15 +40,25 @@ class XLMRobertaEmbeddings(nn.Module):
         if self.type_vocab_size > 0:
             self.token_type_embeddings = nn.Embedding(type_vocab_size, embed_dim, **factory_kwargs)
 
-    def forward(self, input_ids, position_ids=None, token_type_ids=None, task_type=None):
+    def forward(self, input_ids, position_ids=None, token_type_ids=None, adapter_mask=None):
         """
         input_ids: (batch, seqlen)
         position_ids: (batch, seqlen)
         token_type_ids: (batch, seqlen)
         """
         batch_size, seqlen = input_ids.shape
-        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-        embeddings = self.word_embeddings(input_ids, **lora_kwargs)
+        if adapter_mask is not None:
+            unique_tasks = torch.unique(adapter_mask)
+            embedding_dtype = next(self.word_embeddings.parameters()).dtype
+            embeddings = torch.empty(*input_ids.shape, self.word_embeddings.embedding_dim,
+                                            dtype=embedding_dtype, device=input_ids.device)
+            for task_id in unique_tasks:
+                task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                task_input_ids = input_ids[task_indices]
+                task_embeddings = self.word_embeddings(task_input_ids, task_id=task_id)
+                embeddings[task_indices] = task_embeddings
+        else:
+            embeddings = self.word_embeddings(input_ids)
         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)
@@ -58,6 +68,14 @@ class XLMRobertaEmbeddings(nn.Module):
         if self.type_vocab_size > 0:
             if token_type_ids is None:
                 token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=input_ids.device)
-            token_type_embeddings = self.token_type_embeddings(token_type_ids, **lora_kwargs)
-            embeddings = embeddings + token_type_embeddings
+
+            if adapter_mask is not None:
+                unique_tasks = torch.unique(adapter_mask)
+                for task_id in unique_tasks:
+                    task_token_type_embeddings = self.token_type_embeddings(token_type_ids, task_id=task_id)
+                    task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                    embeddings[task_indices] = embeddings[task_indices] + task_token_type_embeddings
+            else:
+                token_type_embeddings = self.token_type_embeddings(token_type_ids)
+                embeddings = embeddings + token_type_embeddings
         return embeddings

mha.py

@@ -590,7 +590,7 @@ class MHA(nn.Module):
         max_seqlen=None,
         mixer_subset=None,
         inference_params=None,
-        task_type=None,
+        cu_adapter_mask=None,
         **kwargs,
     ):
         """
@@ -643,15 +643,31 @@ class MHA(nn.Module):
             inference_params.max_sequence_len if inference_params is not None else max_seqlen
         )
         batch, seqlen = x.shape[:2]
+        lora_kwargs = {}
         if not self.cross_attn and self.num_heads_kv == self.num_heads:
             assert x_kv is None and mixer_subset is None
-            lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-            if not self.return_residual:
-                qkv = self.Wqkv(x, **lora_kwargs)
+
+            if cu_adapter_mask is not None:
+                unique_tasks = torch.unique(cu_adapter_mask)
+                qkv_dtype = next(self.Wqkv.parameters()).dtype
+                qkv = torch.empty(x.shape[0], self.Wqkv.out_features,
+                                         dtype=qkv_dtype, device=x.device)
+                for task_id in unique_tasks:
+                    task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                    task_tensor = x[task_indices]
+                    if not self.return_residual:
+                        task_qkv = self.Wqkv(task_tensor, task_id=task_id)
+                    else:
+                        task_qkv, _ = self.Wqkv(task_tensor, task_id=task_id, residual=True)
+                    qkv[task_indices] = task_qkv
             else:
-                if lora_kwargs:
-                    lora_kwargs['residual'] = True
-                qkv, x = self.Wqkv(x, **lora_kwargs)
+                if not self.return_residual:
+                    qkv = self.Wqkv(x)
+                else:
+                    if hasattr(self.Wqkv, 'parametrizations'):
+                        qkv, x = self.Wqkv(x, residual=True)
+                    else:
+                        qkv, x = self.Wqkv(x)
 
             if self.dwconv:
                 qkv = rearrange(
@@ -738,6 +754,17 @@ class MHA(nn.Module):
             else:
                 context = self._apply_rotary_update_kvcache_attention(q, kv, inference_params)
 
-        lora_kwargs.pop('residual', None)
-        out = self.out_proj(rearrange(context, "... h d -> ... (h d)"), **lora_kwargs)
+        inp = rearrange(context, "... h d -> ... (h d)")
+        if cu_adapter_mask is not None:
+            unique_tasks = torch.unique(cu_adapter_mask)
+            out_dtype = next(self.out_proj.parameters()).dtype
+            out = torch.empty(inp.shape[0], self.out_proj.out_features,
+                                   dtype=out_dtype, device=inp.device)
+            for task_id in unique_tasks:
+                task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                task_tensor = inp[task_indices]
+                task_out = self.out_proj(task_tensor, task_id=task_id)
+                out[task_indices] = task_out
+        else:
+            out = self.out_proj(inp)
         return out if not self.return_residual else (out, x)

mlp.py

@@ -47,12 +47,36 @@ class Mlp(nn.Module):
         self.activation = activation
         self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs)
 
-    def forward(self, x, task_type=None):
-        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-        y = self.fc1(x, **lora_kwargs)
+    def forward(self, x, cu_adapter_mask=None):
+        if cu_adapter_mask is not None:
+            unique_tasks = torch.unique(cu_adapter_mask)
+            fc1_dtype = next(self.fc1.parameters()).dtype
+            y = torch.empty(x.shape[0], self.fc1.out_features,
+                              dtype=fc1_dtype, device=x.device)
+            for task_id in unique_tasks:
+                task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                task_tensor = x[task_indices]
+                task_y = self.fc1(task_tensor, task_id=task_id)
+                y[task_indices] = task_y
+        else:
+            y = self.fc1(x)
+
         y = self.activation(y)
-        y = self.fc2(y, **lora_kwargs)
-        return y if not self.return_residual else (y, x)
+
+        if cu_adapter_mask is not None:
+            unique_tasks = torch.unique(cu_adapter_mask)
+            fc2_dtype = next(self.fc2.parameters()).dtype
+            out = torch.empty(y.shape[0], self.fc2.out_features,
+                              dtype=fc2_dtype, device=y.device)
+            for task_id in unique_tasks:
+                task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                task_tensor = y[task_indices]
+                task_out = self.fc2(task_tensor, task_id=task_id)
+                out[task_indices] = task_out
+        else:
+            out = self.fc2(y)
+
+        return out if not self.return_residual else (out, x)
 
 
 class ParallelMLP(nn.Module):

modeling_lora.py

@@ -161,7 +161,6 @@ class LoRAParametrization(nn.Module):
         rank: int,
         dropout_p: float,
         alpha: float,
-        adaptation_map: dict,
     ):
         if isinstance(layer, nn.Linear):
             parametrize.register_parametrization(
@@ -176,10 +175,9 @@ class LoRAParametrization(nn.Module):
                 ),
             )
 
-            def new_forward(self, input, task_type, residual=False):
-                task_idx = adaptation_map[task_type] if task_type else None
-                if task_idx is not None:
-                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
+            def new_forward(self, input, task_id=None, residual=False):
+                if task_id is not None:
+                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_id)
                 else:
                     weights = self.weight
 
@@ -204,10 +202,9 @@ class LoRAParametrization(nn.Module):
                 ),
             )
 
-            def new_forward(self, input, task_type):
-                task_idx = adaptation_map[task_type] if task_type else None
-                if task_idx is not None:
-                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
+            def new_forward(self, input, task_id=None):
+                if task_id is not None:
+                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_id)
                 else:
                     weights = self.weight
 
@@ -227,7 +224,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         roberta: Optional[XLMRobertaModel] = None
     ):
         super().__init__(config)
-
         if roberta is None:
             self.roberta = XLMRobertaModel(config)
         else:
@@ -318,7 +314,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 rank=rank,
                 dropout_p=dropout_p,
                 alpha=alpha,
-                adaptation_map=self._adaptation_map,
             )
         )
 
@@ -341,6 +336,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     @torch.inference_mode()
     def encode(
         self,
+        sentences: Union[str, List[str]],
         *args,
         task_type: Optional[str] = None,
         **kwargs,
@@ -359,5 +355,9 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
                 f"Alternatively, don't pass the `task_type` argument to disable LoRA."
             )
-
-        return self.roberta.encode(*args, task_type=task_type, **kwargs)
+        adapter_mask = None
+        if task_type:
+            task_id = self._adaptation_map[task_type]
+            num_examples = 1 if isinstance(sentences, str) else len(sentences)
+            adapter_mask = torch.full((num_examples,), task_id, dtype=torch.int32, device=self.device)
+        return self.roberta.encode(sentences, *args, adapter_mask=adapter_mask, **kwargs)

modeling_xlm_roberta.py

@@ -204,18 +204,15 @@ 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, task_type=None):
+    def forward(self, hidden_states, key_padding_mask=None, subset_mask=None, adapter_mask=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
         """
         if key_padding_mask is None or not self.use_flash_attn:
-            mixer_kwargs = (
-                {"key_padding_mask": key_padding_mask.bool()}
-                if key_padding_mask is not None
-                else None
-            )
-            mixer_kwargs['task_type'] = task_type
+            mixer_kwargs = {'adapter_mask': adapter_mask}
+            if key_padding_mask is not None:
+                mixer_kwargs['key_padding_mask'] = key_padding_mask.bool()
             for layer in self.layers:
                 if self._grad_checkpointing:
                     hidden_states = torch.utils.checkpoint.checkpoint(
@@ -230,10 +227,11 @@ class XLMRobertaEncoder(nn.Module):
                 hidden_states = hidden_states[subset_mask]
         else:
             batch, seqlen = hidden_states.shape[:2]
-            hidden_states, indices, cu_seqlens, max_seqlen_in_batch = unpad_input(
-                hidden_states, key_padding_mask
+            hidden_states, indices, cu_seqlens, max_seqlen_in_batch, cu_adapter_mask = unpad_input(
+                hidden_states, key_padding_mask, adapter_mask
             )
-            mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch, "task_type": task_type}
+            mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch, "cu_adapter_mask": cu_adapter_mask}
+
             if subset_mask is None:
                 for layer in self.layers:
                     if self._grad_checkpointing:
@@ -310,13 +308,22 @@ class XLMRobertaPooler(nn.Module):
         self.dense = linear_cls(config.hidden_size, config.hidden_size)
         self.activation = nn.Tanh()
 
-    def forward(self, hidden_states, pool=True, task_type=None):
+    def forward(self, hidden_states, pool=True, adapter_mask=None):
         # We "pool" the model by simply taking the hidden state corresponding
         # to the first token.
-        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-
         first_token_tensor = hidden_states[:, 0] if pool else hidden_states
-        pooled_output = self.dense(first_token_tensor, **lora_kwargs)
+        if adapter_mask is not None:
+            unique_tasks = torch.unique(adapter_mask)
+            pool_dtype = next(self.dense.parameters()).dtype
+            pooled_output = torch.empty(first_token_tensor.shape[0], self.dense.out_features,
+                                            dtype=pool_dtype, device=first_token_tensor.device)
+            for task_id in unique_tasks:
+                task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                task_first_token_tensor = first_token_tensor[task_indices]
+                task_pooled_output = self.dense(task_first_token_tensor, task_id=task_id)
+                pooled_output[task_indices] = task_pooled_output
+        else:
+            pooled_output = self.dense(first_token_tensor)
         pooled_output = self.activation(pooled_output)
         return pooled_output
 
@@ -429,7 +436,6 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             "gelu_fast",
             "gelu_pytorch_tanh",
         ]
-
         self.embeddings = XLMRobertaEmbeddings(
             config.hidden_size,
             config.vocab_size,
@@ -457,6 +463,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
         truncate_dim: Optional[int] = None,
+        adapter_mask: Optional[torch.Tensor] = None,
         task_type: Optional[str] = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
@@ -542,14 +549,14 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             )
         else:
             range_iter = range(0, len(sentences), batch_size)
-        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
+        lora_arguments = {'adapter_mask': adapter_mask} if adapter_mask is not None else {}
         for i in range_iter:
             encoded_input = self.tokenizer(
                 sentences[i : i + batch_size],
                 return_tensors='pt',
                 **tokenizer_kwargs,
             ).to(self.device)
-            token_embs = self.forward(**encoded_input, **lora_kwargs)[0]
+            token_embs = self.forward(**encoded_input, **lora_arguments)[0]
 
             # Accumulate in fp32 to avoid overflow
             token_embs = token_embs.float()
@@ -637,7 +644,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         layer output for these tokens.
         masked_tokens_mask: (batch, seqlen), dtype=torch.bool
         """
-        task_type = kwargs.pop('task_type', None)
+        adapter_mask = kwargs.pop('adapter_mask', None)
         if kwargs:
             for key, value in kwargs.items():
                 if value is not None:
@@ -651,7 +658,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         )
 
         hidden_states = self.embeddings(
-            input_ids, position_ids=position_ids, token_type_ids=token_type_ids, task_type=task_type
+            input_ids, position_ids=position_ids, token_type_ids=token_type_ids, adapter_mask=adapter_mask
         )
         # TD [2022-12:18]: Don't need to force residual in fp32
         # BERT puts embedding LayerNorm before embedding dropout.
@@ -675,12 +682,12 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             subset_mask = None
 
         sequence_output = self.encoder(
-            hidden_states, key_padding_mask=attention_mask, subset_mask=subset_mask, task_type=task_type
+            hidden_states, key_padding_mask=attention_mask, subset_mask=subset_mask, adapter_mask=adapter_mask
         )
 
         if masked_tokens_mask is None:
             pooled_output = (
-                self.pooler(sequence_output, task_type=task_type) if self.pooler is not None else None
+                self.pooler(sequence_output, adapter_mask=adapter_mask) if self.pooler is not None else None
             )
         else:
             # TD [2022-03-01]: the indexing here is very tricky.
@@ -694,7 +701,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
                 pool_input = sequence_output[first_col_mask[subset_mask]]
                 sequence_output = sequence_output[masked_tokens_mask[subset_mask]]
             pooled_output = (
-                self.pooler(pool_input, pool=False, task_type=task_type) if self.pooler is not None else None
+                self.pooler(pool_input, pool=False, adapter_mask=adapter_mask) if self.pooler is not None else None
             )
 
         if not return_dict:

xlm_padding.py

@@ -98,7 +98,7 @@ class IndexFirstAxisResidual(torch.autograd.Function):
 index_first_axis_residual = IndexFirstAxisResidual.apply
 
 
-def unpad_input(hidden_states, attention_mask):
+def unpad_input(hidden_states, attention_mask, adapter_mask=None):
     """
     Arguments:
         hidden_states: (batch, seqlen, ...)
@@ -113,6 +113,9 @@ def unpad_input(hidden_states, attention_mask):
     indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
     max_seqlen_in_batch = seqlens_in_batch.max().item()
     cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+
+    cu_adapter_mask = torch.repeat_interleave(adapter_mask, cu_seqlens[1:] - cu_seqlens[:-1]) if adapter_mask is not None else None
+
     # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
     # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
     # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
@@ -123,6 +126,7 @@ def unpad_input(hidden_states, attention_mask):
         indices,
         cu_seqlens,
         max_seqlen_in_batch,
+        cu_adapter_mask,
     )