use transformer engine v1.1

Azure · Feb 22, 2024 · 1b0d2cf · 1b0d2cf
1 parent eadca11
commit 1b0d2cf
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/dockerfile/torch1.14-cuda11.8.dockerfile b/dockerfile/torch1.14-cuda11.8.dockerfile
@@ -50,7 +50,7 @@ RUN cd third_party/msccl && \
     make install
 # cache TE build to save time in CI
 RUN python3 -m pip install --upgrade pip && \
-    python3 -m pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+    python3 -m pip install git+https://github.com/NVIDIA/TransformerEngine.git@v1.1
 
 ADD . .
 RUN python3 -m pip install . && \

diff --git a/dockerfile/torch2.1-cuda12.2.dockerfile b/dockerfile/torch2.1-cuda12.2.dockerfile
@@ -50,7 +50,7 @@ RUN cd third_party/msccl && \
     make install
 # cache TE build to save time in CI
 RUN python3 -m pip install --upgrade pip && \
-    python3 -m pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+    python3 -m pip install git+https://github.com/NVIDIA/TransformerEngine.git@v1.1
 
 ADD . .
 RUN python3 -m pip install . && \

diff --git a/pyproject.toml b/pyproject.toml
@@ -38,7 +38,7 @@ classifiers=[
 ]
 dependencies = [
     "torch",
-    "transformer-engine@git+https://github.com/NVIDIA/TransformerEngine.git@stable",
+    "transformer-engine@git+https://github.com/NVIDIA/TransformerEngine.git@v1.1",
     "colorlog>=6.7.0",
     "deepspeed==0.13.1",
     "mpi4py",