vLLM based Eval framework by dipannita08 · Pull Request #3531 · AI-Hypercomputer/maxtext

dipannita08 · 2026-03-31T18:42:10Z

Description

Implement a evaluation framework with vllm backend. Requirements, design, further details: go/eval-framework-vllm

The rest of the description includes relevant details and context, examples:

why is this change being made,
the problem being solved and any relevant context,
why this is a good solution,
some information about the specific implementation,
shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/508639301

Tests

Unit tests
E2E tests b/508639301

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

codecov · 2026-03-31T18:48:13Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

github-actions · 2026-03-31T21:19:03Z

🤖 Hi @Rohan-Bierneni, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions

## 📋 Review Summary

The implementation of the vLLM-based evaluation framework is a strong addition to MaxText, providing native support for custom benchmarks, lm-evaluation-harness, and evalchemy. The code is well-structured, but it needs critical updates to correctly support multi-host TPU environments whereリード (lead) rank coordination is essential.

🔍 General Feedback

Rank Coordination: In multi-host TPU setups, client-side operations (warmup, generation, reporting) must be restricted to jax.process_index() == 0 to avoid redundant work and failures on non-lead ranks.
Configurability: Key parameters like request timeouts should be made configurable via the CLI/config files rather than being hardcoded.
Efficiency: Minor optimizations in NLTK data handling and FastAPI request processing would improve the overall robustness and performance of the evaluation tool.

RissyRan

Thanks for the PR! I’ve done a high-level pass, though I haven’t done a deep dive into the code just yet. Have you had a chance to run any benchmarks? I'm curious if you're seeing decent scores. Also, is multi-host benchmarking for large models within the scope of this PR?

entrpn · 2026-05-01T20:42:46Z

@dipannita08 has this code been tested with an nnx checkpoint? @hengtaoguo recently added support for this and it is needed in order to support distilled checkpoints. Hengtao's PR: #3188

RissyRan

Thanks for the change!

Potentially we could have a run and see how it performance? This is Kurt's final presentation shows some period of time for common benchmarks using JetStream. It will be great to see if we could speed up using this tool.

RissyRan · 2026-05-05T01:56:01Z

@@ -0,0 +1,8 @@
+# Base evaluation configuration.


Shall we also include benchmark name or eval dataset in this base yml?

Base config covers server and generation parameters that are shared across all eval runs (temperature, concurrency, tensor_parallel_size, etc.). Benchmark name and dataset-specific settings (e.g. num_samples, max_tokens) live in task-specific configs like mlperf.yml so we can re-use the base config for different benchmarks without modification.

For the harness-based runners (lm_eval, evalchemy), benchmark/task selection is handled by the --tasks CLI arg (--tasks gsm8k mmlu gpqa) rather than config files (examples in the bug I shared above).

dipannita08 · 2026-05-05T22:26:52Z

@dipannita08 has this code been tested with an nnx checkpoint? @hengtaoguo recently added support for this and it is needed in order to support distilled checkpoints. Hengtao's PR: #3188

Yes, the checkpoint loading in model_creation_utils.py defaults to nnx and falls back to Linen if it detects the params.params double-nesting. MaxTextForCausalLM in the adapter is an nnx.Module. Some example runs with Qwen3-30B-A3B NNX checkpoints are in b/508639301.

dipannita08 · 2026-05-05T22:31:09Z

Potentially we could have a run and see how it performance? This is Kurt's final presentation shows some period of time for common benchmarks using JetStream. It will be great to see if we could speed up using this tool.

Please see E2E results in the bug: b/508639301

hengtaoguo

Thanks for the great work Dipannita! Do you plan to add a documentation for this route as a follow up?

hengtaoguo · 2026-05-06T19:45:20Z

+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": question},
+        ]
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


Out of curiosity, is it usually preferred to use tokenizer's chat_template if available?

It is preferred for any custom benchmarks. We defer all of the tokenizer chat templating to the harness if it already supports the benchmark in question.

hengtaoguo · 2026-05-06T19:50:47Z

+_DEFAULT_CONCURRENCY = 64
+_DEFAULT_MAX_TOKENS = 1024
+_DEFAULT_TEMPERATURE = 0.0
+_COMPLETIONS_PATH = "/v1/completions"


For new customized dataset in the future, I assume it needs new endpoints here?

For multi-modal, we will likely need to add an endpoint here, But no new endpoints are needed for custom text-based datasets. The /v1/completions endpoint accepts any pre-formatted prompt string, so a new dataset just needs to implement BenchmarkDataset.sample_requests() and return SampleRequest objects with the prompt already formatted. For chat-format datasets, the /v1/chat/completions endpoint is already implemented in server_manager.py and used by the lm-eval/evalchemy harness path. The custom dataset path uses /v1/completions because prompts are pre-formatted via the chat template in eval_runner.py before they're sent.

hengtaoguo · 2026-05-06T19:54:58Z

+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Async HTTP client for the /v1/completions endpoint.


Excuse my naive question: what does "completions" mean in this scenario? Is this async_client designed specific for completions task, or generalizable to other tasks too?

This is the OpenAI text completion API. It allows us to send a raw prompt string and the model returns the text continuation. The /v1/chat/completions, takes a structured messages array and applies chat templating server-side. We are using /v1/completions in the custom dataset path because prompt formatting (and chat templating) is done client-side in eval_runner.py before the request is sent, so each dataset has control over prompt structure. I assume we will mostly be using the industry standard datasets (eval_harness) for most cases though.

hengtaoguo · 2026-05-06T19:57:07Z

+  from maxtext.eval.runner.async_client import generate_batch
+  from maxtext.eval.runner.common import build_server_manager, maybe_upload_to_gcs, resolve_token
+  from maxtext.eval.runner.warmup import warmup_server
+  from maxtext.eval.scoring.registry import get_scorer


Are there any specific reasons not to import on the toplevel?

These imports are deferred because they pull in optional dependencies so deferring helps with importing and testing helper modules and methods without pulling in vllm deps etc.

NicoGrande · 2026-05-07T23:02:42Z

+      base_url = server.base_url
+
+      # Warmup server.
+      warmup_server(base_url=base_url, model=model_name, sample_requests=requests)


vLLM runs pre-compilation when it receives the first request by default. Is pre-compilation disabled in this framework?

Pre-compilation is not disabled, it is additionally implemented by warmup_server() before any evaluation requests are sent to help speed up the first vLLM request as a pre-vllm-warmup.

The SKIP_JAX_PRECOMPILE=1 environment variable set in server_manager.py only suppresses JAX's startup precompilation pass (which would try to compile shapes before the model config is known), this warmup step then covers exactly the shapes needed for the eval workload. This gives more targeted compilation than JAX's default precompile and avoids wasted compile time on shapes that won't be used.

Sounds good to me! Have you verified that you are not seeing any instances of compilation during generation?

I did check for single host, should be similar for multihost - did not check.

dipannita08 marked this pull request as ready for review March 31, 2026 20:30

dipannita08 requested review from A9isha, NicoGrande, NuojCheng, RissyRan, SurbhiJainUSC, aireenmei, bvandermoon, gagika, gobbleturk, hengtaoguo, igorts-git, jesselu-google, jiangjy1982, khatwanimohit, parambole, richjames0, shralex, suexu1025 and vipannalla as code owners March 31, 2026 20:30

Rohan-Bierneni added the gemini-review label Mar 31, 2026

github-actions Bot reviewed Mar 31, 2026

View reviewed changes

RissyRan reviewed Apr 1, 2026

View reviewed changes

Comment thread src/maxtext/eval/README.md Outdated

Comment thread src/maxtext/eval/README.md Outdated

Comment thread src/maxtext/eval/README.md Outdated

Comment thread src/maxtext/eval/runner/lm_eval_runner.py Outdated

dipannita08 requested a review from abhinavclemson as a code owner April 10, 2026 19:48

dipannita08 force-pushed the eval-framework-01 branch 2 times, most recently from cb0295b to 348c539 Compare April 23, 2026 01:12

vLLM based Eval Framework

03ce83a

dipannita08 force-pushed the eval-framework-01 branch from c2d7d04 to 03ce83a Compare May 1, 2026 17:39

dipannita08 requested a review from RissyRan May 1, 2026 18:13

entrpn reviewed May 1, 2026

View reviewed changes

Comment thread src/maxtext/eval/datasets/base.py Outdated

RissyRan reviewed May 5, 2026

View reviewed changes

hengtaoguo approved these changes May 6, 2026

View reviewed changes

NicoGrande reviewed May 7, 2026

View reviewed changes

Add tokenizer type annotation

351da27

dipannita08 requested review from NicoGrande, RissyRan, entrpn and hengtaoguo May 11, 2026 18:55

hengtaoguo approved these changes May 12, 2026

View reviewed changes

NicoGrande approved these changes May 12, 2026

View reviewed changes

Conversation

dipannita08 commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

codecov Bot commented Mar 31, 2026

Codecov Report

Uh oh!

github-actions Bot commented Mar 31, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

entrpn commented May 1, 2026

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dipannita08 commented May 5, 2026

Uh oh!

dipannita08 commented May 5, 2026

Uh oh!

hengtaoguo left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

dipannita08 commented Mar 31, 2026 •

edited

Loading