separate out flash-attn install (sadly)

docker/Dockerfile

@@ -35,6 +35,7 @@ RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
     else \
         uv pip install --python "$VENV_PYTHON" --no-build-isolation -e .[ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \
     fi && \
+    uv pip install --python "$VENV_PYTHON" --no-build-isolation flash-attn $AXOLOTL_ARGS && \
     "$VENV_PYTHON" scripts/unsloth_install.py | sh && \
     "$VENV_PYTHON" scripts/cutcrossentropy_install.py | sh && \
     uv pip install --python "$VENV_PYTHON" pytest

docker/Dockerfile-base

@@ -48,5 +48,5 @@ RUN git lfs install --skip-repo && \
     pip3 cache purge
 
 RUN if [ "$PYTORCH_VERSION" = "2.6.0" ] && [ "$CUDA" = "124" ] ; then \
-        FLASH_ATTENTION_FORCE_BUILD="TRUE" pip3 install --no-build-isolation flash-attn==2.8.0.post2; \
+        FLASH_ATTENTION_FORCE_BUILD="TRUE" uv pip install --no-build-isolation flash-attn==2.8.0.post2; \
     fi

docker/Dockerfile-tests

@@ -24,13 +24,14 @@ RUN git fetch origin +$GITHUB_REF && \
 
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install --no-build-isolation -e .[deepspeed,flash-attn,mamba-ssm,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        uv pip install --no-build-isolation -e .[deepspeed,mamba-ssm,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \
-        pip install --no-build-isolation -e .[deepspeed,flash-attn,mamba-ssm] $AXOLOTL_ARGS; \
-    fi
+        uv pip install --no-build-isolation -e .[deepspeed,mamba-ssm] $AXOLOTL_ARGS; \
+    fi && \
+    uv pip install --no-build-isolation flash-attn $AXOLOTL_ARGS
 
 # So we can test the Docker image
-RUN pip install pytest
+RUN uv pip install pytest
 
 # fix so that git fetch/pull from remote works
 RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \