Install all deps in Docker images (#1431)

mosaicml · Aug 6, 2024 · 8fe57eb · 8fe57eb
1 parent 6dcc18a
commit 8fe57eb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
@@ -19,10 +19,10 @@ jobs:
         include:
         - name: "2.3.1_cu121"
           base_image: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04
-          dep_groups: "[gpu]"
+          dep_groups: "[all]"
         - name: "2.3.1_cu121_aws"
           base_image: mosaicml/pytorch:2.3.1_cu121-python3.11-ubuntu20.04-aws
-          dep_groups: "[gpu]"
+          dep_groups: "[all]"
     steps:
 
     - name: Checkout

diff --git a/Dockerfile b/Dockerfile
@@ -7,6 +7,8 @@ FROM $BASE_IMAGE
 ARG BRANCH_NAME
 ARG DEP_GROUPS
 
+ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 8.9 9.0"
+
 # 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