refactor: Merge MoECrossEntropyLoss into modalities core

Author: rrutmann

src/modalities/loss_functions.py

@@ -87,6 +87,39 @@ def _parse_arguments(
         return labels, lm_logits
 
 
+class MoECrossEntropyLoss(Loss):
+    """Cross entropy loss with optional MoE auxiliary losses from model layers."""
+
+    def __init__(
+        self,
+        target_key: str,
+        prediction_key: str,
+        model,
+        tag: str = "MoECrossEntropyLoss",
+    ):
+        super().__init__(tag)
+        self.target_key = target_key
+        self.prediction_key = prediction_key
+        self.model = model
+        self.loss_fun = CrossEntropyLoss(reduction="mean")
+
+    def __call__(self, forward_batch: InferenceResultBatch) -> torch.Tensor:
+        labels = forward_batch.get_targets(self.target_key)
+        lm_logits = forward_batch.get_predictions(self.prediction_key)
+
+        labels = labels.to(lm_logits.device)
+        loss = self.loss_fun(
+            lm_logits.contiguous().view(-1, lm_logits.size(-1)),
+            labels.contiguous().long().view(-1),
+        )
+
+        for layer in self.model.layers.values():
+            if hasattr(layer, "aux_loss") and layer.aux_loss is not None:
+                loss = loss + layer.aux_loss.to(loss.dtype)
+
+        return loss
+
+
 def nce_loss(
     embedding1: torch.Tensor, embedding2: torch.Tensor, device: torch.device, is_asymmetric: bool, temperature: float
 ) -> torch.Tensor:

src/modalities/models/moe/__init__.py

@@ -1,4 +1,4 @@
-from modalities.models.moe.loss_functions import MoECrossEntropyLoss
+from modalities.loss_functions import MoECrossEntropyLoss
 from modalities.models.moe.qwen_model import QwenModel, QwenModelConfig
 
 __all__ = [

src/modalities/models/moe/loss_functions.py

@@ -85,7 +85,7 @@
     ProgressSubscriberFactory,
     ResultsSubscriberFactory,
 )
-from modalities.loss_functions import CLMCrossEntropyLoss
+from modalities.loss_functions import CLMCrossEntropyLoss, MoECrossEntropyLoss
 from modalities.models.coca.coca_model import CoCa, CoCaConfig
 from modalities.models.coca.collator import CoCaCollateFnConfig, CoCaCollatorFn
 from modalities.models.components.layer_norms import (
@@ -99,7 +99,6 @@
 from modalities.models.gpt2.llama3_like_initialization import Llama3Initializer, Llama3InitializerConfig
 from modalities.models.huggingface.huggingface_model import HuggingFacePretrainedModel, HuggingFacePretrainedModelConfig
 from modalities.models.model_factory import GPT2ModelFactory, ModelFactory
-from modalities.models.moe.loss_functions import MoECrossEntropyLoss
 from modalities.models.moe.qwen_model import QwenModel, QwenModelConfig
 from modalities.models.parallelism.pipeline_parallelism import ComponentSelectorFromPipeline, PipelineFactory
 from modalities.models.parallelism.pipeline_parallelism_configs import (

src/modalities/registry/components.py

@@ -2,7 +2,7 @@
 from torch.nn import CrossEntropyLoss
 
 from modalities.batch import InferenceResultBatch
-from modalities.models.moe.loss_functions import MoECrossEntropyLoss
+from modalities.loss_functions import MoECrossEntropyLoss
 
 
 class DummyLayer: