Skip to content

GitLab

Unverified Commit 1f2df735 authored by jberchtold-nvidia's avatar jberchtold-nvidia Committed by GitHub
Browse files

Fix JAX and PyTorch wheel builds for v2.6 (#2005) 



* Fix L0_jax_wheel
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>

* Update
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>

* remove commented line
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>

* Reduce usage of --no-deps
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>

* Also fix pytorch wheel build
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>

* Revert test_sanity_import.py changes as it is also used on CPU-only GitHub build jobs
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>

---------
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>
parent 13cae89e
Show whitespace changes
Inline Side-by-side
Showing with 13 additions and 12 deletions
qa/L0_jax_wheel/test.sh
View file @ 1f2df735
...@@ -26,23 +26,24 @@ pip3 uninstall -y transformer-engine transformer-engine-cu12 transformer-engine- ...@@ -26,23 +26,24 @@ pip3 uninstall -y transformer-engine transformer-engine-cu12 transformer-engine-
VERSION=`cat $TE_PATH/build_tools/VERSION.txt` VERSION=`cat $TE_PATH/build_tools/VERSION.txt`
WHL_BASE="transformer_engine-${VERSION}" WHL_BASE="transformer_engine-${VERSION}"
# Core wheel. # Core wheel.
NVTE_RELEASE_BUILD=1 python3 setup.py bdist_wheel || error_exit "Failed to setup bdist_wheel" NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation -vvv --wheel-dir ./dist . || error_exit "Failed to setup bdist_wheel"
wheel unpack dist/* || error_exit "Failed to unpack dist/*" wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl"
sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info" mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info"
wheel pack ${WHL_BASE} || error_exit "Failed to pack ${WHL_BASE}" wheel pack ${WHL_BASE} || error_exit "Failed to pack ${WHL_BASE}"
rm dist/*.whl || error_exit "Failed to remove dist/*.whl" rm dist/*.whl || error_exit "Failed to remove dist/*.whl"
mv *.whl dist/ || error_exit "Failed to move *.whl to dist/" mv *.whl dist/ || error_exit "Failed to move *.whl to dist/"
NVTE_RELEASE_BUILD=1 NVTE_BUILD_METAPACKAGE=1 python3 setup.py bdist_wheel || error_exit "Failed to setup metapackage" NVTE_RELEASE_BUILD=1 NVTE_BUILD_METAPACKAGE=1 pip3 wheel --no-build-isolation --no-deps -vvv --wheel-dir ./dist . || error_exit "Failed to setup metapackage"
cd transformer_engine/jax cd transformer_engine/jax
NVTE_RELEASE_BUILD=1 python3 setup.py sdist || error_exit "Failed to setup sdist" NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation --no-deps -vvv --wheel-dir ./dist . || error_exit "Failed to setup sdist"
pip3 install dist/* || error_exit "Failed to install dist/*" pip3 install --no-build-isolation --no-deps -vvv dist/* || error_exit "Failed to install dist/*"
cd $TE_PATH cd $TE_PATH
pip3 install dist/*.whl --no-deps || error_exit "Failed to install dist/*.whl --no-deps" pip3 install --no-build-isolation --no-deps -vvv dist/*.whl || error_exit "Failed to install dist/*.whl --no-deps"
python3 $TE_PATH/tests/jax/test_sanity_import.py || test_fail "test_sanity_import.py" python3 $TE_PATH/tests/jax/test_sanity_import.py || test_fail "test_sanity_import.py"
......
qa/L0_pytorch_wheel/test.sh
View file @ 1f2df735
...@@ -27,22 +27,22 @@ VERSION=`cat $TE_PATH/build_tools/VERSION.txt` ...@@ -27,22 +27,22 @@ VERSION=`cat $TE_PATH/build_tools/VERSION.txt`
WHL_BASE="transformer_engine-${VERSION}" WHL_BASE="transformer_engine-${VERSION}"
# Core wheel. # Core wheel.
NVTE_RELEASE_BUILD=1 python3 setup.py bdist_wheel || error_exit "Failed to setup bdist_wheel" NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation -vvv --wheel-dir ./dist . || error_exit "Failed to setup bdist_wheel"
wheel unpack dist/* || error_exit "Failed to unpack dist/*" wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl"
sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA"
mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info" mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info"
wheel pack ${WHL_BASE} || error_exit "Failed to pack ${WHL_BASE}" wheel pack ${WHL_BASE} || error_exit "Failed to pack ${WHL_BASE}"
rm dist/*.whl || error_exit "Failed to remove dist/*.whl" rm dist/*.whl || error_exit "Failed to remove dist/*.whl"
mv *.whl dist/ || error_exit "Failed to move *.whl to dist/" mv *.whl dist/ || error_exit "Failed to move *.whl to dist/"
NVTE_RELEASE_BUILD=1 NVTE_BUILD_METAPACKAGE=1 python3 setup.py bdist_wheel || error_exit "Failed to setup metapackage" NVTE_RELEASE_BUILD=1 NVTE_BUILD_METAPACKAGE=1 pip3 wheel --no-build-isolation --no-deps -vvv --wheel-dir ./dist . || error_exit "Failed to setup metapackage"
cd transformer_engine/pytorch cd transformer_engine/pytorch
NVTE_RELEASE_BUILD=1 python3 setup.py sdist || error_exit "Failed to setup sdist" NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation --no-deps -vvv --wheel-dir ./dist . || error_exit "Failed to setup sdist"
pip3 install dist/* || error_exit "Failed to install dist/*" pip3 install --no-build-isolation --no-deps -vvv dist/* || error_exit "Failed to install dist/*"
cd $TE_PATH cd $TE_PATH
pip3 install dist/*.whl --no-deps || error_exit "Failed to install dist/*.whl --no-deps" pip3 install --no-build-isolation --no-deps -vvv dist/*.whl || error_exit "Failed to install dist/*.whl --no-deps"
python3 $TE_PATH/tests/pytorch/test_sanity_import.py || test_fail "test_sanity_import.py" python3 $TE_PATH/tests/pytorch/test_sanity_import.py || test_fail "test_sanity_import.py"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment