AI-Hypercomputer / maxtext

A simple, performant and scalable Jax LLM!
Apache License 2.0
1.54k stars 295 forks source link

Updating Jax Stable Stack BaseImage path #1050

Closed parambole closed 2 days ago

parambole commented 3 days ago

Updating path to point to the latest location (changed for better nomenclature)

Tested:

Ran the following command successfully.

bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxtext_gpu_jax_stable_stack_0.4.35 MODE=stable_stack DEVICE=gpu PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxtext_jax_stable_stack_0.4.35 BASEIMAGE=us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/gpu:jax0.4.35-cuda_dl24.10-rev1 MAXTEXT_REQUIREMENTS_FILE=requirements_with_jax_stable_stack.txt