fix: warn user to install mamba_ssm package (#1019)

2024-01-10 16:50:56 +09:00
parent 9e3f0cb5a7
commit d69ba2b0b7
4 changed files with 24 additions and 10 deletions
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -20,9 +20,9 @@ WORKDIR /workspace/axolotl

 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,$AXOLOTL_EXTRAS]; \
+        pip install -e .[deepspeed,flash-attn,mamba-ssm,$AXOLOTL_EXTRAS]; \
    else \
-        pip install -e .[deepspeed,flash-attn]; \
+        pip install -e .[deepspeed,flash-attn,mamba-ssm]; \
    fi

 # So we can test the Docker image