MoE Refactoring

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/model_factory.py

@@ -5,6 +5,7 @@
 import itertools
 import json
 import time
+import warnings
 from dataclasses import asdict, dataclass
 from functools import partial
 from pathlib import Path
@@ -18,6 +19,7 @@
 from torch.distributed.fsdp import FSDPModule as FSDP2
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP1
 from torch.distributed.fsdp import ShardingStrategy
+from torch.distributed.tensor import DTensor
 from torch.distributed.tensor.parallel import (
     ColwiseParallel,
     PrepareModuleInput,
@@ -41,9 +43,10 @@
     TransformerMLP,
 )
 from modalities.models.model import ActivationType
+from modalities.models.parallelism.expert_parallelism import ExpertParallel
 from modalities.nn.model_initialization.initialization_if import ModelInitializationIF
 from modalities.running_env.env_utils import FSDP2MixedPrecisionSettings, MixedPrecisionSettings
-from modalities.running_env.fsdp.device_mesh import ParallelismDegrees
+from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_mesh_for_parallelism_method
 from modalities.running_env.fsdp.fsdp_auto_wrapper import FSDPTransformerAutoWrapPolicyFactory
 from modalities.training.activation_checkpointing.activation_checkpointing import (
     ActivationCheckpointing,
@@ -253,6 +256,95 @@ def get_fsdp2_wrapped_model(
         )
         return model
 
+    @staticmethod
+    def get_ep_wrapped_model(
+        model,
+        block_names: list[str],
+        device_mesh: DeviceMesh,
+        mixed_precision_settings: FSDP2MixedPrecisionSettings,
+    ) -> nn.Module:
+        """Apply expert parallel wrapping for selected MoE blocks."""
+        block_types = []
+        missing_block_names = []
+        for name in block_names:
+            block_type = get_module_class_from_name(model, name)
+            if block_type is None:
+                missing_block_names.append(name)
+            else:
+                block_types.append(block_type)
+
+        if len(missing_block_names) > 0 and (
+            not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
+        ):
+            warnings.warn(
+                f"Could not resolve some requested MoE block names and they will be ignored: {missing_block_names}",
+                stacklevel=2,
+            )
+
+        block_types = tuple(block_types)
+
+        if len(block_types) == 0:
+            raise ValueError(f"None of the requested MoE block names were found: {block_names}")
+
+        ep_mesh = get_mesh_for_parallelism_method(device_mesh, ParallelismDegrees.EP)
+        target_dtype = mixed_precision_settings.param_dtype.value
+
+        wrapped_blocks = 0
+        for module in model.modules():
+            if isinstance(module, block_types):
+                if hasattr(module, "experts"):
+                    ep_target = module
+                elif (ffn := getattr(module, "ffn", None)) is not None and hasattr(ffn, "experts"):
+                    ep_target = ffn
+                else:
+                    raise ValueError(
+                        f"Module {type(module).__name__} has no EP-compatible experts location. "
+                        "Expected `experts` or `ffn.experts`."
+                    )
+
+                if getattr(ep_target, "_ep_enabled", False):
+                    continue
+
+                experts = ep_target.experts
+                missing = [a for a in ("w1", "w2") if not hasattr(experts, a)]
+                if missing:
+                    raise ValueError(
+                        f"Module {type(ep_target).__name__}.experts is not grouped-experts compatible. "
+                        f"Missing: {missing}"
+                    )
+                if experts.w1.ndim != 3 or experts.w2.ndim != 3:
+                    raise ValueError(
+                        f"Expected grouped expert parameters with ndim=3. Got w1.ndim={experts.w1.ndim}, "
+                        f"w2.ndim={experts.w2.ndim}"
+                    )
+
+                ep_target._ep_mesh = ep_mesh
+                ep_target._ep_group = ep_mesh.get_group()
+                ep_target._ep_size = ep_mesh.size()
+                ep_target._ep_rank = ep_mesh.get_local_rank()
+
+                ep_target.experts = ExpertParallel()._apply(ep_target.experts, ep_mesh)
+                setattr(ep_target.experts, "_ep_enabled", True)
+
+                for pname, p in list(ep_target.experts._parameters.items()):
+                    if isinstance(p, DTensor) and p.dtype != target_dtype:
+                        local = p.to_local().to(target_dtype)
+                        ep_target.experts._parameters[pname] = nn.Parameter(
+                            DTensor.from_local(local, p.device_mesh, p.placements, run_check=False),
+                            requires_grad=p.requires_grad,
+                        )
+
+                wrapped_blocks += 1
+
+        if wrapped_blocks == 0:
+            raise ValueError(f"No blocks matched the requested types: {[t.__name__ for t in block_types]}")
+
+        model._ep_wrapped = True
+        model._ep_mesh = ep_mesh
+        model._ep_num_wrapped_blocks = wrapped_blocks
+
+        return model
+
     @staticmethod
     def get_weight_initialized_model(model: nn.Module, model_initializer: ModelInitializationIF) -> nn.Module:
         """

src/modalities/models/moe/__init__.py

@@ -1,10 +1,8 @@
-from modalities.models.moe.loss_functions import MoECrossEntropyLoss
-from modalities.models.moe.model_factory import get_ep_wrapped_model
+from modalities.loss_functions import MoECrossEntropyLoss
 from modalities.models.moe.qwen_model import QwenModel, QwenModelConfig
 
 __all__ = [
     "MoECrossEntropyLoss",
     "QwenModel",
     "QwenModelConfig",
-    "get_ep_wrapped_model",
 ]

src/modalities/registry/components.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,8 +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.model_factory import get_ep_wrapped_model
 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 (
@@ -196,7 +194,7 @@ class ComponentEntity:
     # models
     ComponentEntity("model", "gpt2", GPT2ModelFactory.get_gpt2_model, GPT2LLMConfig),
     ComponentEntity("model", "moe", QwenModel, QwenModelConfig),
-    ComponentEntity("model", "ep_wrapped", get_ep_wrapped_model, EPWrappedModelConfig),
+    ComponentEntity("model", "ep_wrapped", ModelFactory.get_ep_wrapped_model, EPWrappedModelConfig),
     ComponentEntity(
         "model", "gpt2_tp", maybe_model_list(GPT2ModelFactory.get_gpt2_tensor_parallelized_model), GPT2ModelTPConfig
     ),

tests/dataloader/distributed/test_distributed_multidim_dataloader.py

@@ -93,6 +93,7 @@ def _test_dataloader_produces_different_samples_in_different_dp_ranks(process_id
             tensor_parallel_degree=2,
             pipeline_parallel_degree=2,
             context_parallel_degree=1,
+            expert_parallel_degree=1,
             enable_loss_parallel=False,
             world_size=world_size,
         )
@@ -117,6 +118,7 @@ def _test_dataloader_produces_same_samples_in_connected_non_dp_ranks(process_id:
             tensor_parallel_degree=2,
             pipeline_parallel_degree=2,
             context_parallel_degree=1,
+            expert_parallel_degree=1,
             enable_loss_parallel=False,
             world_size=world_size,
         )

tests/models/moe/test_loss_functions.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: