update return dict.

modeling_mpt.py

@@ -134,8 +134,8 @@ class MPTModel(MPTPreTrainedModel):
             attention_mask = attention_mask.bool()
         if prefix_mask is not None:
             prefix_mask = prefix_mask.bool()
-        if not return_dict:
-            raise NotImplementedError('return_dict False is not implemented yet for MPT')
+        # if not return_dict:
+        #     raise NotImplementedError('return_dict False is not implemented yet for MPT')
         if output_attentions:
             raise NotImplementedError('output_attentions is not implemented yet for MPT')
         if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0] and self.training:
@@ -184,6 +184,9 @@ class MPTModel(MPTPreTrainedModel):
             if past_key_values is not None:
                 past_key_values[b_idx] = past_key_value
         x = self.norm_f(x)
+        if not return_dict:
+            output = (x,) + (tuple(past_key_values),)
+            return output
         return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states)
 
     def param_init_fn(self, module):
@@ -234,6 +237,9 @@ class MPTForCausalLM(MPTPreTrainedModel):
     def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None):
         return_dict = return_dict if return_dict is not None else self.config.return_dict
         use_cache = use_cache if use_cache is not None else self.config.use_cache
+        
+        past_key_values = list(past_key_values) if past_key_values is not None else None
+        
         outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
         logits = F.linear(outputs.last_hidden_state, self.transformer.wte.weight)
         if self.logit_scale is not None:
@@ -245,6 +251,10 @@ class MPTForCausalLM(MPTPreTrainedModel):
             labels = torch.roll(labels, shifts=-1)
             labels[:, -1] = -100
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1))
+
+        if not return_dict:
+            output = (logits,) + (tuple(outputs[1]),)
+            return (loss,) + output if loss is not None else output
         return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states)
 
     def param_init_fn(self, module):