wip

.github/workflows/docker.yaml

@@ -3,6 +3,12 @@ on:
   push:
     branches:
     - main
+  pull_request:
+    branches:
+      - main
+    paths:
+      - ./Dockerfile
+      - .github/workflows/docker.yaml
   workflow_dispatch: {}
 jobs:
   docker-build:

setup.py

@@ -91,6 +91,12 @@
     # PyPI does not support direct dependencies, so we remove this line before uploading from PyPI
     'xentropy-cuda-lib@git+https://github.com/HazyResearch/[email protected]#subdirectory=csrc/xentropy',
 ]
+extra_deps['gpu-flash2'] = [
+    'flash-attn==2.3.1',
+    'mosaicml-turbo==0.0.4',
+    # PyPI does not support direct dependencies, so we remove this line before uploading from PyPI
+    'xentropy-cuda-lib@git+https://github.com/HazyResearch/[email protected]#subdirectory=csrc/xentropy',
+]
 
 extra_deps['peft'] = [
     'loralib==0.1.1',  # lora core
@@ -107,7 +113,8 @@
 ]
 extra_deps['all-cpu'] = set(
     dep for key, deps in extra_deps.items() for dep in deps if 'gpu' not in key)
-extra_deps['all'] = set(dep for deps in extra_deps.values() for dep in deps)
+extra_deps['all'] = set(dep for key, deps in extra_deps.items() for dep in deps if key != 'gpu-flash2')
+extra_deps['all-flash2'] = set(dep for key, deps in extra_deps.items() for dep in deps if key != 'gpu')
 
 setup(
     name=_PACKAGE_NAME,