From feaf12999adb57fa88662f695d1a395e5fba4c6b Mon Sep 17 00:00:00 2001 From: Daniel King Date: Fri, 6 Oct 2023 14:59:06 -0700 Subject: [PATCH] wip --- .github/workflows/docker.yaml | 6 ++++++ setup.py | 9 ++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index 8e30554475..aa740530c0 100644 --- a/.github/workflows/docker.yaml +++ b/.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: diff --git a/setup.py b/setup.py index be5b6708a3..7e355ad2fb 100644 --- a/setup.py +++ b/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/flash-attention.git@v1.0.9#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/flash-attention.git@v2.3.1#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,