commit change

mosaicml · Nov 10, 2023 · 3b02b01 · 3b02b01
1 parent 3f32995
commit 3b02b01
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
@@ -26,13 +26,15 @@ jobs:
         - name: '2.1.0_cu121'
           base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04
           dep_groups: '[gpu]'
+        - name: '2.1.0_cu121_flash2'
+          base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04
+          dep_groups: '[gpu-flash2]'
         - name: '2.1.0_cu121_aws'
           base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04-aws
           dep_groups: '[gpu]'
-        - name: '2.1.0_cu121_flash2'
-          base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04
+        - name: '2.1.0_cu121_flash2_aws'
+          base_image: mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04-aws
           dep_groups: '[gpu-flash2]'
-
     steps:
     - name: Maximize Build Space on Worker
       uses: easimon/maximize-build-space@v4

diff --git a/README.md b/README.md
@@ -120,6 +120,8 @@ You can select a specific commit hash such as `mosaicml/llm-foundry:1.13.1_cu117
 | `mosaicml/llm-foundry:2.0.1_cu118-latest`                   | 2.0.1          | 11.8         | Yes                                 |
 | `mosaicml/llm-foundry:2.1.0_cu121-latest`                   | 2.1.0          | 12.1         | Yes (flash attention v1)            |
 | `mosaicml/llm-foundry:2.1.0_cu121_flash2-latest`            | 2.1.0          | 12.1         | Yes (flash attention v2)            |
+| `mosaicml/llm-foundry:2.1.0_cu121_aws-latest`               | 2.1.0          | 12.1         | Yes (flash attention v2)            |
+| `mosaicml/llm-foundry:2.1.0_cu121_flash2_aws-latest`        | 2.1.0          | 12.1         | Yes (flash attention v2)            |
 
 
 # Installation