Add labels into forward

mosaic_gpt.py

@@ -238,6 +238,7 @@ class MosaicGPT(PreTrainedModel):
             input_ids: torch.LongTensor,
             past_key_values: Optional[List[Tuple[torch.FloatTensor]]] = None,
             attention_mask: Optional[torch.ByteTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
             prefix_mask: Optional[torch.ByteTensor] = None,
             sequence_id: Optional[torch.LongTensor] = None,
             return_dict: Optional[bool] = None,
@@ -370,9 +371,27 @@ class MosaicGPT(PreTrainedModel):
                 )
             logits *= self.logit_scale
 
-        return CausalLMOutputWithPast(logits=logits,
-                                      past_key_values=past_key_values,
-                                      hidden_states=all_hidden_states)
+        # compute loss from logits
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(
+                    -1, self.transformer.wte.num_embeddings
+                ),
+                shift_labels.view(-1),
+            )
+            return CausalLMOutputWithPast(loss=loss, logits=logits,
+                                          past_key_values=past_key_values,
+                                          hidden_states=all_hidden_states)
+
+        else:
+            return CausalLMOutputWithPast(logits=logits,
+                                          past_key_values=past_key_values,
+                                          hidden_states=all_hidden_states)
 
     # Param Initialization, needed for device='meta' fast initialization
     def param_init_fn(self, module):