Refactor lm_head losses

fast_llm/layers/language_model/head.py

@@ -370,11 +370,13 @@ def _logits_cross_entropy_forward_backward(
                 logits_scale_factor=self._config.logits_scale_factor,
                 target_format=TargetFormat.labels,
             )
+            if self.training and losses is not None:
+                losses[self._ce_loss_name_unscaled].append(lm_loss.detach())
             lm_loss = lm_loss * self._config.language_model_loss_factor
         else:
             lm_loss, lm_grad = None, None
 
-        if distillation_target is not None and self._config.distillation_loss_factor > 0.0:
+        if distillation_target is not None:
             if self._config.distillation_loss_implementation == DistillationLossImpl.reverse_kl:
                 distillation_loss, distillation_grad = reverse_kl_forward_backward(
                     logits.flatten(0, -2),
@@ -405,18 +407,19 @@ def _logits_cross_entropy_forward_backward(
                 raise ValueError(
                     f"Invalid distillation loss implementation: {self._config.distillation_loss_implementation}"
                 )
+            if self.training and losses is not None:  # we keep track of unscaled losses for model comparison purposes
+                losses[self._distillation_loss_name_unscaled].append(distillation_loss.detach())
             distillation_loss = distillation_loss * self._config.distillation_loss_factor
-        else:
-            distillation_loss, distillation_grad = None, None
-
-        # TODO: de-allocate earlier.
-        del logits
 
         # TODO: Accumulate grads in-place to reduce memory and compute overhead.
         grad = _add_tensors(dpo_grad, lm_grad, distillation_grad)
 
         # TODO: Return individual losses?
         loss = _add_tensors(dpo_loss, lm_loss, distillation_loss)
+
+        # TODO: de-allocate earlier.
+        del logits
+
         if self.training and losses is not None:
             if dpo_loss is not None:
                 losses[self._dpo_loss_name].append(dpo_loss.detach())
@@ -434,6 +437,13 @@ def _loss_name(self) -> str:
             name = f"{name}_{self._prediction_distance}"
         return name
 
+    @functools.cached_property
+    def _ce_loss_name_unscaled(self) -> str:
+        name = "language_model_loss_unscaled"
+        if self._prediction_distance > 0:
+            name = f"{name}_{self._prediction_distance}"
+        return name
+
     @functools.cached_property
     def _z_loss_name(self) -> str:
         name = "z_loss"
@@ -462,8 +472,24 @@ def _distillation_loss_name(self) -> str:
             name = f"{name}_{self._prediction_distance}"
         return name
 
+    @functools.cached_property
+    def _distillation_loss_name_unscaled(self) -> str:
+        name = "distillation_loss_unscaled"
+        if self._prediction_distance > 0:
+            name = f"{name}_{self._prediction_distance}"
+        return name
+
     def get_loss_definitions(self, count: int = 1) -> list[LossDef]:
         loss_defs = [LossDef(name=self._loss_name, formatted_name=_format_name(self._loss_name), count=count)]
+        if self._config.distillation_model is None or self._config.language_model_loss_factor > 0.0:
+            # unscaled CE loss (NTP)
+            loss_defs = [
+                LossDef(
+                    name=self._ce_loss_name_unscaled,
+                    formatted_name=_format_name(self._ce_loss_name_unscaled),
+                    count=count,
+                )
+            ]
         if self._config.logit_z_loss:
             loss_defs.append(
                 LossDef(name=self._z_loss_name, formatted_name=_format_name(self._z_loss_name), count=count)
@@ -481,6 +507,15 @@ def get_loss_definitions(self, count: int = 1) -> list[LossDef]:
                     count=count,
                 )
             )
+            # unscaled distillation loss for comparison purposes
+            loss_defs.append(
+                LossDef(
+                    name=self._distillation_loss_name_unscaled,
+                    formatted_name=_format_name(self._distillation_loss_name_unscaled),
+                    count=count,
+                )
+            )
+            # if we mix distillation loss and CE loss for NTP, we want to log both
             if self._config.language_model_loss_factor > 0.0:
                 loss_defs.append(
                     LossDef(