diff --git a/src/MaxText/examples/grpo_llama3_demo.py b/src/MaxText/examples/grpo_llama3_demo.py index ee5997b902..a1e070ef76 100644 --- a/src/MaxText/examples/grpo_llama3_demo.py +++ b/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", diff --git a/src/MaxText/examples/install_tunix_vllm_requirement.sh b/src/MaxText/examples/install_tunix_vllm_requirement.sh index 0a89f96675..19a6205455 100644 --- a/src/MaxText/examples/install_tunix_vllm_requirement.sh +++ b/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 \