rm TE

.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

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