Merge branch 'main' into pr/31

block.py

@@ -233,17 +233,7 @@ class Block(nn.Module):
                     is_rms_norm=isinstance(self.norm1, RMSNorm),
                 )
             if not isinstance(self.mlp, nn.Identity):
-                task_type = mixer_kwargs.get('task_type')
-                if task_type:
-                    if isinstance(task_type, tuple):
-                        assert mixer_kwargs['cu_seqlens'].shape[0] % 9 == 1
-                        split_index = int((mixer_kwargs['cu_seqlens'].shape[0] - 1) / 9)
-                        split = mixer_kwargs['cu_seqlens'][split_index]
-                        mlp_out = self.mlp(hidden_states, task_type=mixer_kwargs.get('task_type'), split=split)
-                    else:
-                        mlp_out = self.mlp(hidden_states, task_type=task_type)
-                else:
-                    mlp_out = self.mlp(hidden_states)
+                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,25 +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
-        if isinstance(task_type, tuple):
-            assert input_ids.shape[0] % 9 == 0
-            split = int(input_ids.shape[0] / 9)
-            tensor1 = input_ids[:split, :]
-            tensor2 = input_ids[split:, :]
-            emb1 = self.word_embeddings(tensor1, task_type=task_type[0])
-            emb2 = self.word_embeddings(tensor2, task_type=task_type[1])
-            embeddings = torch.cat((emb1, emb2), dim=0)
+        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:
-            lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-            embeddings = self.word_embeddings(input_ids, **lora_kwargs)
-
+            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)
@@ -68,18 +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)
-            if isinstance(task_type, tuple):
-                assert embeddings.shape[0] % 9 == 0
-                split = int(embeddings.shape[0] / 9)
-                emb1 = embeddings[:split, :, :]
-                emb2 = embeddings[split:, :, :]
-                token_type_embs1 = self.token_type_embeddings(token_type_ids, task_type=task_type[0])
-                token_type_embs2 = self.token_type_embeddings(token_type_ids, task_type=task_type[1])
-                emb1 = emb1 + token_type_embs1
-                emb2 = emb2 + token_type_embs2
-                embeddings = torch.cat((emb1, emb2), dim=0)
+
+            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:
-                lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-                token_type_embeddings = self.token_type_embeddings(token_type_ids, **lora_kwargs)
+                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,
     ):
         """
@@ -647,35 +647,27 @@ class MHA(nn.Module):
         if not self.cross_attn and self.num_heads_kv == self.num_heads:
             assert x_kv is None and mixer_subset is None
 
-            split = None
-            if isinstance(task_type, tuple):
-                assert cu_seqlens.shape[0] % 9 == 1
-                split_index = int((cu_seqlens.shape[0] - 1) / 9)
-                split = cu_seqlens[split_index]
-
-            lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-
-            if not self.return_residual:
-                if isinstance(task_type, tuple):
-                    tensor1 = x[:split, :]
-                    tensor2 = x[split:, :]
-                    qkv1 = self.Wqkv(tensor1, task_type=task_type[0])
-                    qkv2 = self.Wqkv(tensor2, task_type=task_type[1])
-                    qkv = torch.cat((qkv1, qkv2), dim=0)
-                else:
-                    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 isinstance(task_type, tuple):
-                    tensor1 = x[:split, :]
-                    tensor2 = x[split:, :]
-                    qkv1, tensor1 = self.Wqkv(tensor1, task_type=task_type[0], residual=True)
-                    qkv2, tensor2 = self.Wqkv(tensor2, task_type=task_type[1], residual=True)
-                    qkv = torch.cat((qkv1, qkv2), dim=0)
-                    x = torch.cat((tensor1, tensor2), dim=0)
+                if not self.return_residual:
+                    qkv = self.Wqkv(x)
                 else:
-                    if lora_kwargs:
-                        lora_kwargs['residual'] = True
-                    qkv, x = self.Wqkv(x, **lora_kwargs)
+                    if hasattr(self.Wqkv, 'parametrizations'):
+                        qkv, x = self.Wqkv(x, residual=True)
+                    else:
+                        qkv, x = self.Wqkv(x)
 
             if self.dwconv:
                 qkv = rearrange(
@@ -762,14 +754,17 @@ class MHA(nn.Module):
             else:
                 context = self._apply_rotary_update_kvcache_attention(q, kv, inference_params)
 
-        lora_kwargs.pop('residual', None)
         inp = rearrange(context, "... h d -> ... (h d)")
-        if isinstance(task_type, tuple):
-            tensor1 = inp[:split, :]
-            tensor2 = inp[split:, :]
-            out1 = self.out_proj(tensor1, task_type=task_type[0])
-            out2 = self.out_proj(tensor2, task_type=task_type[1])
-            out = torch.cat((out1, out2), dim=0)
+        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, **lora_kwargs)
+            out = self.out_proj(inp)
         return out if not self.return_residual else (out, x)

mlp.py

@@ -47,30 +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, split=None):
-        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
-        if split:
-            assert isinstance(task_type, tuple)
-            tensor1 = x[:split, :]
-            tensor2 = x[split:, :]
-            y1 = self.fc1(tensor1, task_type=task_type[0])
-            y2 = self.fc1(tensor2, task_type=task_type[1])
-            y = torch.cat((y1, y2), dim=0)
+    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, **lora_kwargs)
+            y = self.fc1(x)
 
         y = self.activation(y)
 
-        if split:
-            assert isinstance(task_type, tuple)
-            tensor1 = y[:split, :]
-            tensor2 = y[split:, :]
-            y1 = self.fc2(tensor1, task_type=task_type[0])
-            y2 = self.fc2(tensor2, task_type=task_type[1])
-            y = torch.cat((y1, y2), dim=0)
+        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:
-            y = self.fc2(y, **lora_kwargs)
-        return y if not self.return_residual else (y, x)
+            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
 
@@ -325,7 +322,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 rank=rank,
                 dropout_p=dropout_p,
                 alpha=alpha,
-                adaptation_map=self._adaptation_map,
             )
         )
 
@@ -348,6 +344,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     @torch.inference_mode()
     def encode(
         self,
+        sentences: Union[str, List[str]],
         *args,
         task_type: Optional[str] = None,
         **kwargs,
@@ -366,5 +363,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,16 +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 = {'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(
@@ -228,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:
@@ -308,24 +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
-
-        if isinstance(task_type, tuple):
-            assert first_token_tensor.shape[0] % 9 == 0
-            split = int(first_token_tensor.shape[0] / 9)
-            tensor1 = first_token_tensor[:split, :]
-            tensor2 = first_token_tensor[split:, :]
-            pooled_out1 = self.dense(tensor1, task_type=task_type[0])
-            pooled_out2 = self.dense(tensor2, task_type=task_type[0])
-            pooled_output = torch.cat((pooled_out1, pooled_out2), dim=0)
+        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, **lora_kwargs)
-
+            pooled_output = self.dense(first_token_tensor)
         pooled_output = self.activation(pooled_output)
         return pooled_output
 
@@ -438,7 +436,6 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             "gelu_fast",
             "gelu_pytorch_tanh",
         ]
-
         self.embeddings = XLMRobertaEmbeddings(
             config.hidden_size,
             config.vocab_size,
@@ -467,6 +464,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]:
@@ -552,14 +550,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()
@@ -657,7 +655,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:
@@ -671,7 +669,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.
@@ -695,12 +693,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.
@@ -714,7 +712,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,
     )