fix: adapter masks

Signed-off-by: Meow <[email protected]>

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, cu_adapter_mask=mixer_kwargs.get('cu_adapter_mask'))
+                mlp_out = self.mlp(hidden_states, adapter_mask=mixer_kwargs.get('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:

mha.py

@@ -590,7 +590,7 @@ class MHA(nn.Module):
         max_seqlen=None,
         mixer_subset=None,
         inference_params=None,
-        cu_adapter_mask=None,
+        adapter_mask=None,
         **kwargs,
     ):
         """
@@ -647,13 +647,13 @@ 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
 
-            if cu_adapter_mask is not None:
-                unique_tasks = torch.unique(cu_adapter_mask)
+            if adapter_mask is not None:
+                unique_tasks = torch.unique(adapter_mask)
                 qkv_dtype = next(self.Wqkv.parameters()).dtype
-                qkv = torch.empty(x.shape[0], self.Wqkv.out_features,
+                qkv = torch.empty(*x.shape[:-1], 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_indices = (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)
@@ -755,13 +755,13 @@ class MHA(nn.Module):
                 context = self._apply_rotary_update_kvcache_attention(q, kv, inference_params)
 
         inp = rearrange(context, "... h d -> ... (h d)")
-        if cu_adapter_mask is not None:
-            unique_tasks = torch.unique(cu_adapter_mask)
+        if adapter_mask is not None:
+            unique_tasks = torch.unique(adapter_mask)
             out_dtype = next(self.out_proj.parameters()).dtype
-            out = torch.empty(inp.shape[0], self.out_proj.out_features,
+            out = torch.empty(*inp.shape[:-1], 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_indices = (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

mlp.py

@@ -47,14 +47,14 @@ class Mlp(nn.Module):
         self.activation = activation
         self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs)
 
-    def forward(self, x, cu_adapter_mask=None):
-        if cu_adapter_mask is not None:
-            unique_tasks = torch.unique(cu_adapter_mask)
+    def forward(self, x, adapter_mask=None):
+        if adapter_mask is not None:
+            unique_tasks = torch.unique(adapter_mask)
             fc1_dtype = next(self.fc1.parameters()).dtype
-            y = torch.empty(x.shape[0], self.fc1.out_features,
+            y = torch.empty(*x.shape[:-1], 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_indices = (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
@@ -63,13 +63,13 @@ class Mlp(nn.Module):
 
         y = self.activation(y)
 
-        if cu_adapter_mask is not None:
-            unique_tasks = torch.unique(cu_adapter_mask)
+        if adapter_mask is not None:
+            unique_tasks = torch.unique(adapter_mask)
             fc2_dtype = next(self.fc2.parameters()).dtype
-            out = torch.empty(y.shape[0], self.fc2.out_features,
+            out = torch.empty(*y.shape[:-1], 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_indices = (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

modeling_xlm_roberta.py

@@ -230,7 +230,7 @@ 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, "cu_adapter_mask": cu_adapter_mask}
+            mixer_kwargs = {"cu_seqlens": cu_seqlens, "max_seqlen": max_seqlen_in_batch, "adapter_mask": cu_adapter_mask}
 
             if subset_mask is None:
                 for layer in self.layers: