Add DP axis, test multihost, fix tests

Author: dipannita08

src/maxtext/eval/README.md

@@ -86,8 +86,9 @@ python -m maxtext.eval.runner.run \
 | `--model_name` | MaxText model name (e.g. `llama3.1-8b`) |
 | `--hf_path` | HF model ID or local path |
 | `--max_model_len` | vLLM max context length. |
-| `--tensor_parallel_size` | Total number of chips |
-| `--expert_parallel_size` | Number of EP chips |
+| `--tensor_parallel_size` | Chips per model replica |
+| `--expert_parallel_size` | Chips for the expert mesh axis |
+| `--data_parallel_size` | Number of model replicas |
 | `--hbm_memory_utilization` | Fraction of HBM reserved for KV cache |
 | `--hf_token` | HF token (or set `HF_TOKEN` env var) |
 | `--hf_mode` | HF safetensors mode, no MaxText checkpoint loading |

src/maxtext/eval/runner/async_client.py

@@ -31,7 +31,7 @@
 _DEFAULT_MAX_TOKENS = 1024
 _DEFAULT_TEMPERATURE = 0.0
 _COMPLETIONS_PATH = "/v1/completions"
-_REQUEST_TIMEOUT_S = 600 # (TODO): Check if this is reasoanable.
+_REQUEST_TIMEOUT_S = 600
 
 
 @dataclass

src/maxtext/eval/runner/common.py

@@ -65,6 +65,7 @@ def build_server_manager(cfg: dict, token: str | None) -> "VllmServerManager":
     max_num_seqs = int(max_num_seqs)
 
   expert_parallel_size = int(cfg.get("expert_parallel_size") or 1)
+  data_parallel_size = int(cfg.get("data_parallel_size") or 1)
   hbm_memory_utilization = float(cfg.get("hbm_memory_utilization") or 0.3)
 
   server_env = {"HF_TOKEN": token} if token else None
@@ -77,6 +78,7 @@ def build_server_manager(cfg: dict, token: str | None) -> "VllmServerManager":
       port=server_port,
       tensor_parallel_size=tensor_parallel_size,
       expert_parallel_size=expert_parallel_size,
+      data_parallel_size=data_parallel_size,
       max_model_len=max_model_len,
       max_num_batched_tokens=max_num_batched_tokens,
       max_num_seqs=max_num_seqs,
@@ -122,6 +124,12 @@ def add_server_args(parser: argparse.ArgumentParser) -> None:
           "Chips allocated to the expert mesh axis (EP). "
       ),
   )
+  parser.add_argument(
+      "--data_parallel_size",
+      type=int,
+      default=1,
+      help="Number of model replicas (DP).",
+  )
   parser.add_argument(
       "--hbm_memory_utilization",
       type=float,

src/maxtext/eval/runner/harness_runner.py

@@ -42,39 +42,13 @@
 
 logger = logging.getLogger(__name__)
 
-_TASK_MAP_LM_EVAL: dict[str, str] = {
-    "mmlu": "mmlu",
-    "gpqa": "gpqa_diamond",
-    "math": "hendrycks_math",
-    "gsm8k": "gsm8k",
-}
-
-_TASK_MAP_EVALCHEMY: dict[str, str] = {
-    "ifeval": "ifeval",
-    "alpacaeval": "alpaca_eval_v2",
-    "arena_hard": "arena_hard",
-    "mtbench": "mt_bench",
-    "wildbench": "wildbench",
-    "mixeval": "mixeval",
-    "zeroeval": "zeroeval",
-    "math500": "math_500",
-    "aime24": "aime2024",
-    "aime25": "aime2025",
-    "amc23": "amc2023",
-    "gpqa_diamond": "gpqa_diamond",
-    "humaneval": "humaneval",
-    "livecodebench": "livecodebench",
-    "gsm8k": "gsm8k",
-}
-
-
-def _map_results(raw_results: dict, tasks: list[str], task_map: dict[str, str]) -> dict:
+
+def _map_results(raw_results: dict, tasks: list[str]) -> dict:
   """Extract per-task accuracy metrics from lm-eval / evalchemy output."""
   scores: dict[str, float] = {}
   results_section = raw_results.get("results", {})
   for task in tasks:
-    lm_task = task_map.get(task, task)
-    task_r = results_section.get(lm_task, {})
+    task_r = results_section.get(task, {})
 
     acc = None
     for key in (
@@ -125,7 +99,6 @@ def run_harness(cfg: dict, hf_token: str | None = None) -> dict:
 
   Raises:
     ImportError: If lm_eval (or evalchemy for that backend) is not installed.
-    ValueError: If a requested task name is not in the backend's task map.
   """
   # pylint: disable=import-outside-toplevel
   try:
@@ -154,19 +127,8 @@ def run_harness(cfg: dict, hf_token: str | None = None) -> dict:
   gcs_results_path = cfg.get("gcs_results_path")
   token = resolve_token(cfg, hf_token)
 
-  task_map = _TASK_MAP_EVALCHEMY if backend == "evalchemy" else _TASK_MAP_LM_EVAL
   lm_model_type = "local-chat-completions" if backend == "evalchemy" else "local-completions"
 
-  lm_tasks: list[str] = []
-  for t in tasks:
-    lm_task = task_map.get(t)
-    if lm_task is None:
-      raise ValueError(
-          f"No {backend} task mapping for '{t}'. "
-          f"Known tasks: {list(task_map.keys())}"
-      )
-    lm_tasks.append(lm_task)
-
   with build_server_manager(cfg, token) as server:
     import jax as _jax
     from jax.experimental import multihost_utils as _multihost_utils
@@ -191,14 +153,14 @@ def run_harness(cfg: dict, hf_token: str | None = None) -> dict:
       logger.info(
           "Running %s tasks %s via %s at %s",
           backend,
-          lm_tasks,
+          tasks,
           lm_model_type,
           server.base_url,
       )
       raw_results = lm_eval_lib.simple_evaluate(
           model=lm_model_type,
           model_args=model_args,
-          tasks=lm_tasks,
+          tasks=tasks,
           num_fewshot=num_fewshot,
           limit=num_samples,
           log_samples=False,
@@ -214,7 +176,7 @@ def run_harness(cfg: dict, hf_token: str | None = None) -> dict:
   if not is_rank0:
     return {}
 
-  scores = _map_results(raw_results, tasks, task_map)
+  scores = _map_results(raw_results, tasks)
   logger.info("%s scores: %s", backend, scores)
 
   output = write_results(
@@ -253,9 +215,8 @@ def _build_arg_parser() -> argparse.ArgumentParser:
       nargs="+",
       default=["mmlu"],
       help=(
-          "Benchmark task names. "
-          "lm_eval choices: " + ", ".join(_TASK_MAP_LM_EVAL) + ". "
-          "evalchemy choices: " + ", ".join(_TASK_MAP_EVALCHEMY) + "."
+          "lm-eval task names passed directly to simple_evaluate. "
+          "Any task registered in lm-eval or evalchemy is accepted (e.g. gsm8k, mmlu, gpqa_diamond, ifeval, math_500)."
       ),
   )
   parser.add_argument(

src/maxtext/eval/runner/server_manager.py

@@ -210,6 +210,7 @@ def __init__(
       port: int = 8000,
       tensor_parallel_size: int = 4,
       expert_parallel_size: int = 1,
+      data_parallel_size: int = 1,
       max_model_len: int = 4096,
       dtype: str = "bfloat16",
       max_num_batched_tokens: int | None = None,
@@ -233,6 +234,7 @@ def __init__(
     self.port = port
     self.tensor_parallel_size = tensor_parallel_size
     self.expert_parallel_size = expert_parallel_size
+    self.data_parallel_size = data_parallel_size
     self.max_model_len = max_model_len
     self.dtype = dtype
     self.max_num_batched_tokens = max_num_batched_tokens
@@ -272,6 +274,7 @@ def start(self) -> None:
     vllm_kwargs: dict = {
         "model": self.model_path,
         "tensor_parallel_size": ici_tp,
+        "data_parallel_size": self.data_parallel_size,
         "max_model_len": self.max_model_len,
         "dtype": self.dtype,
         "gpu_memory_utilization": self.hbm_memory_utilization,
@@ -318,9 +321,10 @@ def start(self) -> None:
           vllm_kwargs[_k] = _v
 
     logger.info(
-        "Initializing in-process vLLM (tp=%d, ep=%d, max_len=%d)...",
+        "Initializing in-process vLLM (tp=%d, ep=%d, dp=%d, max_len=%d)...",
         ici_tp,
         ici_ep,
+        self.data_parallel_size,
         self.max_model_len,
     )
     self._llm = LLM(**vllm_kwargs)

src/maxtext/eval/scoring/rouge_scorer.py

@@ -18,6 +18,11 @@
 
 import numpy as np
 
+import nltk  # pylint: disable=import-outside-toplevel
+
+nltk.download("punkt", quiet=True)
+nltk.download("punkt_tab", quiet=True)
+
 
 def score_batch(
     responses: list[str],
@@ -43,10 +48,7 @@ def score_batch(
     )
 
   import evaluate  # pylint: disable=import-outside-toplevel
-  import nltk  # pylint: disable=import-outside-toplevel
 
-  nltk.download("punkt", quiet=True)
-  nltk.download("punkt_tab", quiet=True)
   metric = evaluate.load("rouge")
 
   preds = []