Add YaRN RoPE configuration support and enhance evaluation parameters

src/maxtext/checkpoint_conversion/to_huggingface.py

@@ -275,6 +275,32 @@ def _validate_or_update_architecture(hf_config, max_config, override: bool):
     raise ValueError(error_msg)
 
 
+def _apply_yarn_rope_config(hf_config, max_config) -> None:
+  """Apply MaxText YaRN RoPE values without dropping HF model-specific fields."""
+  rope_scaling = getattr(hf_config, "rope_scaling", None)
+  if isinstance(rope_scaling, dict):
+    rope_scaling = dict(rope_scaling)
+  else:
+    rope_scaling = {}
+
+  rope_type_key = "type" if "type" in rope_scaling and "rope_type" not in rope_scaling else "rope_type"
+  rope_scaling.update(
+      {
+          "beta_fast": float(getattr(max_config, "beta_fast", 32.0)),
+          "beta_slow": float(getattr(max_config, "beta_slow", 1.0)),
+          "factor": float(getattr(max_config, "rope_factor", 32.0)),
+          "original_max_position_embeddings": getattr(max_config, "original_max_position_embeddings", 4096),
+          "rope_theta": getattr(max_config, "rope_max_timescale"),
+          rope_type_key: "yarn",
+          "truncate": getattr(max_config, "rope_truncate", False),
+      }
+  )
+
+  hf_config.rope_theta = getattr(max_config, "rope_max_timescale")
+  hf_config.rope_scaling = rope_scaling
+  hf_config.rope_parameters = rope_scaling
+
+
 def _transform_weights_to_adapter(param_map, state_dict):
   """Extracts standalone PEFT weights from MaxText state dict."""
   processed_params_list = []
@@ -404,6 +430,10 @@ def main(argv: Sequence[str]) -> None:
   # Validate architecture consistency (raising ValueError on mismatch) or override HF config if specified.
   _validate_or_update_architecture(hf_config_obj, config, override=FLAGS.override_model_architecture)
 
+  # Ensure YaRN rope params are preserved if specified in the maxtext config
+  if getattr(config, "rope_type", None) == "yarn":
+    _apply_yarn_rope_config(hf_config_obj, config)
+
   # 2. Load Tokenizer
   if model_key not in HF_IDS:
     raise ValueError(f"HF Tokenizer ID not found for model key: {model_key}")

src/maxtext/eval/runner/harness_runner.py

@@ -123,6 +123,9 @@ def run_harness(cfg: dict, hf_token: str | None = None) -> dict:
   num_fewshot = cfg.get("num_fewshot", 0)
   num_samples = cfg.get("num_samples")
   gcs_results_path = cfg.get("gcs_results_path")
+  apply_chat_template = cfg.get("apply_chat_template", False)
+  fewshot_as_multiturn = cfg.get("fewshot_as_multiturn", False)
+  gen_kwargs = cfg.get("gen_kwargs")
   token = resolve_token(cfg, hf_token)
 
   lm_model_type = "local-chat-completions" if backend == "evalchemy" else "local-completions"
@@ -156,14 +159,21 @@ def run_harness(cfg: dict, hf_token: str | None = None) -> dict:
           lm_model_type,
           server.base_url,
       )
-      raw_results = lm_eval_lib.simple_evaluate(
-          model=lm_model_type,
-          model_args=model_args,
-          tasks=tasks,
-          num_fewshot=num_fewshot,
-          limit=num_samples,
-          log_samples=False,
-      )
+      simple_eval_kwargs: dict = {
+          "model": lm_model_type,
+          "model_args": model_args,
+          "tasks": tasks,
+          "num_fewshot": num_fewshot,
+          "limit": num_samples,
+          "log_samples": False,
+      }
+      if apply_chat_template:
+        simple_eval_kwargs["apply_chat_template"] = True
+      if fewshot_as_multiturn:
+        simple_eval_kwargs["fewshot_as_multiturn"] = True
+      if gen_kwargs:
+        simple_eval_kwargs["gen_kwargs"] = gen_kwargs
+      raw_results = lm_eval_lib.simple_evaluate(**simple_eval_kwargs)
 
     # All ranks block here until rank-0 finishes evaluation. Non-rank-0 hosts
     # keep their in-process LLM alive so rank-0's llm.generate() calls can
@@ -221,6 +231,39 @@ def _build_arg_parser() -> argparse.ArgumentParser:
   )
   parser.add_argument("--num_fewshot", type=int, default=0, help="Few-shot examples per task.")
   parser.add_argument("--num_samples", type=int, help="Limit samples per task (None = full dataset).")
+  parser.add_argument(
+      "--apply_chat_template",
+      action="store_true",
+      default=False,
+      help=(
+          "Wrap each prompt in the tokenizer's chat template before evaluation. "
+          "Required for instruction-tuned / chat models (e.g. GPT-OSS harmony format). "
+          "Without this, MMLU/GSM8K prompts arrive as bare text the model was never "
+          "trained on -> near-random scores."
+      ),
+  )
+  parser.add_argument(
+      "--fewshot_as_multiturn",
+      action="store_true",
+      default=False,
+      help=(
+          "When --apply_chat_template and --num_fewshot > 0: format each few-shot "
+          "example as a separate user/assistant chat turn instead of one big context. "
+          "Recommended for chat models with few-shot MMLU."
+      ),
+  )
+  parser.add_argument(
+      "--gen_kwargs",
+      type=str,
+      default=None,
+      help=(
+          "Comma-separated key=value overrides for generation tasks (generate_until). "
+          "Passed directly to lm-eval simple_evaluate(gen_kwargs=...). "
+          "No effect on loglikelihood tasks (e.g. MMLU). "
+          "Example for GPT-OSS GSM8K CoT: 'until=[],max_gen_toks=1024' "
+          "(clears default stop sequences and raises the token budget)."
+      ),
+  )
   return parser
 
 

src/maxtext/eval/runner/server_manager.py

@@ -27,9 +27,16 @@
 
 logger = logging.getLogger(__name__)
 
+
 _HEALTH_ENDPOINT = "/health"
 
 
+def _top_logprobs_dict(tokenizer: Any, lp_dict: dict | None) -> "dict[str, float] | None":
+  if not lp_dict:
+    return None
+  return {tokenizer.decode([tid]): lp.logprob for tid, lp in lp_dict.items()}
+
+
 def _build_app(llm: Any) -> Any:
   """Return a FastAPI app that wraps an in-process vLLM LLM instance."""
   import fastapi  # pylint: disable=import-outside-toplevel
@@ -48,9 +55,12 @@ async def completions(request: fastapi.Request):
     body = await request.json()
 
     raw_prompt = body.get("prompt", "")
-    prompts = raw_prompt if isinstance(raw_prompt, list) else [raw_prompt]
+    if isinstance(raw_prompt, list) and raw_prompt and isinstance(raw_prompt[0], int):
+      prompts = [raw_prompt]
+    else:
+      prompts = raw_prompt if isinstance(raw_prompt, list) else [raw_prompt]
     model_name = body.get("model", "")
-    max_tokens = int(body.get("max_tokens") or 256)
+    max_tokens = int(body["max_tokens"]) if body.get("max_tokens") is not None else 256
     temperature = float(body.get("temperature") or 0.0)
     logprobs_n = body.get("logprobs")  # int | None
     echo = bool(body.get("echo", False))
@@ -80,6 +90,7 @@ async def completions(request: fastapi.Request):
       if logprobs_n is not None:
         tok_strings: list[str] = []
         tok_lps: list[float | None] = []
+        tok_tops: list[dict[str, float] | None] = []
         tok_offsets: list[int] = []
         running_offset = 0
 
@@ -93,6 +104,7 @@ async def completions(request: fastapi.Request):
             lp_dict = prompt_lps[pos] if pos < len(prompt_lps) else None
             lp_val = lp_dict[tok_id].logprob if (lp_dict and tok_id in lp_dict) else None
             tok_lps.append(lp_val)
+            tok_tops.append(_top_logprobs_dict(tokenizer, lp_dict))
 
         gen_lps = gen.logprobs or []
         for pos, tok_id in enumerate(gen.token_ids):
@@ -103,15 +115,20 @@ async def completions(request: fastapi.Request):
           lp_dict = gen_lps[pos] if pos < len(gen_lps) else None
           lp_val = lp_dict[tok_id].logprob if (lp_dict and tok_id in lp_dict) else None
           tok_lps.append(lp_val)
+          tok_tops.append(_top_logprobs_dict(tokenizer, lp_dict))
 
         logprobs_payload = {
             "tokens": tok_strings,
             "token_logprobs": tok_lps,
-            "top_logprobs": None,
+            "top_logprobs": tok_tops,
             "text_offset": tok_offsets,
         }
 
-      text_out = (prompts[idx] + gen.text) if echo else gen.text
+      if echo:
+        p = prompts[idx]
+        text_out = (tokenizer.decode(p) if isinstance(p, list) else p) + gen.text
+      else:
+        text_out = gen.text
       choices.append(
           {
               "text": text_out,
@@ -143,7 +160,7 @@ async def chat_completions(request: fastapi.Request):  # pylint: disable=unused-
     body = await request.json()
     messages = body.get("messages", [])
     model_name = body.get("model", "")
-    max_tokens = int(body.get("max_tokens") or 256)
+    max_tokens = int(body["max_tokens"]) if body.get("max_tokens") is not None else 256
     temperature = float(body.get("temperature") or 0.0)
     stop = body.get("stop")
 

tests/unit/eval/test_build_app.py

@@ -223,6 +223,57 @@ def test_chat_completions_applies_template(self):
     passed_messages = call_args[0][0] if call_args[0] else call_args[1].get("conversation")
     self.assertIsNotNone(passed_messages)
 
+  def test_completions_token_id_prompt_normalised_as_single(self):
+    """A list-of-ints prompt must be treated as one token-ID prompt, not a batch of ints."""
+    mock_output = _make_mock_output(generated_text="out", prompt_token_ids=[1, 2, 3], generated_token_ids=[7])
+    self.mock_llm.generate.return_value = [mock_output]
+
+    resp = self.client.post(
+        "/v1/completions",
+        json={"model": "m", "prompt": [1, 2, 3], "max_tokens": 5},
+    )
+    self.assertEqual(resp.status_code, 200)
+    data = resp.json()
+    self.assertEqual(len(data["choices"]), 1)
+    self.assertEqual(data["choices"][0]["text"], "out")
+    # generate() must be called with a single-element list containing the token-ID list
+    call_prompts = self.mock_llm.generate.call_args[0][0]
+    self.assertEqual(call_prompts, [[1, 2, 3]])
+
+  def test_completions_top_logprobs_populated(self):
+    """When logprobs is requested, top_logprobs entries must be non-None dicts."""
+    mock_output = _make_mock_output(generated_text="hi", prompt_token_ids=[1], generated_token_ids=[4, 5])
+    mock_output.outputs[0].logprobs = [
+        {4: SimpleNamespace(logprob=-0.5), 6: SimpleNamespace(logprob=-1.0)},
+        {5: SimpleNamespace(logprob=-0.8)},
+    ]
+    self.mock_llm.generate.return_value = [mock_output]
+
+    resp = self.client.post(
+        "/v1/completions",
+        json={"model": "m", "prompt": "x", "max_tokens": 5, "logprobs": 2},
+    )
+    self.assertEqual(resp.status_code, 200)
+    lp = resp.json()["choices"][0]["logprobs"]
+    self.assertIsNotNone(lp)
+    self.assertIsInstance(lp["top_logprobs"][0], dict)
+    self.assertGreater(len(lp["top_logprobs"][0]), 0)
+
+  def test_completions_echo_with_token_id_prompt(self):
+    """echo=True with a token-ID prompt must decode the prompt and prepend it to text."""
+    mock_output = _make_mock_output(generated_text=" world", prompt_token_ids=[1, 2, 3], generated_token_ids=[4])
+    self.mock_llm.generate.return_value = [mock_output]
+
+    resp = self.client.post(
+        "/v1/completions",
+        json={"model": "m", "prompt": [1, 2, 3], "max_tokens": 5, "echo": True},
+    )
+    self.assertEqual(resp.status_code, 200)
+    text = resp.json()["choices"][0]["text"]
+    # tokenizer.decode([1, 2, 3]) → "tok1tok2tok3" (see _make_mock_llm side_effect)
+    self.assertTrue(text.startswith("tok1tok2tok3"), f"Expected decoded prompt prefix, got: {text!r}")
+    self.assertIn(" world", text)
+
 
 if __name__ == "__main__":
   unittest.main()

tests/unit/eval/test_lm_eval_runner.py

@@ -94,5 +94,52 @@ def test_empty_tasks_list(self):
     self.assertEqual(scores, {})
 
 
+_REQUIRED_ARGS = [
+    "--model_name",
+    "llama3.1-8b",
+    "--hf_path",
+    "meta-llama/Llama-3.1-8B",
+    "--base_output_directory",
+    "gs://bucket/",
+    "--run_name",
+    "test_run",
+    "--max_model_len",
+    "4096",
+]
+
+
+class TestArgParser(unittest.TestCase):
+  """Tests for _build_arg_parser — new flags added in the vllm_eval commits."""
+
+  def setUp(self):
+    from maxtext.eval.runner.harness_runner import _build_arg_parser  # pylint: disable=import-outside-toplevel
+
+    self.parser = _build_arg_parser()
+
+  def test_apply_chat_template_default_false(self):
+    args = self.parser.parse_args(_REQUIRED_ARGS)
+    self.assertFalse(args.apply_chat_template)
+
+  def test_apply_chat_template_flag_sets_true(self):
+    args = self.parser.parse_args(_REQUIRED_ARGS + ["--apply_chat_template"])
+    self.assertTrue(args.apply_chat_template)
+
+  def test_fewshot_as_multiturn_default_false(self):
+    args = self.parser.parse_args(_REQUIRED_ARGS)
+    self.assertFalse(args.fewshot_as_multiturn)
+
+  def test_fewshot_as_multiturn_flag_sets_true(self):
+    args = self.parser.parse_args(_REQUIRED_ARGS + ["--fewshot_as_multiturn"])
+    self.assertTrue(args.fewshot_as_multiturn)
+
+  def test_gen_kwargs_default_none(self):
+    args = self.parser.parse_args(_REQUIRED_ARGS)
+    self.assertIsNone(args.gen_kwargs)
+
+  def test_gen_kwargs_passthrough(self):
+    args = self.parser.parse_args(_REQUIRED_ARGS + ["--gen_kwargs", "until=[],max_gen_toks=1024"])
+    self.assertEqual(args.gen_kwargs, "until=[],max_gen_toks=1024")
+
+
 if __name__ == "__main__":
   unittest.main()

tests/unit/hf_checkpoint_conversion_test.py

@@ -12,14 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-""" Tests for kernels """
+"""Tests for kernels"""
 
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 import numpy as np
 from maxtext.utils.max_utils import permute_to_match_maxtext_rope, unpermute_from_match_maxtext_rope
 from maxtext.checkpoint_conversion import to_huggingface as to_hf
 from maxtext.checkpoint_conversion.to_huggingface import (
+    _apply_yarn_rope_config,
     _get_lora_delta,
     _transform_weights_to_adapter,
     _transform_weights_to_full_model,
@@ -61,6 +63,37 @@ def test_huggingface_to_maxtext_back_to_huggingface_flow(self):
     if not np.array_equal(wq2, wq4):
       print("Test failed: wq2 does not match wq4")
 
+  def test_apply_yarn_rope_config_preserves_existing_hf_fields(self):
+    hf_config = SimpleNamespace(
+        rope_theta=10_000,
+        rope_scaling={
+            "beta_fast": 32.0,
+            "beta_slow": 1.0,
+            "factor": 40.0,
+            "mscale": 0.707,
+            "mscale_all_dim": 0.707,
+            "original_max_position_embeddings": 4096,
+            "rope_theta": 10_000,
+            "type": "yarn",
+        },
+    )
+    max_config = SimpleNamespace(
+        beta_fast=32,
+        beta_slow=1,
+        rope_factor=40,
+        original_max_position_embeddings=4096,
+        rope_max_timescale=10_000,
+        rope_truncate=True,
+    )
+
+    _apply_yarn_rope_config(hf_config, max_config)
+
+    self.assertEqual(hf_config.rope_theta, 10_000)
+    self.assertEqual(hf_config.rope_scaling["type"], "yarn")
+    self.assertEqual(hf_config.rope_scaling["mscale"], 0.707)
+    self.assertEqual(hf_config.rope_scaling["mscale_all_dim"], 0.707)
+    self.assertTrue(hf_config.rope_scaling["truncate"])
+
 
 class MaxTextToHFLoRAConversionTest(unittest.TestCase):
   """Tests the conversion modes (Base, Adapter, Merged) in to_huggingface with LoRA support."""