feat: adapter masking finished

Signed-off-by: Meow <ongjackm@gmail.com>

block.py

@@ -233,17 +233,17 @@ 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 cu_adapter_mask:
+                #     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)
                 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,40 +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, adapter_mask=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).tolist()
-            torch_dtype = next(self.word_embeddings.parameters()).dtype
-            embeddings = torch.empty(*input_ids.shape, self.word_embeddings.embedding_dim, dtype=torch_dtype).to(input_ids.device)
-            for task in unique_tasks:
-                indices = (adapter_mask == task).nonzero(as_tuple=True)[0]
-                inp = input_ids[indices]
-                lora_kwargs = {'task_type': task} if task is not None else {}
-                emb = self.word_embeddings(inp, **lora_kwargs)
-                embeddings[indices] = emb
-
-            exit(0)
+            embedding_dtype = next(self.word_embeddings.parameters()).dtype
+            embeddings = torch.empty(*input_ids.shape, self.word_embeddings.embedding_dim,
+                                            dtype=embedding_dtype).to(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_type=task_id)
+                embeddings[task_indices] = task_embeddings
         else:
-            unique_task = torch.unique(adapter_mask)[0]
-            task1_indices = (adapter_mask == unique_task).nonzero(as_tuple=True)[0]
-            input1 = input_ids[task1_indices]
-            lora_kwargs = {'task_type': unique_task} if unique_task is not None else {}
-            embeddings = self.word_embeddings(input1, **lora_kwargs)
-
+            embeddings = self.word_embeddings(input_ids)
 
         if self.max_position_embeddings > 0:
             if position_ids is None:
@@ -84,19 +69,15 @@ 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).tolist()
+                for task_id in unique_tasks:
+                    task_token_type_embeddings = self.token_type_embeddings(token_type_ids, task_type=task_id)
+                    task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0]
+                    embeddings[task_indices] = embeddings[task_indices] + task_token_type_embeddings
             else:
-                unique_task = torch.unique(adapter_mask)[0]
-                lora_kwargs = {'task_type': unique_task} if unique_task 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,
     ):
         """
@@ -643,39 +643,27 @@ 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
 
-            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).tolist()
+                qkv_dtype = next(self.Wqkv.parameters()).dtype
+                qkv = torch.empty(x.shape[0], self.Wqkv.out_features,
+                                         dtype=qkv_dtype).to(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_type=task_id)
+                    else:
+                        task_qkv, _ = self.Wqkv(task_tensor, task_type=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)
+                    qkv, x = self.Wqkv(x)
 
             if self.dwconv:
                 qkv = rearrange(
@@ -762,14 +750,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).tolist()
+            out_dtype = next(self.out_proj.parameters()).dtype
+            out = torch.empty(inp.shape[0], self.out_proj.out_features,
+                                   dtype=out_dtype).to(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_type=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).tolist()
+            fc1_dtype = next(self.fc1.parameters()).dtype
+            y = torch.empty(x.shape[0], self.fc1.out_features,
+                              dtype=fc1_dtype).to(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_type=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).tolist()
+            fc2_dtype = next(self.fc2.parameters()).dtype
+            out = torch.empty(y.shape[0], self.fc2.out_features,
+                              dtype=fc2_dtype).to(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_type=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.fc1(y)
+
+        return out if not self.return_residual else (out, x)
 
 
 class ParallelMLP(nn.Module):

modeling_xlm_roberta.py

@@ -204,18 +204,16 @@ 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, adapter_mask=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(
@@ -233,7 +231,8 @@ class XLMRobertaEncoder(nn.Module):
             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, "cu_adapter_mask": cu_adapter_mask}
+            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,24 +309,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).tolist()
+            pool_dtype = next(self.dense.parameters()).dtype
+            pooled_output = torch.empty(first_token_tensor.shape[0], self.dense.out_features,
+                                            dtype=pool_dtype).to(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_type=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
 
@@ -440,7 +437,6 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             "gelu_fast",
             "gelu_pytorch_tanh",
         ]
-
         self.embeddings = XLMRobertaEmbeddings(
             config.hidden_size,
             config.vocab_size,
@@ -648,7 +644,6 @@ 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():
@@ -663,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, adapter_mask=adapter_mask
+            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.
@@ -687,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, adapter_mask=adapter_mask
+            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.
@@ -706,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: