RFC: refactor causal lm loss to handle lm_head in loss function

src/transformers/loss/loss_utils.py

@@ -16,6 +16,7 @@
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from torch.nn import BCEWithLogitsLoss, MSELoss
 
 from .loss_d_fine import DFineForObjectDetectionLoss
@@ -46,11 +47,22 @@ def ForCausalLMLoss(
     logits,
     labels,
     vocab_size: int,
+    hidden_states: Optional[torch.Tensor] = None,
+    lm_head_weight: Optional[torch.Tensor] = None,
+    logits_to_keep: Optional[int] = None,
     num_items_in_batch: Optional[torch.Tensor] = None,
     ignore_index: int = -100,
     shift_labels: Optional[torch.Tensor] = None,
     **kwargs,
 ) -> torch.Tensor:
+    if hidden_states and lm_head_weight:
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        # compute linear
+        logits = F.linear(
+            hidden_states[:, slice_indices, :],
+            lm_head_weight,
+        )
+
     # Upcast to float if we need to compute the loss to avoid potential precision issues
     logits = logits.float()
 

src/transformers/models/llama/modeling_llama.py

@@ -562,13 +562,24 @@ def forward(
         )
 
         hidden_states = outputs.last_hidden_state
-        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
 
         loss = None
+        logits = None
         if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+            loss = self.loss_function(
+                logits=None,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                hidden_states=hidden_states,
+                lm_head_weight=self.lm_head.weight,
+                logits_to_keep=logits_to_keep,
+                **kwargs,
+            )
+        else:
+            # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+
 
         return CausalLMOutputWithPast(
             loss=loss,