From d282f324815bd8d312cbe94c05859a08cf28a12c Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 14 Jan 2026 12:03:55 -0500 Subject: [PATCH] don't install deepspeed in arm64 images (#3357) --- docker/Dockerfile | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 116361dcd..d80cede55 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -6,6 +6,7 @@ ARG AXOLOTL_EXTRAS="" ARG AXOLOTL_ARGS="" ARG CUDA="118" ARG PYTORCH_VERSION="2.1.2" +ARG TARGETARCH ENV PYTORCH_VERSION=$PYTORCH_VERSION @@ -20,13 +21,17 @@ RUN git clone --depth=1 https://github.com/axolotl-ai-cloud/axolotl.git WORKDIR /workspace/axolotl -# If AXOLOTL_EXTRAS is set, append it in brackets -RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \ - pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \ +# If AXOLOTL_EXTRAS is set, append it in brackets; don't install deepspeed with arm64 +RUN if [ "$TARGETARCH" = "arm64" ]; then \ + BASE_EXTRAS="flash-attn,ring-flash-attn,optimizers,ray"; \ else \ - pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray] $AXOLOTL_ARGS; \ + BASE_EXTRAS="deepspeed,flash-attn,ring-flash-attn,optimizers,ray"; \ fi && \ - python scripts/unsloth_install.py | sh && \ + if [ "$AXOLOTL_EXTRAS" != "" ]; then \ + pip install --no-build-isolation -e .[$BASE_EXTRAS,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \ + else \ + pip install --no-build-isolation -e .[$BASE_EXTRAS] $AXOLOTL_ARGS; \ + fi && \ python scripts/unsloth_install.py | sh && \ python scripts/cutcrossentropy_install.py | sh && \ pip install pytest && \ pip cache purge