Add vllm arg to enable expert parallelism for MoE models

Author: dipannita08

src/maxtext/eval/runner/eval_runner.py

@@ -165,6 +165,10 @@ def run_eval(cfg: dict, hf_token: str | None = None) -> dict:
 
   # Start vLLM server.
   server_env = {"HF_TOKEN": token} if token else None
+  additional_vllm_kwargs = {}
+  if cfg.get("enable_expert_parallel"):
+    additional_vllm_kwargs["enable_expert_parallel"] = True
+
   with VllmServerManager(
       model_path=hf_path,
       checkpoint_path=checkpoint_path if use_maxtext_adapter else None,
@@ -176,6 +180,7 @@ def run_eval(cfg: dict, hf_token: str | None = None) -> dict:
       max_num_batched_tokens=max_num_batched_tokens,
       max_num_seqs=max_num_seqs,
       env=server_env,
+      additional_vllm_kwargs=additional_vllm_kwargs or None,
   ) as server:
     base_url = server.base_url
 
@@ -255,6 +260,16 @@ def _build_arg_parser() -> argparse.ArgumentParser:
   parser.add_argument("--server_host", help="vLLM server host.")
   parser.add_argument("--server_port", type=int, help="vLLM server port.")
   parser.add_argument("--hf_mode", action="store_true", help="Use HF safetensors mode.")
+  parser.add_argument(
+      "--enable_expert_parallel",
+      action="store_true",
+      help=(
+          "Enable expert parallelism in vLLM. Required for MoE models such as "
+          "qwen3-30b-a3b, qwen3-235b-a22b, deepseek-v3, etc. Without this flag "
+          "tpu-inference omits the 'expert' mesh axis and MaxText's MoE sharding "
+          "raises KeyError."
+      ),
+  )
   parser.add_argument("--hf_token", help="HuggingFace token for gated models.")
   parser.add_argument(
       "--log_level",

src/maxtext/eval/runner/evalchemy_runner.py

@@ -33,7 +33,7 @@
       --tensor_parallel_size 4 \\
       --hf_token $HF_TOKEN
 
-Requires: pip install evalchemy
+Requires: pip install git+https://github.com/mlfoundations/evalchemy.git
 """
 
 from __future__ import annotations
@@ -200,6 +200,9 @@ def run_evalchemy(cfg: dict, hf_token: str | None = None) -> dict:
     lm_eval_tasks.append(lm_eval_task)
 
   server_env = {"HF_TOKEN": token} if token else None
+  additional_vllm_kwargs = {}
+  if cfg.get("enable_expert_parallel"):
+    additional_vllm_kwargs["enable_expert_parallel"] = True
 
   with VllmServerManager(
       model_path=hf_path,
@@ -212,6 +215,7 @@ def run_evalchemy(cfg: dict, hf_token: str | None = None) -> dict:
       max_num_batched_tokens=max_num_batched_tokens,
       max_num_seqs=max_num_seqs,
       env=server_env,
+      additional_vllm_kwargs=additional_vllm_kwargs or None,
   ) as server:
     warmup_server(base_url=server.base_url, model=model_name)
 
@@ -329,6 +333,16 @@ def _build_arg_parser() -> argparse.ArgumentParser:
       action="store_true",
       help="HF safetensors mode.",
   )
+  parser.add_argument(
+      "--enable_expert_parallel",
+      action="store_true",
+      help=(
+          "Enable expert parallelism in vLLM. Required for MoE models such as "
+          "qwen3-30b-a3b, qwen3-235b-a22b, deepseek-v3, etc. Without this flag "
+          "tpu-inference omits the 'expert' mesh axis and MaxText's MoE sharding "
+          "raises KeyError."
+      ),
+  )
   parser.add_argument(
       "--hf_token",
       help="HuggingFace token for gated tokenizers.",

src/maxtext/eval/runner/lm_eval_runner.py

@@ -142,6 +142,9 @@ def run_lm_eval(cfg: dict, hf_token: str | None = None) -> dict:
     lm_tasks.append(lm_task)
 
   server_env = {"HF_TOKEN": token} if token else None
+  additional_vllm_kwargs = {}
+  if cfg.get("enable_expert_parallel"):
+    additional_vllm_kwargs["enable_expert_parallel"] = True
 
   with VllmServerManager(
       model_path=hf_path,
@@ -154,6 +157,7 @@ def run_lm_eval(cfg: dict, hf_token: str | None = None) -> dict:
       max_num_batched_tokens=max_num_batched_tokens,
       max_num_seqs=max_num_seqs,
       env=server_env,
+      additional_vllm_kwargs=additional_vllm_kwargs or None,
   ) as server:
     warmup_server(base_url=server.base_url, model=model_name)
 
@@ -213,6 +217,16 @@ def _build_arg_parser() -> argparse.ArgumentParser:
   parser.add_argument("--num_samples", type=int, help="Limit samples per task (None = full dataset).")
   parser.add_argument("--hf_token", help="HuggingFace token for gated tokenizers.")
   parser.add_argument("--hf_mode", action="store_true", help="HF safetensors mode.")
+  parser.add_argument(
+      "--enable_expert_parallel",
+      action="store_true",
+      help=(
+          "Enable expert parallelism in vLLM. Required for MoE models such as "
+          "qwen3-30b-a3b, qwen3-235b-a22b, deepseek-v3, etc. Without this flag "
+          "tpu-inference omits the 'expert' mesh axis and MaxText's MoE sharding "
+          "raises KeyError."
+      ),
+  )
   parser.add_argument("--gcs_results_path", help="Optional GCS path to upload results.")
   parser.add_argument("--log_level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"])
   return parser

src/maxtext/eval/runner/server_manager.py

@@ -183,7 +183,7 @@ class VllmServerManager:
   """Manages an in-process vLLM-TPU LLM with an OpenAI-compatible HTTP layer.
 
   Args:
-    model_path: HF model ID or local path. 
+    model_path: HF model ID or local path.
     checkpoint_path: MaxText orbax checkpoint path.
     maxtext_model_name: MaxText model name (e.g. "llama3.1-8b").
     host: Hostname the HTTP server binds to (rank-0 only).
@@ -195,6 +195,7 @@ class VllmServerManager:
     max_num_seqs: Max concurrent sequences (None = vLLM default).
     startup_timeout: Seconds to wait for /health to return healthy.
     env: Optional environment-variable overrides.
+    additional_vllm_kwargs: Extra kwargs merged into the vLLM LLM() constructor.
   """
 
   def __init__(
@@ -211,6 +212,7 @@ def __init__(
       max_num_seqs: int | None = None,
       startup_timeout: int = 600,
       env: dict[str, str] | None = None,
+      additional_vllm_kwargs: dict | None = None,
   ):
     if checkpoint_path and not maxtext_model_name:
       raise ValueError("maxtext_model_name is required when checkpoint_path is set.")
@@ -226,6 +228,7 @@ def __init__(
     self.max_num_seqs = max_num_seqs
     self.startup_timeout = startup_timeout
     self.env = env
+    self.additional_vllm_kwargs = additional_vllm_kwargs or {}
 
     self._llm: Any | None = None
     self._uvicorn_server: Any | None = None
@@ -271,6 +274,9 @@ def start(self) -> None:
     else:
       vllm_kwargs["load_format"] = "auto"
 
+    if self.additional_vllm_kwargs:
+      vllm_kwargs.update(self.additional_vllm_kwargs)
+
     logger.info(
         "Initializing in-process vLLM (tp=%d, max_len=%d)...",
         self.tensor_parallel_size,