Merge pull request #2472 from AI-Hypercomputer:vllm_fix

PiperOrigin-RevId: 817239502

Author: Google-ML-Automation

src/MaxText/examples/grpo_llama3_demo.py

@@ -458,7 +458,7 @@ def get_ref_maxtext_model(config):
 config_policy = pyconfig.initialize(
     [
         "",
-        os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml"),
+        f"{HOME}/maxtext/src/MaxText/configs/base.yml",
     ],
     base_output_directory="dummy",  # This is not used in Tunix.
     run_name="test-tunix-maxtext-llama3.1-8b",  # This is not used in Tunix.

src/MaxText/examples/install_tunix_vllm_requirement.sh

@@ -39,14 +39,14 @@ VLLM_TARGET_DEVICE="tpu" pip install --no-cache-dir --pre \
     --find-links https://storage.googleapis.com/libtpu-releases/index.html \
     --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
     --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html \
-    vllm==0.10.2rc2.dev59+gdcb28a332.tpu
+    vllm==0.11.1rc1.dev292+g1b86bd8e1.tpu
 
 # Install tpu-commons from the artifact registry
 pip install --no-cache-dir --pre \
     --index-url https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/ \
     --extra-index-url https://pypi.org/simple/ \
     --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
     --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
-    tpu-commons==0.1.1
+    tpu-commons==0.1.2
 
 pip install numba==0.61.2