diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 94868d47fd..294f832bca 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -22,7 +22,6 @@ jobs: - name: "2.4.0_cu124_aws" base_image: mosaicml/pytorch:2.4.0_cu124-python3.11-ubuntu20.04-aws dep_groups: "[all]" - te_commit: 901e5d2 aws_extension: "_aws" steps: - name: Checkout source @@ -63,10 +62,8 @@ jobs: ${{ env.DOCKER_TAG }} ${{ env.DOCKER_LATEST_TAG }} cache-from: type=registry,ref=${{ env.IMAGE_CACHE }},scope=global - cache-to: type=registry,ref=${{ env.IMAGE_CACHE }},mode=max,scope=global build-args: | BRANCH_NAME=${{ github.ref_name }} BASE_IMAGE=${{ matrix.base_image }} DEP_GROUPS=${{ matrix.dep_groups }} - TE_COMMIT=${{ matrix.te_commit }} KEEP_FOUNDRY=true diff --git a/Dockerfile b/Dockerfile index 83838df286..f2566cd3cc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,14 +6,10 @@ FROM $BASE_IMAGE ARG BRANCH_NAME ARG DEP_GROUPS -ARG TE_COMMIT ARG KEEP_FOUNDRY=false ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 8.9 9.0" -# Install TransformerEngine -RUN NVTE_FRAMEWORK=pytorch CMAKE_BUILD_PARALLEL_LEVEL=4 MAX_JOBS=4 pip install git+https://github.com/NVIDIA/TransformerEngine.git@$TE_COMMIT - # Check for changes in setup.py. # If there are changes, the docker cache is invalidated and a fresh pip installation is triggered. ADD https://raw.githubusercontent.com/mosaicml/llm-foundry/$BRANCH_NAME/setup.py setup.py