From 33a9fd6557be1606d043e0aa527271a673d9cd53 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Thu, 4 Jul 2024 20:11:39 -0400
Subject: [PATCH] make sure to install optimizer packages

---
 cicd/Dockerfile.jinja | 4 ++--
 docker/Dockerfile     | 4 ++--
 setup.py              | 6 ++++++
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/cicd/Dockerfile.jinja b/cicd/Dockerfile.jinja
index ce03e08c2..20fa238a4 100644
--- a/cicd/Dockerfile.jinja
+++ b/cicd/Dockerfile.jinja
@@ -24,9 +24,9 @@ RUN git fetch origin +$GITHUB_REF && \
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN pip install causal_conv1d
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers] $AXOLOTL_ARGS; \
     fi
 
 # So we can test the Docker image
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 6b9cf7d4c..34ce11f7d 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -22,9 +22,9 @@ WORKDIR /workspace/axolotl
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN pip install causal_conv1d
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \
-        pip install -e .[deepspeed,flash-attn,mamba-ssm,galore] $AXOLOTL_ARGS; \
+        pip install -e .[deepspeed,flash-attn,mamba-ssm,optimizers] $AXOLOTL_ARGS; \
     fi
 
 # So we can test the Docker image
diff --git a/setup.py b/setup.py
index c7b4e15de..d2bea3bdb 100644
--- a/setup.py
+++ b/setup.py
@@ -101,5 +101,11 @@ def parse_requirements():
         "galore": [
             "galore_torch",
         ],
+        "optimizers": [
+            "galore_torch",
+            "lion-pytorch==0.1.2",
+            "lomo-optim==0.1.1",
+            "torch-optimi==0.2.1",
+        ],
     },
 )