From f116054dbed654cbb280764457fa9a78fe003b51 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Mon, 9 Sep 2024 10:24:07 -0700 Subject: [PATCH] TE multithread build (#1009) --- .github/container/Dockerfile.jax | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/container/Dockerfile.jax b/.github/container/Dockerfile.jax index c85bee347..726656a7a 100644 --- a/.github/container/Dockerfile.jax +++ b/.github/container/Dockerfile.jax @@ -62,6 +62,7 @@ pip install ninja && rm -rf ~/.cache/pip # TransformerEngine now needs JAX at build time git-clone.sh ${URLREF_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE} pushd ${SRC_PATH_TRANSFORMER_ENGINE} +export NVTE_BUILD_THREADS_PER_JOB=8 python setup.py bdist_wheel && rm -rf build ls "${SRC_PATH_TRANSFORMER_ENGINE}/dist" EOF