|
20 | 20 |
|
21 | 21 | # For instructions on building the MaxText Docker image, please refer to the https://maxtext.readthedocs.io/en/latest/build_maxtext.html. |
22 | 22 |
|
23 | | -PACKAGE_DIR="${PACKAGE_DIR:-src}" |
24 | | -echo "PACKAGE_DIR: $PACKAGE_DIR" |
25 | | -TESTS_DIR="${TESTS_DIR:-tests}" |
26 | | -echo "TESTS_DIR: $TESTS_DIR" |
| 23 | +# Build docker image with stable dependencies |
| 24 | +## bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE={{gpu|tpu}} MODE=stable |
| 25 | + |
| 26 | +# Build docker image with nightly dependencies |
| 27 | +## bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE={{gpu|tpu}} MODE=nightly |
| 28 | + |
| 29 | +# Build docker image with stable dependencies and, a pinned JAX_VERSION for TPUs |
| 30 | +## bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.4.13 |
| 31 | + |
| 32 | +# Build docker image with a pinned JAX_VERSION and, a pinned LIBTPU_VERSION for TPUs |
| 33 | +## bash src/dependencies/scripts/docker_build_dependency_image.sh MODE={{stable|nightly}} JAX_VERSION=0.8.1 LIBTPU_VERSION=0.0.31.dev20251119+nightly |
| 34 | + |
| 35 | +# Build docker image with a custom libtpu.so for TPUs |
| 36 | +# Note: libtpu.so file must be present in the root directory of the MaxText repository |
| 37 | +## bash src/dependencies/scripts/docker_build_dependency_image.sh MODE={{stable|nightly}} |
| 38 | + |
| 39 | +# Build docker image with nightly dependencies and, a pinned JAX_VERSION for GPUs |
| 40 | +# Available versions listed at https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jax |
| 41 | +## bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=gpu MODE=nightly JAX_VERSION=0.4.36.dev20241109 |
| 42 | + |
| 43 | +# ================================== |
| 44 | +# POST-TRAINING BUILD EXAMPLES |
| 45 | +# ================================== |
| 46 | + |
| 47 | +# Build docker image with post-training dependencies |
| 48 | +## bash src/dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-training |
| 49 | + |
| 50 | +if [ "${BASH_SOURCE-}" ]; then |
| 51 | + this_file="${BASH_SOURCE[0]}" |
| 52 | +elif [ "${ZSH_VERSION-}" ]; then |
| 53 | + # shellcheck disable=SC2296 |
| 54 | + this_file="${(%):-%x}" |
| 55 | +else |
| 56 | + this_file="${0}" |
| 57 | +fi |
| 58 | + |
| 59 | +MAXTEXT_REPO_ROOT="${MAXTEXT_REPO_ROOT:-$(CDPATH='' cd -- "$(dirname -- "${this_file}")"'/../../..' && pwd)}" |
27 | 60 |
|
28 | 61 | # Enable "exit immediately if any command fails" option |
29 | 62 | set -e |
|
0 commit comments