Cleanup Dockerfile-notebook

Dockerfile-notebook

@@ -12,27 +12,18 @@ RUN apt update && \
     apt update && \
     apt install -y --no-install-recommends --allow-change-held-packages libnccl2 libnccl-dev
 
-# Install torch
+# Install torch and axolotl
 USER jovyan
 COPY requirements.txt notebook-requirements.txt /tmp/llm-finetune/
 RUN pip install -U pip wheel setuptools && \
     pip uninstall -y axolotl && \
-    pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt
+    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install --no-cache-dir -U --use-pep517 -r /tmp/llm-finetune/notebook-requirements.txt
 
 # Setup editable packages
 USER root
 RUN mkdir -p /packages && \
     chown -R jovyan:users /packages
 
-# Install axolotl
-USER jovyan
-RUN cd /packages && \
-    git clone https://github.com/truefoundry/axolotl && \
-    cd axolotl/ && \
-    git checkout b8db5db0fea9a1dad15338c1daf73a04f647caf4 && \
-    cd /packages/axolotl/ && \
-    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --use-pep517 --no-build-isolation --no-cache-dir -e .[flash-attn,mamba-ssm,optimizers,lion-pytorch,galore]
-
 # Install axolotl_truefoundry plugin with our requirements overrides over axolotl
 COPY --chown=jovyan:users plugins/axolotl_truefoundry /packages/axolotl_truefoundry
 RUN cd /packages/axolotl_truefoundry/ && \