expert parallelism config

Author: khatwanimohit

tunix/generate/vllm_sampler.py

@@ -64,6 +64,7 @@ class VllmConfig:
   mesh: jax.sharding.Mesh = None
   data_parallel_size: int = -1
   tensor_parallel_size: int = -1
+  expert_parallel_size: int = 1
 
   # vLLM engine args that can be directly passed in without additional processing, e.g. max_model_len, async_scheduling, etc.
   engine_kwargs: dataclasses.InitVar[Optional[Dict[str, Any]]] = None
@@ -204,25 +205,49 @@ def _find_total_size(self, mesh: jax.sharding.Mesh) -> int:
     # since vllm doesn't support DP yet, simply return the total rank size.
     return math.prod(mesh.shape.values())
 
-  def _vllm_config(self, config: VllmConfig):
-    """Setup vllm config from Tunix Vllm config."""
-    args = config._processed_engine_kwargs.copy()
-
+  def _configure_sharding(
+      self, config: VllmConfig, args: Dict[str, Any]
+  ) -> None:
+    """Resolves parallelism sizes and sets the sharding config in args."""
     tensor_parallel_size = config.tensor_parallel_size
     data_parallel_size = config.data_parallel_size
+    expert_parallel_size = config.expert_parallel_size
     total_mesh_devices = self._find_total_size(config.mesh)
 
+    if total_mesh_devices % expert_parallel_size != 0:
+      raise ValueError(
+          f"Total mesh devices ({total_mesh_devices}) must be divisible by"
+          f" expert_parallel_size ({expert_parallel_size})."
+      )
+
     if config.tensor_parallel_size == -1 and config.data_parallel_size == -1:
-      tensor_parallel_size = total_mesh_devices
+      tensor_parallel_size = total_mesh_devices // expert_parallel_size
       data_parallel_size = 1
     elif config.tensor_parallel_size == -1:
-      tensor_parallel_size = total_mesh_devices // data_parallel_size
+      tensor_parallel_size = (
+          total_mesh_devices // (data_parallel_size * expert_parallel_size)
+      )
     elif config.data_parallel_size == -1:
-      data_parallel_size = total_mesh_devices // tensor_parallel_size
+      data_parallel_size = (
+          total_mesh_devices // (tensor_parallel_size * expert_parallel_size)
+      )
 
     args["data_parallel_size"] = data_parallel_size
     args["tensor_parallel_size"] = tensor_parallel_size
 
+    device_indexes = config.mesh.device_ids.flatten().tolist()
+    args["additional_config"]["sharding"] = {
+        "sharding_strategy": {
+            "expert_parallelism": expert_parallel_size,
+            "device_indexes": device_indexes,
+            "enable_dp_attention": config.enable_dp_attention,
+        }
+    }
+
+  def _vllm_config(self, config: VllmConfig):
+    """Setup vllm config from Tunix Vllm config."""
+    args = config._processed_engine_kwargs.copy()
+
     # Init vLLM model with random weights to speed up bootstrap time, because
     # model weights are synced from trainer later on
     if config.init_with_random_weights:
@@ -235,14 +260,7 @@ def _vllm_config(self, config: VllmConfig):
     if config.lora_config is not None:
       args["additional_config"]["lora_config"] = config.lora_config
 
-    device_indexes = config.mesh.device_ids.flatten().tolist()
-
-    args["additional_config"]["sharding"] = {
-        "sharding_strategy": {
-            "device_indexes": device_indexes,
-            "enable_dp_attention": config.enable_dp_attention,
-        }
-    }
+    self._configure_sharding(config, args)
 
     return args
 

tunix/rl/rollout/base_rollout.py

@@ -111,6 +111,7 @@ class RolloutConfig:
   # Parallelism configs.
   tensor_parallel_size: int = -1
   data_parallel_size: int = -1
+  expert_parallel_size: int = 1
 
   # vLLM specific rollout configs.
 

tunix/rl/rollout/vllm_rollout.py

@@ -49,6 +49,7 @@ def __init__(
             tpu_backend_type=rollout_config.rollout_vllm_tpu_backend_type,
             additional_config=rollout_config.rollout_vllm_additional_config,
             enable_dp_attention=rollout_config.rollout_vllm_enable_dp_attention,
+            expert_parallel_size=rollout_config.expert_parallel_size,
             hbm_utilization=rollout_config.rollout_vllm_hbm_utilization,
             lora_config=rollout_config.rollout_vllm_lora_config,
             mesh=mesh,