Upload model

modeling_t5mimo.py

@@ -277,7 +277,7 @@ class T5Attention(nn.Module):
         relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
         return relative_buckets
 
-    def compute_bias(self, query_length, key_length, device=None):
+    def compute_bias(self, query_length, key_length, device=None, multivar_dim=None):
         """Compute binned relative position bias"""
         if device is None:
             device = self.relative_attention_bias.weight.device
@@ -293,7 +293,10 @@ class T5Attention(nn.Module):
         values = self.relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)
         values = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)
         if self.config.is_mimo: 
+            if multivar_dim == None:
+                raise ValueError(f"multivar_dim can not be None when config.is_mimo=True")
             values = values.unsqueeze(0)# shape (1, 1, num_heads, query_length, key_length) 
+            values = values.repeat(1, multivar_dim, 1, 1, 1)  # shape (1, multivar_dim, num_heads, query_length, key_length)
         return values
 
     def forward(
@@ -319,6 +322,7 @@ class T5Attention(nn.Module):
             batch_size, multivar_dim, seq_length = hidden_states.shape[:3]
         else:
             batch_size, seq_length = hidden_states.shape[:2]
+            multivar_dim=None
         real_seq_length = seq_length
         
         if past_key_value is not None:
@@ -406,13 +410,13 @@ class T5Attention(nn.Module):
         if position_bias is None:
             if not self.has_relative_attention_bias:
                 if self.config.is_mimo:
-                    position_bias = torch.zeros((1,1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype)
+                    position_bias = torch.zeros((1,multivar_dim, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype)
                 else:
                     position_bias = torch.zeros((1, self.n_heads, real_seq_length, key_length), device=scores.device, dtype=scores.dtype)
                 if self.gradient_checkpointing and self.training:
                     position_bias.requires_grad = True
             else:
-                position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device)
+                position_bias = self.compute_bias(real_seq_length, key_length, device=scores.device, multivar_dim=multivar_dim)
 
 
             # if key and values are already calculated
@@ -924,8 +928,9 @@ class T5Stack(T5PreTrainedModel):
         # ourselves in which case we just need to make it broadcastable to all heads.
 
         if self.config.is_mimo:
-            extended_attention_mask = self.get_extended_attention_mask(attention_mask, (input_shape[0], input_shape[2]))
-            extended_attention_mask = extended_attention_mask.transpose(1,2).unsqueeze(2)
+            extended_attention_mask = self.get_extended_attention_mask(attention_mask[:,0,:], (input_shape[0], input_shape[2]))
+            extended_attention_mask = extended_attention_mask.unsqueeze(0)  
+            extended_attention_mask = extended_attention_mask.repeat(1, input_shape[1], 1, 1, 1) 
         else:
             extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape)
 
@@ -943,11 +948,9 @@ class T5Stack(T5PreTrainedModel):
                 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device, dtype=torch.long)
 
             if self.config.is_mimo:
-
-                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask).transpose(1,2)
-
-                encoder_extended_attention_mask = encoder_extended_attention_mask.unsqueeze(2)
-
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+                encoder_extended_attention_mask = encoder_extended_attention_mask.unsqueeze(0)  
+                encoder_extended_attention_mask = encoder_extended_attention_mask.repeat(1, input_shape[1], 1, 1, 1) 
             else:
                 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
      
@@ -1488,13 +1491,6 @@ class T5MIMOForConditionalGeneration(T5PreTrainedModel):
             if decoder_attention_mask is not None:
                 decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device)
 
-        if hidden_states is not None and decoder_input_ids is not None:
-            if len(hidden_states.shape) == 4:
-                batch_size, multivar_seqs, seq_length , model_dim = hidden_states.shape
-                if len(decoder_input_ids.shape) == 2:
-                    decoder_input_ids = decoder_input_ids.unsqueeze(1).repeat(1, multivar_seqs, 1)
-
-
 
 
         # Decode