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
@@ -210,15 +211,20 @@ def _vllm_config(self, config: VllmConfig):
 
     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 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

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.
 
@@ -149,6 +150,9 @@ class RolloutConfig:
   # axes, which can help reduce memory usage for large models with few KV heads.
   rollout_vllm_enable_dp_attention: bool = False
 
+  # Whether to enable expert parallelism for vLLM rollout engine.
+  rollout_vllm_enable_expert_parallelism: bool = False
+
   # Maximum number of batched tokens allowed in vLLM. This allows for pending prefill requests
   # to be batched along with decode requests if enough tokens are available. Only used when
   # chunked prefill is enabled.

tunix/rl/rollout/vllm_rollout.py

@@ -61,6 +61,7 @@ def __init__(
                 ),
                 "tensor_parallel_size": rollout_config.tensor_parallel_size,
                 "data_parallel_size": rollout_config.data_parallel_size,
+                "expert_parallel_size": rollout_config.expert_parallel_size,
                 "max_num_batched_tokens": (
                     rollout_config.rollout_vllm_max_num_batched_tokens
                 ),