diff --git a/.github/actions/build-pytorch-wheel/Dockerfile b/.github/actions/build-pytorch-wheel/Dockerfile index 05ac20db0d..d3b0bd6f98 100644 --- a/.github/actions/build-pytorch-wheel/Dockerfile +++ b/.github/actions/build-pytorch-wheel/Dockerfile @@ -36,10 +36,18 @@ RUN CUDA_MAJOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'}) && \ dpkg -i cuda-keyring_1.1-1_all.deb && \ rm cuda-keyring_1.1-1_all.deb && \ apt-get update && \ - apt-get install -y cuda-toolkit-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} cudnn-cuda-$CUDA_MAJOR_VERSION libcudnn$CUDNN_MAJOR_VERSION-cuda-$CUDA_MAJOR_VERSION libnccl2 libnccl-dev cmake + apt-get install -y \ + cuda-compiler-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ + cuda-libraries-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ + cuda-nvtx-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \ + cudnn-cuda-$CUDA_MAJOR_VERSION \ + libcudnn$CUDNN_MAJOR_VERSION-cuda-$CUDA_MAJOR_VERSION \ + libnccl2 \ + libnccl-dev \ + cmake # Install PyTorch RUN export MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'}) && \ export MATRIX_TORCH_VERSION=$(echo $TORCH_VERSION | awk -F \. {'print $1 "." $2'}) && \ - export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; versions = {'2.5': (118, 124), '2.6': (118, 126), '2.7': (118, 128), '2.8': (126, 129), '2.9': (126, 130)}; minv, maxv = versions[env['MATRIX_TORCH_VERSION']]; print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)") && \ + export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; versions = {'2.5': (118, 124), '2.6': (118, 126), '2.7': (118, 128), '2.8': (126, 129), '2.9': (126, 130), '2.11': (130, 130)}; minv, maxv = versions[env['MATRIX_TORCH_VERSION']]; print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)") && \ pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} diff --git a/.github/workflows/attach-wheels-to-release.yml b/.github/workflows/attach-wheels-to-release.yml index fb56622d77..bf520ac64f 100644 --- a/.github/workflows/attach-wheels-to-release.yml +++ b/.github/workflows/attach-wheels-to-release.yml @@ -78,7 +78,27 @@ jobs: "torch-version": ["2.8.0"], "cuda-version": ["12.9.1"], "cudnn-version": ["9"], - "cxx11_abi": ["TRUE"] + "cxx11_abi": ["TRUE"], + "include": [ + { + "os": "ubuntu-22.04", + "release-version": "${{ github.event.release.tag_name }}", + "python-version": "3.12", + "torch-version": "2.11.0", + "cuda-version": "13.0.0", + "cudnn-version": "9", + "cxx11_abi": "TRUE" + }, + { + "os": "ubuntu-22.04-arm", + "release-version": "${{ github.event.release.tag_name }}", + "python-version": "3.12", + "torch-version": "2.11.0", + "cuda-version": "13.0.0", + "cudnn-version": "9", + "cxx11_abi": "TRUE" + } + ] }' | jq -rc) else MATRIX=$(echo '{ @@ -166,6 +186,71 @@ jobs: asset_name: ${{ steps.build-pytorch-wheel.outputs.wheel_name }} asset_content_type: application/* + build_core_wheel: + name: Build Core Wheel (ubuntu-22.04-arm, CUDA 13) + runs-on: ubuntu-22.04-arm + needs: pre-flight + if: ${{ github.event_name == 'release' || inputs.ngc-image == '' }} + steps: + - name: 'Checkout' + uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 + with: + ref: ${{ github.event_name == 'release' && github.event.release.tag_name || inputs.release-version }} + submodules: recursive + + - name: Move /var/lib/docker/ + run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" + + - name: Maximize build space + uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 + with: + root-reserve-mb: 5120 + temp-reserve-mb: 32 + swap-size-mb: 10240 + remove-dotnet: 'true' + remove-android: 'true' + remove-haskell: 'true' + remove-codeql: 'true' + build-mount-path: '/var/lib/docker/' + + - name: Restore /var/lib/docker/ + run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" + + - name: Build CUDA 13 core wheel + id: build-core-wheel + run: | + docker build \ + --build-arg CUDA_MAJOR=13 \ + --build-arg CUDA_MINOR=0 \ + --build-arg BUILD_METAPACKAGE=false \ + --build-arg BUILD_COMMON=true \ + --build-arg BUILD_PYTORCH=false \ + --build-arg BUILD_JAX=false \ + --tag transformer-engine-core-wheel \ + --file build_tools/wheel_utils/Dockerfile.aarch \ + . + + mkdir -p dist + docker run --rm \ + --volume "$(pwd)/dist:/wheelhouse" \ + transformer-engine-core-wheel + + wheel_path=$(find dist -maxdepth 1 -type f \ + -name 'transformer_engine_cu13-*-manylinux_2_28_aarch64.whl' \ + -print -quit) + test -n "$wheel_path" + echo "wheel_name=$(basename "$wheel_path")" | tee -a "$GITHUB_OUTPUT" + + - name: Upload release asset + uses: actions/upload-release-asset@e8f9f06c4b078e705bd2ea027f0926603fc9b4d5 # v1.0.2 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + upload_url: ${{ needs.pre-flight.outputs.release-assets-url }} + asset_path: ./dist/${{ steps.build-core-wheel.outputs.wheel_name }} + asset_name: ${{ steps.build-core-wheel.outputs.wheel_name }} + asset_content_type: application/* + build_wheels_for_ngc: name: Build Wheels for NGC PyTorch images runs-on: ${{ matrix.os }} diff --git a/build_tools/wheel_utils/Dockerfile.aarch b/build_tools/wheel_utils/Dockerfile.aarch index c040dadcdb..d0b2ea74ed 100644 --- a/build_tools/wheel_utils/Dockerfile.aarch +++ b/build_tools/wheel_utils/Dockerfile.aarch @@ -25,13 +25,12 @@ ENV CUDA_MAJOR=${CUDA_MAJOR} RUN dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo RUN dnf -y install cuda-compiler-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \ cuda-libraries-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \ - cuda-libraries-devel-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 + cuda-libraries-devel-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \ + cuda-nvtx-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 RUN dnf -y install --allowerasing cudnn9-cuda-${CUDA_MAJOR} RUN dnf clean all RUN rm -rf /var/cache/dnf/* RUN echo "/usr/local/cuda/lib64" >> /etc/ld.so.conf.d/999_nvidia_cuda.conf -RUN dnf -y install cuda-toolkit-${CUDA_MAJOR} -RUN dnf clean all RUN dnf -y install glog.aarch64 glog-devel.aarch64 RUN dnf -y install libnccl libnccl-devel libnccl-static