feat: add dedicated EP dimension to device mesh

Author: gbesposito

config_files/training/config_lorem_ipsum_long_moe_ep_fsdp2.yaml

@@ -184,7 +184,7 @@ device_mesh:
     device_type: cuda
     data_parallel_replicate_degree: 1
     data_parallel_shard_degree: -1
-    tensor_parallel_degree: 4
+    expert_parallel_degree: 4
     world_size: ${settings.cuda_env.world_size}
 
 dp_degree:
@@ -238,7 +238,6 @@ ep_model:
     device_mesh:
       instance_key: device_mesh
       pass_type: BY_REFERENCE
-    ep_mesh_dim_name: tp
     block_names: [TransformerBlock]
 
 ac_model:

src/modalities/config/config.py

@@ -338,7 +338,6 @@ class EPWrappedModelConfig(BaseModel):
     model: PydanticPytorchModuleOrListType
     block_names: list[str]
     device_mesh: PydanticDeviceMeshIFType
-    ep_mesh_dim_name: str | None = None
 
 
 class DebuggingEnrichedModelConfig(BaseModel):

src/modalities/models/moe/model_factory.py

@@ -5,28 +5,12 @@
 import torch.nn as nn
 from torch.distributed._composable.fsdp import MixedPrecisionPolicy
 from torch.distributed.device_mesh import DeviceMesh
+
 from modalities.models.parallelism.expert_parallelism import ExpertParallel
+from modalities.running_env.fsdp.device_mesh import ParallelismDegrees, get_mesh_for_parallelism_method
 from modalities.util import get_module_class_from_name
 
 
-# TODO refactor these funtions into a utils
-def _resolve_ep_mesh(device_mesh: DeviceMesh, ep_mesh_dim_name: str | None) -> DeviceMesh:
-    mesh_dim_names = tuple(device_mesh.mesh_dim_names or ())
-
-    if ep_mesh_dim_name is not None:
-        if ep_mesh_dim_name not in mesh_dim_names:
-            raise ValueError(f"ep_mesh_dim_name='{ep_mesh_dim_name}' not in mesh_dim_names={mesh_dim_names}")
-        return device_mesh[ep_mesh_dim_name]
-
-    if len(mesh_dim_names) <= 1:
-        return device_mesh
-
-    raise ValueError(
-        "DeviceMesh has multiple dimensions. Pass ep_mesh_dim_name explicitly. "
-        f"Available dimensions: {mesh_dim_names}"
-    )
-
-
 def _validate_moe_block_for_ep(module) -> None:
     if not hasattr(module, "experts"):
         raise ValueError(f"Module {type(module).__name__} has no 'experts' attribute")
@@ -64,16 +48,10 @@ def _attach_ep_metadata(module, ep_mesh) -> None:
     setattr(module, "_ep_rank", ep_mesh.get_local_rank())
 
 
-def _apply_ep(module, ep_mesh) -> None:
-    module.experts = ExpertParallel()._apply(module.experts, ep_mesh)
-    setattr(module.experts, "_ep_enabled", True)
-
-
 def get_ep_wrapped_model(
     model,
     block_names: list[str],
     device_mesh: DeviceMesh,
-    ep_mesh_dim_name: str | None = None,
     mp_param_dtype=torch.bfloat16,
     mp_reduce_dtype=torch.bfloat16,
 ) -> nn.Module:
@@ -97,7 +75,7 @@ def get_ep_wrapped_model(
     if len(block_types) == 0:
         raise ValueError(f"None of the requested MoE block names were found: {block_names}")
 
-    ep_mesh = _resolve_ep_mesh(device_mesh, ep_mesh_dim_name)
+    ep_mesh = get_mesh_for_parallelism_method(device_mesh, ParallelismDegrees.EP)
     MixedPrecisionPolicy(param_dtype=mp_param_dtype, reduce_dtype=mp_reduce_dtype)
 
     wrapped_blocks = 0
@@ -115,7 +93,9 @@ def get_ep_wrapped_model(
 
             _validate_moe_block_for_ep(ep_target_module)
             _attach_ep_metadata(ep_target_module, ep_mesh)
-            _apply_ep(ep_target_module, ep_mesh)
+
+            ep_target_module.experts = ExpertParallel()._apply(ep_target_module.experts, ep_mesh)
+            setattr(ep_target_module.experts, "_ep_enabled", True)
 
             wrapped_blocks += 1
 

src/modalities/running_env/fsdp/device_mesh.py

@@ -21,13 +21,15 @@ class DeviceMeshConfig(BaseModel):
     tensor_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1
     pipeline_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1
     context_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1
+    expert_parallel_degree: Annotated[int, Field(strict=True, gt=0)] = 1
     enable_loss_parallel: Optional[bool] = False
     world_size: Annotated[int, Field(strict=True, gt=0)]
 
     @model_validator(mode="after")
     def _validate(self):
         for d in (
             self.context_parallel_degree,
+            self.expert_parallel_degree,
             self.tensor_parallel_degree,
             self.pipeline_parallel_degree,
         ):
@@ -50,6 +52,7 @@ def _validate(self):
             self.data_parallel_shard_degree = self.world_size // (
                 self.data_parallel_replicate_degree
                 * self.context_parallel_degree
+                * self.expert_parallel_degree
                 * self.tensor_parallel_degree
                 * self.pipeline_parallel_degree
             )
@@ -58,12 +61,14 @@ def _validate(self):
             self.data_parallel_replicate_degree = self.world_size // (
                 self.data_parallel_shard_degree
                 * self.context_parallel_degree
+                * self.expert_parallel_degree
                 * self.tensor_parallel_degree
                 * self.pipeline_parallel_degree
             )
         if (
             self.data_parallel_shard_degree
             * self.data_parallel_replicate_degree
+            * self.expert_parallel_degree
             * self.tensor_parallel_degree
             * self.pipeline_parallel_degree
             * self.context_parallel_degree
@@ -72,6 +77,7 @@ def _validate(self):
             raise ConfigError(
                 f"Invalid parallel dims: data_parallel_shard_degree({self.data_parallel_shard_degree}) * "
                 f"data_parallel_replicate_degree({self.data_parallel_replicate_degree}) * "
+                f"expert_parallel_degree({self.expert_parallel_degree}) * "
                 f"tensor_parallel_degree({self.tensor_parallel_degree}) *"
                 f"* pipeline_parallel_degree({self.pipeline_parallel_degree}) *"
                 f"context_parallel_degree({self.context_parallel_degree})!= WORLD_SIZE({self.world_size})"
@@ -85,6 +91,7 @@ class ParallelismDegrees(Enum):
     DP_REPLICATE = "dp_replicate"
     DP_SHARD = "dp_shard"
     CP = "cp"
+    EP = "ep"
     TP = "tp"
     PP = "pp"
 
@@ -96,6 +103,7 @@ def get_device_mesh(
     tensor_parallel_degree: int,
     pipeline_parallel_degree: int,
     context_parallel_degree: int,
+    expert_parallel_degree: int,
     enable_loss_parallel: bool,
     world_size: int,
 ) -> DeviceMesh:
@@ -109,6 +117,7 @@ def get_device_mesh(
         tensor_parallel_degree (int): The tensor parallel degree.
         pipeline_parallel_degree (int): The pipeline parallel degree.
         context_parallel_degree (int): The context parallel degree.
+        expert_parallel_degree (int): The expert parallel degree.
         enable_loss_parallel (bool): Whether to enable loss parallelism.
         world_size (int): The world size.
 
@@ -123,13 +132,15 @@ def get_device_mesh(
             data_parallel_replicate_degree,
             data_parallel_shard_degree,
             context_parallel_degree,
+            expert_parallel_degree,
             tensor_parallel_degree,
         ],
         [
             ParallelismDegrees.PP.value,
             ParallelismDegrees.DP_REPLICATE.value,
             ParallelismDegrees.DP_SHARD.value,
             ParallelismDegrees.CP.value,
+            ParallelismDegrees.EP.value,
             ParallelismDegrees.TP.value,
         ],
         strict=True,