expert parallelism config

Author: khatwanimohit

tests/generate/utils_test.py

@@ -20,6 +20,7 @@
 from jax import sharding
 import jax.numpy as jnp
 import numpy as np
+from unittest import mock
 from tunix.generate import utils
 from tunix.rl import reshard
 
@@ -1185,5 +1186,47 @@ def test_transfer_state_directly_scanned_layers_casting(self):
     )
 
 
+class ResolveParallelismSizesTest(parameterized.TestCase):
+
+  def _make_mesh(self, total_devices):
+    """Returns a mock mesh with the given total device count."""
+    mesh = mock.MagicMock()
+    mesh.shape = {"axis": total_devices}
+    return mesh
+
+  @parameterized.named_parameters(
+      ("tp_and_dp_inferred_no_ep", 8, -1, -1, 1, 8, 1, 1),
+      ("tp_and_dp_inferred_with_ep", 8, -1, -1, 2, 4, 1, 2),
+      ("tp_inferred_with_ep", 8, -1, 2, 2, 2, 2, 2),
+      ("dp_inferred_with_ep", 8, 2, -1, 2, 2, 2, 2),
+      ("all_explicit", 8, 4, 2, 1, 4, 2, 1),
+  )
+  def test_resolve_parallelism_sizes(
+      self,
+      total_devices,
+      tp_in,
+      dp_in,
+      ep_in,
+      expected_tp,
+      expected_dp,
+      expected_ep,
+  ):
+    mesh = self._make_mesh(total_devices)
+    tp, dp, ep = utils.resolve_parallelism_sizes(
+        mesh=mesh,
+        tensor_parallel_size=tp_in,
+        data_parallel_size=dp_in,
+        expert_parallel_size=ep_in,
+    )
+    self.assertEqual(tp, expected_tp)
+    self.assertEqual(dp, expected_dp)
+    self.assertEqual(ep, expected_ep)
+
+  def test_resolve_parallelism_sizes_indivisible_ep_raises(self):
+    mesh = self._make_mesh(8)
+    with self.assertRaisesRegex(ValueError, "expert_parallel_size"):
+      utils.resolve_parallelism_sizes(mesh=mesh, expert_parallel_size=3)
+
+
 if __name__ == "__main__":
   absltest.main()

tests/generate/vllm_sampler_test.py

@@ -358,5 +358,58 @@ async def dispatch_requests():
     )
 
 
+class VllmSamplerConfigTest(absltest.TestCase):
+  """Unit tests for VllmSampler config plumbing (no hardware required)."""
+
+  def _make_mock_mesh(self, total_devices):
+    mesh = mock.MagicMock()
+    mesh.shape = {"axis": total_devices}
+    mesh.device_ids.flatten.return_value.tolist.return_value = list(
+        range(total_devices)
+    )
+    return mesh
+
+  def _make_sampler(self, config):
+    with mock.patch("tunix.generate.vllm_sampler.LLM"), mock.patch(
+        "tunix.generate.vllm_sampler.tok_adapter.TokenizerAdapter"
+    ):
+      return vllm_sampler.VllmSampler(
+          tokenizer=mock.MagicMock(), config=config
+      )
+
+  def test_expert_parallel_size_plumbed_to_sharding(self):
+    mesh = self._make_mock_mesh(8)
+    config = vllm_sampler.VllmConfig(
+        mesh=mesh,
+        expert_parallel_size=2,
+        init_with_random_weights=False,
+    )
+    sampler = self._make_sampler(config)
+
+    sharding_strategy = sampler.args["additional_config"]["sharding"][
+        "sharding_strategy"
+    ]
+    # EP=2 should appear in the sharding strategy passed to vLLM.
+    self.assertEqual(sharding_strategy["expert_parallelism"], 2)
+    # With 8 total devices and EP=2, TP should be inferred as 4 and DP as 1.
+    self.assertEqual(sampler.args["tensor_parallel_size"], 4)
+    self.assertEqual(sampler.args["data_parallel_size"], 1)
+
+  def test_default_expert_parallel_size_is_one(self):
+    mesh = self._make_mock_mesh(8)
+    config = vllm_sampler.VllmConfig(
+        mesh=mesh,
+        init_with_random_weights=False,
+    )
+    sampler = self._make_sampler(config)
+
+    sharding_strategy = sampler.args["additional_config"]["sharding"][
+        "sharding_strategy"
+    ]
+    self.assertEqual(sharding_strategy["expert_parallelism"], 1)
+    self.assertEqual(sampler.args["tensor_parallel_size"], 8)
+    self.assertEqual(sampler.args["data_parallel_size"], 1)
+
+
 if __name__ == "__main__":
   absltest.main()

tunix/generate/utils.py

@@ -991,6 +991,50 @@ def intersect_trees(
   gc.collect()
 
 
+def resolve_parallelism_sizes(
+    mesh: jax.sharding.Mesh,
+    tensor_parallel_size: int = -1,
+    data_parallel_size: int = -1,
+    expert_parallel_size: int = 1,
+) -> tuple[int, int, int]:
+  """Resolves tensor, data, and expert parallelism sizes from the mesh.
+
+  Any size passed as -1 is inferred from the total number of mesh devices and
+  the other sizes. Raises ValueError if the mesh size is not divisible by
+  expert_parallel_size.
+
+  Args:
+    mesh: The JAX device mesh.
+    tensor_parallel_size: Desired tensor parallelism degree, or -1 to infer.
+    data_parallel_size: Desired data parallelism degree, or -1 to infer.
+    expert_parallel_size: Desired expert parallelism degree.
+
+  Returns:
+    A tuple of (tensor_parallel_size, data_parallel_size, expert_parallel_size).
+  """
+  total_mesh_devices = math.prod(mesh.shape.values())
+
+  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 tensor_parallel_size == -1 and data_parallel_size == -1:
+    tensor_parallel_size = total_mesh_devices // expert_parallel_size
+    data_parallel_size = 1
+  elif tensor_parallel_size == -1:
+    tensor_parallel_size = (
+        total_mesh_devices // (data_parallel_size * expert_parallel_size)
+    )
+  elif data_parallel_size == -1:
+    data_parallel_size = (
+        total_mesh_devices // (tensor_parallel_size * expert_parallel_size)
+    )
+
+  return tensor_parallel_size, data_parallel_size, expert_parallel_size
+
+
 def verify_state_closeness(golden_state, state, atol=1e-2):
   """Check if the golden NNX state is close to the other NNX state.
 

tunix/generate/vllm_sampler.py

@@ -17,7 +17,6 @@
 import atexit
 import dataclasses
 from itertools import count
-import math
 import os
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -64,6 +63,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
@@ -199,30 +199,10 @@ def load_checkpoint(self, path_or_weights: str | jaxtyping.PyTree):
     else:
       raise NotImplementedError("Only support in memory weight sync as of now.")
 
-  def _find_total_size(self, mesh: jax.sharding.Mesh) -> int:
-    """Finds the tensor parallel size from the mesh."""
-    # 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()
 
-    tensor_parallel_size = config.tensor_parallel_size
-    data_parallel_size = config.data_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
-      data_parallel_size = 1
-    elif config.tensor_parallel_size == -1:
-      tensor_parallel_size = total_mesh_devices // data_parallel_size
-    elif config.data_parallel_size == -1:
-      data_parallel_size = total_mesh_devices // tensor_parallel_size
-
-    args["data_parallel_size"] = data_parallel_size
-    args["tensor_parallel_size"] = tensor_parallel_size
-
     # 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,10 +215,19 @@ 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()
+    tp, dp, ep = utils.resolve_parallelism_sizes(
+        mesh=config.mesh,
+        tensor_parallel_size=config.tensor_parallel_size,
+        data_parallel_size=config.data_parallel_size,
+        expert_parallel_size=config.expert_parallel_size,
+    )
+    args["tensor_parallel_size"] = tp
+    args["data_parallel_size"] = dp
 
+    device_indexes = config.mesh.device_ids.flatten().tolist()
     args["additional_config"]["sharding"] = {
         "sharding_strategy": {
+            "expert_parallelism": ep,
             "device_indexes": device_indexes,
             "enable_dp_attention": config.enable_dp_attention,
         }
@@ -414,7 +403,6 @@ def __call__(
         sampling_params.top_p = top_p
       if top_k is not None:
         sampling_params.top_k = top_k
-
       if seed is not None:
         sampling_params.seed = seed
 

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

@@ -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
                 ),