Use shell script to drive jaxlib build #909

yhtang · 2023-03-30T19:01:07Z

Together with jax-ml/jax#15205, this should eliminate the need to apply the build_jaxlib/update_build_scripts.patch each time we rebase. Note the --editable option will only be usable once alpa is rebased onto JAX 0.4.8.

The original update_build_scripts.patch file is meant to provide two functionalities:

build jaxlib as an 'editable' package. This can now be done by

cd build
./build.sh  # build wheel
./build.sh --editable  # build editable installation

Ask Bazel to prefer local installations NCCL. This can now be done by

./build.sh --nccl local  # prefer local NCCL
./build.sh --nccl download  # prefer to use XLA's pinned version

gjoliver

nice cleanup.
let's figure out when we should merge this.

gjoliver · 2023-03-31T16:54:36Z

build_jaxlib/build.sh

+
+if [[ "${EDITABLE}" == "1" ]]; then
+    BUILD_PARAM="${BUILD_PARAM} --editable"
+fi


why don't we do this directly in the flag parsing while loop above?
you then don't need the additional EDITABLE variable.

gjoliver · 2023-03-31T16:55:02Z

build_jaxlib/build.sh

+
+if [[ "$DEBUG" == "1" ]]; then
+    BUILD_PARAM="${BUILD_PARAM} --bazel_options=-c --bazel_options=dbg --bazel_options=--strip=never --bazel_options=--cxxopt=-g --bazel_options=--cxxopt=-O0"
+fi


same for --debug I guess ...

gjoliver · 2023-03-31T16:55:37Z

build_jaxlib/build.sh

+print_var TF_NCCL_VERSION
+print_var CC_OPT_FLAGS
+
+print_var BUILD_PARAM


BUILD_PARAM is printed twice. get rid of the one above I guess?

gjoliver · 2023-03-31T16:56:48Z

build_jaxlib/local_cuda_arch

@@ -0,0 +1,30 @@
+//usr/bin/env nvcc -odir `mktemp -d` -x cu --run "$0" ${1:+--run-args "${@:1}"} ; exit $?


is this a new util?

use Shell script to drive jaxlib build

c093723

yhtang changed the title ~~Use Shell script to drive jaxlib build~~ Use shell script to drive jaxlib build Mar 30, 2023

gjoliver reviewed Mar 31, 2023

View reviewed changes

merrymercy force-pushed the main branch from 21232f0 to b8078a9 Compare December 9, 2023 16:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use shell script to drive jaxlib build #909

Use shell script to drive jaxlib build #909

yhtang commented Mar 30, 2023

gjoliver left a comment

gjoliver Mar 31, 2023

gjoliver Mar 31, 2023

gjoliver Mar 31, 2023

gjoliver Mar 31, 2023

		@@ -0,0 +1,30 @@
		//usr/bin/env nvcc -odir `mktemp -d` -x cu --run "$0" ${1:+--run-args "${@:1}"} ; exit $?

Use shell script to drive jaxlib build #909

Are you sure you want to change the base?

Use shell script to drive jaxlib build #909

Conversation

yhtang commented Mar 30, 2023

gjoliver left a comment

Choose a reason for hiding this comment

gjoliver Mar 31, 2023

Choose a reason for hiding this comment

gjoliver Mar 31, 2023

Choose a reason for hiding this comment

gjoliver Mar 31, 2023

Choose a reason for hiding this comment

gjoliver Mar 31, 2023

Choose a reason for hiding this comment