TransformerEngine Image Build (#1204)

* install

* fix

* try chucks branch

* fix image

* pin

* fix lint

* revert image size

* remove verbose

* switch to 8wide

---------

.github/workflows/docker.yaml

@@ -12,7 +12,7 @@ on:
   workflow_dispatch: {}
 jobs:
   docker-build:
-    runs-on: mosaic-4wide
+    runs-on: mosaic-8wide
     if: github.repository_owner == 'mosaicml'
     strategy:
       matrix:

Dockerfile

@@ -12,6 +12,9 @@ ARG DEP_GROUPS
 ADD https://raw.githubusercontent.com/mosaicml/llm-foundry/$BRANCH_NAME/setup.py setup.py
 RUN rm setup.py
 
+# Install TransformerEngine
+RUN NVTE_FRAMEWORK=pytorch CMAKE_BUILD_PARALLEL_LEVEL=4 MAX_JOBS=4 pip install git+https://github.com/NVIDIA/TransformerEngine.git@05eb6deb31c1b48e9f4380d18fe95f3c38e84335
+
 # Install and uninstall foundry to cache foundry requirements
 RUN git clone -b $BRANCH_NAME https://github.com/mosaicml/llm-foundry.git
 RUN pip install --no-cache-dir "./llm-foundry${DEP_GROUPS}"