[DRAFT] GRPO integration test

src/MaxText/examples/grpo_llama3_demo.py

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

src/MaxText/examples/install_tunix_vllm_requirement.sh

@@ -29,17 +29,9 @@ pip install aiohttp==3.12.15
 pip install keyring keyrings.google-artifactregistry-auth
 
 # Install vLLM for Jax and TPUs from the artifact registry
-VLLM_TARGET_DEVICE="tpu" 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/ \
-    --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-    --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
-    --find-links https://storage.googleapis.com/libtpu-wheels/index.html \
-    --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
+git clone https://github.com/vllm-project/vllm.git /tmp/vllm
+cd /tmp/vllm
+VLLM_TARGET_DEVICE="tpu" pip install -e .
 
 # Install tpu-commons from the artifact registry
 pip install --no-cache-dir --pre \