Fix mistake in parallel layers

decoder_only_t5/modeling.py

@@ -5,7 +5,7 @@ import torch
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from transformers.models.t5 import modeling_t5
-from transformers.modeling_outputs import Seq2SeqLMOutput
+from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.utils import (
     add_start_docstrings_to_model_forward,
     logging,
@@ -167,22 +167,28 @@ class DecoderOnlyT5Attention(modeling_t5.T5Attention):
         )  # (batch_size, n_heads, seq_length, dim_per_head)
 
         # get key/value states
-        key_states = project(
-            hidden_states,
-            self.k,
-            key_value_states,
-            past_key_value[0] if past_key_value is not None else None,
+        key_states = repeat_kv(
+            project(
+                hidden_states,
+                self.k,
+                key_value_states,
+                past_key_value[0] if past_key_value is not None else None,
+            ),
+            self.n_kv_groups,
         )
-        value_states = project(
-            hidden_states,
-            self.v,
-            key_value_states,
-            past_key_value[1] if past_key_value is not None else None,
+        value_states = repeat_kv(
+            project(
+                hidden_states,
+                self.v,
+                key_value_states,
+                past_key_value[1] if past_key_value is not None else None,
+            ),
+            self.n_kv_groups,
         )
 
         # compute scores
         scores = torch.matmul(
-            query_states, repeat_kv(key_states, self.n_kv_groups).transpose(3, 2)
+            query_states, key_states.transpose(3, 2)
         )  # equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
 
         if position_bias is None:
@@ -345,8 +351,9 @@ class DecoderOnlyT5Block(modeling_t5.T5Block):
 
         ff_layer = self.layer[-1]
         if self.parallel_layers:
+            # https://github.com/google/flaxformer/blob/ea17eb012a1d340ddff017b7a534c2162aaec34c/flaxformer/architectures/t5/t5_architecture.py#L563-L568
             x = self.layer[0].layer_norm(hidden_states)
-            ff_output = ff_layer(hidden_states)
+            ff_output = ff_layer(x)
         else:
             x = hidden_states
 
@@ -418,7 +425,7 @@ class DecoderOnlyT5Block(modeling_t5.T5Block):
             attention_outputs = attention_outputs + cross_attention_outputs[2:]
 
         if self.parallel_layers:
-            # https://github.com/google/flaxformer/blob/ea17eb012a1d340ddff017b7a534c2162aaec34c/flaxformer/architectures/t5/t5_architecture.py#L295
+            # https://github.com/google/flaxformer/blob/ea17eb012a1d340ddff017b7a534c2162aaec34c/flaxformer/architectures/t5/t5_architecture.py#L534-L578
             hidden_states = x + ff_output
             hidden_states *= 2**-0.5
             hidden_states = hidden_states + self.layer[0].dropout(hidden_states)
@@ -508,27 +515,21 @@ class DecoderOnlyT5Model(modeling_t5.T5ForConditionalGeneration):
 
     @add_start_docstrings_to_model_forward(modeling_t5.T5_INPUTS_DOCSTRING)
     @replace_return_docstrings(
-        output_type=Seq2SeqLMOutput, config_class=_CONFIG_FOR_DOC
+        output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
     )
     def forward(
         self,
-        _input_ids: Optional[torch.LongTensor] = None,
+        input_ids: Optional[torch.LongTensor] = None,
         attention_mask: Optional[torch.FloatTensor] = None,
-        decoder_input_ids: Optional[torch.LongTensor] = None,
-        decoder_attention_mask: Optional[torch.BoolTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        decoder_head_mask: Optional[torch.FloatTensor] = None,
-        cross_attn_head_mask: Optional[torch.Tensor] = None,
-        encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        position_ids: Optional[torch.LongTensor] = None,
         past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
-        _inputs_embeds: Optional[torch.FloatTensor] = None,
-        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-    ) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
             Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ...,
@@ -548,43 +549,31 @@ class DecoderOnlyT5Model(modeling_t5.T5ForConditionalGeneration):
         if self.model_parallel:
             torch.cuda.set_device(self.decoder.first_device)
 
-        if (
-            labels is not None
-            and decoder_input_ids is None
-            and decoder_inputs_embeds is None
-        ):
-            # get decoder inputs from shifting lm labels to the right
-            decoder_input_ids = self._shift_right(labels)
-
         # Set device for model parallelism
         if self.model_parallel:
             torch.cuda.set_device(self.decoder.first_device)
-            if decoder_input_ids is not None:
-                decoder_input_ids = decoder_input_ids.to(self.decoder.first_device)
+            if input_ids is not None:
+                input_ids = input_ids.to(self.decoder.first_device)
             if attention_mask is not None:
                 attention_mask = attention_mask.to(self.decoder.first_device)
-            if decoder_attention_mask is not None:
-                decoder_attention_mask = decoder_attention_mask.to(
-                    self.decoder.first_device
-                )
 
         # Decode
-        decoder_outputs = self.decoder(
-            input_ids=decoder_input_ids,
-            attention_mask=decoder_attention_mask,
-            inputs_embeds=decoder_inputs_embeds,
+        outputs = self.decoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
             past_key_values=past_key_values,
-            # encoder_hidden_states=hidden_states,
-            encoder_attention_mask=attention_mask,
-            head_mask=decoder_head_mask,
-            cross_attn_head_mask=cross_attn_head_mask,
+            encoder_hidden_states=None,
+            encoder_attention_mask=None,
+            head_mask=None,
+            cross_attn_head_mask=None,
             use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
 
-        sequence_output = decoder_outputs[0]
+        sequence_output = outputs[0]
 
         # Set device for model parallelism
         if self.model_parallel:
@@ -608,13 +597,13 @@ class DecoderOnlyT5Model(modeling_t5.T5ForConditionalGeneration):
             # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666
 
         if not return_dict:
-            output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs
+            output = (lm_logits,) + outputs[1:]
             return ((loss,) + output) if loss is not None else output
 
-        return Seq2SeqLMOutput(
+        return CausalLMOutputWithPast(
             loss=loss,
             logits=lm_logits,
-            past_key_values=decoder_outputs.past_key_values,
-            decoder_hidden_states=decoder_outputs.hidden_states,
-            decoder_attentions=decoder_outputs.attentions,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
         )