feat: Allow for multiple parallelism methods in get_parallel_degree

Author: rrutmann

config_files/training/config_lorem_ipsum_long_fsdp2.yaml

@@ -183,13 +183,13 @@ device_mesh:
     world_size: ${settings.cuda_env.world_size}
 
 dp_degree:
-  component_key: device_mesh
+  component_key: number_conversion
   variant_key: parallel_degree
   config: # get the parallel degree from the device mesh
     device_mesh:
       instance_key: device_mesh
       pass_type: BY_REFERENCE
-    parallelism_method: dp_shard
+    parallelism_methods: [dp_shard, dp_replicate]
 
 app_state:
   component_key: app_state

src/modalities/config/config.py

@@ -485,7 +485,7 @@ class GPT2MFUCalculatorConfig(BaseModel):
 
 class ParallelDegreeConfig(BaseModel):
     device_mesh: PydanticDeviceMeshIFType
-    parallelism_method: ParallelismDegrees
+    parallelism_methods: list[ParallelismDegrees]
 
 
 def load_app_config_dict(

src/modalities/main.py

@@ -141,7 +141,8 @@ def run(self, components: TrainingComponentsInstantiationModel):
 
         # Trainer
         dp_degree = get_parallel_degree(
-            device_mesh=components.device_mesh, parallelism_method=ParallelismDegrees.DP_SHARD
+            device_mesh=components.device_mesh,
+            parallelism_methods=[ParallelismDegrees.DP_SHARD, ParallelismDegrees.DP_REPLICATE],
         )
         global_num_tokens_per_train_step = (
             components.settings.step_profile.local_train_micro_batch_size

src/modalities/registry/components.py

@@ -177,7 +177,7 @@ class ComponentEntity:
     ),
     # Device mesh
     ComponentEntity("device_mesh", "default", get_device_mesh, DeviceMeshConfig),
-    ComponentEntity("device_mesh", "parallel_degree", get_parallel_degree, ParallelDegreeConfig),
+    ComponentEntity("number_conversion", "parallel_degree", get_parallel_degree, ParallelDegreeConfig),
     # weight initializers
     ComponentEntity(
         "model_initialization",

src/modalities/running_env/fsdp/device_mesh.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from math import prod
 from typing import Annotated, Optional
 
 from pydantic import BaseModel, Field, model_validator
@@ -129,19 +130,20 @@ def get_device_mesh(
     return device_mesh
 
 
-def get_parallel_degree(device_mesh: DeviceMesh, parallelism_method: ParallelismDegrees) -> int:
+def get_parallel_degree(device_mesh: DeviceMesh, parallelism_methods: list[ParallelismDegrees]) -> int:
     """Gets the number of parallel ranks (i.e., the parallelism degree)
     from the device mesh for a specific parallelism method.
     Args:
         device_mesh (DeviceMesh): The device mesh.
-        parallelism_method (ParallelismDegrees): The parallelism method.
+        parallelism_methods (list[ParallelismDegrees]): The parallelism methods.
     Returns:
         int: The number of parallel ranks for the specified parallelism method.
     """
     if device_mesh.mesh_dim_names is None:
         raise ValueError("device_mesh.mesh_dim_names is None")
 
-    if parallelism_method.value not in device_mesh.mesh_dim_names:
-        return 1
-    else:
-        return device_mesh.size(device_mesh.mesh_dim_names.index(parallelism_method.value))
+    return prod(
+        device_mesh.size(device_mesh.mesh_dim_names.index(method.value))
+        for method in parallelism_methods
+        if method.value in device_mesh.mesh_dim_names
+    )