Merge pull request #47 from coreweave/es/extras-xformers

feat(torch-extras): Add `xformers`

torch-extras/Dockerfile

@@ -4,6 +4,7 @@ ARG BASE_IMAGE
 ARG DEEPSPEED_VERSION="0.10.3"
 ARG FLASH_ATTN_VERSION="2.0.2"
 ARG APEX_COMMIT="38a12698bc3cc95987bca270bcd6d025bb0be346"
+ARG XFORMERS_VERSION="0.0.22.post7"
 
 FROM alpine/git:2.36.3 as flash-attn-downloader
 WORKDIR /git
@@ -215,6 +216,20 @@ RUN --mount=type=bind,from=apex-downloader,source=/git/apex,target=apex/,rw \
 WORKDIR /wheels
 
 
+FROM builder-base as xformers-builder
+
+ARG XFORMERS_VERSION
+
+RUN python3 -m pip install -U --no-cache-dir \
+      setuptools wheel pip && \
+    CC=$(realpath -e ./compiler) \
+      MAX_JOBS=$(($(./effective_cpu_count.sh) / 2 + 1)) \
+      python3 -m pip wheel -w /wheels -v \
+      --no-cache-dir --no-build-isolation --no-deps \
+      --no-binary=xformers \
+      xformers==${XFORMERS_VERSION}
+
+
 FROM ${BASE_IMAGE}
 
 RUN apt-get -qq update && \
@@ -227,3 +242,5 @@ RUN --mount=type=bind,from=flash-attn-builder,source=/wheels,target=/tmp/wheels
     python3 -m pip install --no-cache-dir /tmp/wheels/*.whl
 RUN --mount=type=bind,from=apex-builder,source=/wheels,target=/tmp/wheels \
     python3 -m pip install --no-cache-dir /tmp/wheels/*.whl
+RUN --mount=type=bind,from=xformers-builder,source=/wheels,target=/tmp/wheels \
+    python3 -m pip install --no-cache-dir /tmp/wheels/*.whl