[Tunix] Support DeepSWE evaluation pipeline with MaxText and vLLM on TPU by susanbao · Pull Request #2094 · google/tunix

susanbao · 2026-09-02T21:17:28Z

Resolves #

It's a good idea to open an issue first for discussion.

Description

This PR enhances the DeepSWE evaluation pipeline (examples/deepswe/eval_deepswe.py) to support both HuggingFace and MaxText models using vLLM on TPU clusters:

DeepSWE Evaluation Pipeline (examples/deepswe/eval_deepswe.py):
- Adds CLI argument parsing (--model_version, --model_source, --model_absolute_path, --mesh_fsdp, --mesh_tp, --scan_layers, etc.) with fallback to environment variables.
- Supports loading MaxText models (e.g. Qwen3.5-35B-A3B) directly into VllmSampler or via AutoModel.
- Integrates 3D MRoPE position normalization patch for Qwen3-Omni-Moe thinker models in MaxText.
- Dynamically configures Kubernetes nodeSelector via patch_kubernetes_runtime for GKE nodepools.
- Implements Pass@k metric computation (Pass@1, Pass@4, Pass@5) using an unbiased combinatorial estimator across instance groups.
- Adds direct GCS upload support for evaluation results using google-cloud-storage with gcloud storage cp fallback.
- Adds garbage collection (gc.collect()) after model calls for memory management.
AutoModel (tunix/models/automodel.py):
- Whitelists checkpoint_storage_concurrent_gb in AutoModel.from_pretrained kwargs for Orbax checkpoint restore configuration.
Trajectory Engine (tunix/rl/agentic/trajectory/trajectory_collect_engine.py):
- Safely extracts prompt tokens from either left_padded_prompt_tokens or padded_prompt_tokens on SamplerOutput to prevent prompt desync.

Reference

Colab Notebook

Checklist

I have added all the necessary unit tests for my change.
I have verified that my change does not break existing code and all unit tests pass.
I have added all appropriate doc-strings/documentation.
My PR is based on the latest changes of the main branch (if unsure, rebase the code).
I have signed the Contributor License Agreement.
I have followed Contribution Guidelines.

Note: Standard CPU unit tests, package builds, and documentation checks will run automatically on pull requests. Once the PR is approved and ready for submission, maintainers will add the ready-to-submit label to trigger full TPU testing.

- examples/deepswe/eval_deepswe.py: - Add CLI arguments for model, dataset, and runtime configuration - Add MaxText checkpoint loading support in vLLM sampler and AutoModel - Add patch for 3D MRoPE position normalization with Qwen thinker models - Configure Kubernetes nodeSelector for GKE nodepools - Compute pass@k metrics (pass@1, pass@4, pass@5) with unbiased estimator - Support direct GCS upload for evaluation results - tunix/models/automodel.py: - Whitelist checkpoint_storage_concurrent_gb in AutoModel kwargs - tunix/rl/agentic/trajectory/trajectory_collect_engine.py: - Safely extract prompt tokens from either left_padded_prompt_tokens or padded_prompt_tokens

RissyRan

Thanks for the change! A few comments.

RissyRan · 2026-09-02T21:58:40Z

+
+    def _patched_qwen_mrope_call(self, inputs, position, *args, **kwargs):
+      if position is not None:
+        # Case 1: vLLM (3, N, 1) -> (N, 1, 3)


Could you help point the source for cases to be handled for user readability? i.e. why those cases exist

# Case 1: vLLM (3, N, 1) -> (N, 1, 3) # Case 2: vLLM (3, N) -> (N, 1, 3) # Case 3: (N,) -> (N, 1, 3) # Case 4: (B, S) -> (B, S, 3) # Case 5: (B, S, 1) -> (B, S, 3)

I didn't see the else branch, is this expected?

How do we test them?

Thanks for pointing this out! I will extract this logic into a dedicated helper function with clear docstrings explaining each case:

Why Cases 1 & 2 exist (vLLM MRoPE format): For Qwen thinker/MRoPE models, vLLM provides input_positions with the 3 MRoPE coordinates (temporal, height, width) as the leading dimension: (3, num_tokens) or (3, num_tokens, 1). MaxText's inputs tensor has shape (num_tokens, 1, num_heads, head_dim). If not transposed, computing rotary embeddings yields (3, num_tokens, 1, head_dim), which causes JAX broadcasting failure (N, 1, 16, 64) * (3, N, 1, 64). Transposing to (N, 1, 3) aligns axis 0 to the token sequence dimension.

Why Cases 3–5 exist (Standard 1D/2D text positions): When text positions come in as (N,), (B, S), or (B, S, 1), MaxText asserts position.shape[-1] == 3. These cases duplicate 1D text positions across the 3 MRoPE coordinate axes.

Why there is no else branch: If position is already (..., 3) (e.g. (B, S, 3) or (N, 1, 3)), it is already in MaxText's native format and passes through to the original call untouched. I will add an explicit comment and fallback log to make this behavior obvious.

End-to-End Verification: In our cluster tests on TPU v5p-32 running Qwen3.5-35B-A3B and Qwen3-32B, this patch was exercised live across SWE-Bench tasks. Without it, vLLM crashes on the very first forward step with ValueError: Position IDs must be 2D (batch, seq) or 3D (batch, seq, 3) or broadcasting mismatch (8192, 1, 16, 64) vs (3, 8192, 1, 64).

…in eval_deepswe.py - Add str2bool helper for robust boolean flag parsing - Add --dataset_num_proc with automatic fallback to single-process on failure - Add --allow_split_physical_axes (defaults to False to avoid perf degradation) - Add --checkpoint_storage_concurrent_gb (defaults to 32) - Add --weight_dtype, --prefuse_moe_weights, --maxtext_attention, --enable_continue_decode

s-noghabi · 2026-09-03T00:42:51Z

+
+    _orig_qwen_mrope_call = Qwen3OmniMoeThinkerTextRotaryEmbedding.__call__
+
+    def _patched_qwen_mrope_call(self, inputs, position, *args, **kwargs):


can this be placed in a helper function, ideally not in this file? maybe maxtex_utils.py? it degrades the readability

s-noghabi · 2026-09-03T00:46:00Z

-logger.info("Kubernetes connection verified.")

-# ========================== Model ==========================
+def patch_kubernetes_runtime():


is this patch_kubernetes_runtime() a change that would be needed both in eval and main deepswe? should we move this a to helper outside eval? maybe r2e_gym_helper.py?

+1. Nitin had similar stuff here https://github.com/google/tunix/pull/2051/changes#diff-1a6299c4c9d1a03ec4a737371fb4cd43e6280c5a797d324ab5ebeb9df92b97ea

can you reuse?

s-noghabi · 2026-09-03T00:48:17Z

-    TP_SIZE,
-)
-
-if MODEL_VERSION == "Qwen/Qwen3-4B-Instruct-2507":


why remove these checks?

s-noghabi · 2026-09-03T00:48:37Z

-    MODEL_PATH, model_config, mesh, dtype=jnp.float32
+    total_mesh_devices,
 )
-sft_utils.show_hbm_usage()


ditto, why remove these logging and info?

s-noghabi · 2026-09-03T00:49:31Z


 logger.info("Creating sampler with engine=%s ...", ROLLOUT_ENGINE)

-if ROLLOUT_ENGINE == "vanilla":


why remove vanilla? we use this in our team

s-noghabi · 2026-09-03T00:58:33Z

-      self.agent.trajectory.prompt_tokens = (  # pyrefly: ignore[missing-attribute]
-          rollout_output.left_padded_prompt_tokens[0]
-      )
+    prompt_tokens = getattr(


why use getattr? won't rollout_output always have left_padded_prompt_tokens?

s-noghabi · 2026-09-03T01:00:20Z

please make sure you add tests for all changes added,, especially to tunix/ core directories.

s-noghabi · 2026-09-03T01:01:55Z

+        ):
+          position = jnp.transpose(position.squeeze(-1), (1, 0))[:, None, :]
+        # Case 2: vLLM (3, N) -> (N, 1, 3)
+        elif position.ndim == 2 and position.shape[0] == 3:


The check elif position.ndim == 2 and position.shape[0] == 3: assumes that any 2D position tensor with leading dimension 3 is a vLLM 3D coordinate tuple (3, N). If a standard 2D batch of shape (batch_size, seq_len) is evaluated with batch_size = 3, this branch incorrectly triggers, no?

s-noghabi · 2026-09-03T01:02:47Z

            echo=False,
            eos_tokens=qwen_eos_tokens,
        )
+    gc.collect()


calling gc collect synchronously inside model_call on every single generation step seems like a poor performance choice

s-noghabi · 2026-09-03T01:03:19Z

      mesh=mesh,
      hbm_utilization=VLLM_HBM_UTILIZATION,
-      init_with_random_weights=VLLM_INIT_RANDOM_WEIGHTS,
+      init_with_random_weights=False,


why false and not VLLM_INIT_RANDOM_WEIGHTS

tianshub · 2026-09-03T17:00:02Z

 #!/usr/bin/env python
 """DeepSWE evaluation with deepscaler-style task-level parallelism.

-This script intentionally does not modify the existing eval entrypoint.


why are these doc string removed?

tianshub · 2026-09-03T17:01:23Z

 import time

+# Path Setup before JAX
+workdir = os.getcwd()


why do we need this? if you have the venv with dependencies installed, all these should be available no?

tianshub · 2026-09-03T17:01:54Z

+  from guarded_swe_env import GuardedSWEEnv
+  from swe_agent import SWEAgent
+  from swe_env import SWEEnv
+except ImportError:


in what case the import won't be available?

tianshub · 2026-09-03T17:05:56Z

 os.environ["TOKENIZERS_PARALLELISM"] = "true"

+# Register MaxText vLLM adapter if using a MaxText model
+if MODEL_SOURCE == "maxtext":


please put this into a maxtext_model_helper file

tianshub · 2026-09-03T17:06:50Z

-    cache_dir=DATASET_CACHE,
-    num_proc=32,
-)
+if DATASET_NUM_PROC > 1:


why do we need this knob? do you see data loading being slow?

tianshub · 2026-09-03T17:08:20Z

+parser_cli.add_argument(
+    "--node_selector_val",
+    type=str,
+    default=os.getenv("NODE_SELECTOR_VAL", "cpu-np"),


why do we want to have two ways to config the CLI, one with py args one with env var, and default one to another? can we just reconcile to one mechanism?

tianshub · 2026-09-03T17:11:19Z

-logger.info("Kubernetes connection verified.")

-# ========================== Model ==========================
+def patch_kubernetes_runtime():


+1. Nitin had similar stuff here https://github.com/google/tunix/pull/2051/changes#diff-1a6299c4c9d1a03ec4a737371fb4cd43e6280c5a797d324ab5ebeb9df92b97ea

can you reuse?

tianshub · 2026-09-03T17:12:53Z

-TP_SIZE = 8
-mesh_devices = np.array(devices[:TP_SIZE]).reshape(1, TP_SIZE)
+total_mesh_devices = MESH_FSDP * MESH_TP
+if total_mesh_devices > len(devices):


a lot of these functions (similar for tokenizer) can just go to a util file. We should keep the recipe script as lean as possible.

susanbao requested review from abheesht17, hgao327, jiangyangmu, lc5211, s-noghabi, sizhit2, tianshub and wang2yn84 as code owners September 2, 2026 21:17

github-actions Bot assigned lc5211 Sep 2, 2026

RissyRan reviewed Sep 2, 2026

View reviewed changes

s-noghabi reviewed Sep 3, 2026

View reviewed changes

tianshub requested changes Sep 3, 2026

View reviewed changes


		_orig_qwen_mrope_call = Qwen3OmniMoeThinkerTextRotaryEmbedding.__call__

		def _patched_qwen_mrope_call(self, inputs, position, args, *kwargs):


		logger.info("Creating sampler with engine=%s ...", ROLLOUT_ENGINE)

		if ROLLOUT_ENGINE == "vanilla":

Uh oh!

Conversation

susanbao commented Sep 2, 2026

Description

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!

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!

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!

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!

5 participants