Update dockerfiles

.gitignore

@@ -17,3 +17,4 @@ mlruns/
 output/
 outputs/
 *.egg-info/
+.ipynb_checkpoints/

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
           - --line-length=120
           - --safe
           - --target-version=py310
-        stages: [commit]
+        stages: [pre-commit]
   - repo: https://github.com/pycqa/isort
     rev: 5.12.0
     hooks:

Dockerfile

@@ -1,22 +1,31 @@
 # https://hub.docker.com/layers/winglian/axolotl/main-20241104-py3.11-cu121-2.3.1/images/sha256-790297fa1d71f8f1590c73ca4505ca39fe7dfa2886b2b6862199a6df679bf8e4?context=explore
 FROM winglian/axolotl@sha256:cffbcc4993e80301a8918062f8136a6ac402877fd6c29f1168be563e543aee4d
+SHELL ["/bin/bash", "-c"]
 USER root
+
+# Install torch
 COPY requirements.txt /tmp/
 RUN pip install -U pip wheel setuptools && \
     pip uninstall -y axolotl && \
     pip install --no-cache-dir -U -r /tmp/requirements.txt
+
+# Install axolotl
 RUN mkdir -p /packages && \
     cd /packages && \
     git clone https://github.com/truefoundry/axolotl && \
     cd axolotl/ && \
-    git checkout e16f637d079ef5d56321a240ef0547a50c37b708
-RUN cd /packages/axolotl/ && \
+    git checkout e16f637d079ef5d56321a240ef0547a50c37b708 && \
+    cd /packages/axolotl/ && \
     MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation --no-cache-dir -e .[flash-attn,mamba-ssm,fused-dense-lib,optimizers,lion-pytorch,galore] && \
-    pip install --no-cache-dir -U -r /tmp/requirements.txt && \
     rm -rf /root/.cache/pip
+
+# Install axolotl_truefoundry plugin with our requirements overrides over axolotl
 COPY plugins/axolotl_truefoundry /packages/axolotl_truefoundry
 RUN cd /packages/axolotl_truefoundry/ && \
-    pip install --no-cache-dir -e . && \
+    pip install --no-cache-dir -U -r /tmp/requirements.txt -e . && \
     rm -rf /root/.cache/pip
+
 WORKDIR /app
+
+# Add source code for finetuning
 COPY . /app

Dockerfile-notebook

@@ -2,32 +2,42 @@ FROM tfy.jfrog.io/tfy-images/jupyter:0.3.4-cu121-py3.11.10-sudo
 SHELL ["/bin/bash", "-c"]
 ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
 ENV DEBIAN_FRONTEND=noninteractive
+
+# upgrade libnccl
 USER root
 RUN apt update && \
     apt install -y --no-install-recommends git curl wget && \
     wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb -O /tmp/cuda-keyring_1.1-1_all.deb && \
     dpkg -i /tmp/cuda-keyring_1.1-1_all.deb && \
     apt update && \
     apt install -y --no-install-recommends --allow-change-held-packages libnccl2 libnccl-dev
+
+# Install torch
 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
+
+# Setup editable packages
 USER root
 RUN mkdir -p /packages && \
     chown -R jovyan:users /packages
+
+# Install axolotl
 USER jovyan
-RUN pip install -U pip setuptools wheel && \
-    cd /packages && \
+RUN cd /packages && \
     git clone https://github.com/truefoundry/axolotl && \
     cd axolotl/ && \
-    git checkout e16f637d079ef5d56321a240ef0547a50c37b708
-RUN cd /packages/axolotl/ && \
-    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation --no-cache-dir -e .[flash-attn,mamba-ssm,fused-dense-lib,optimizers,lion-pytorch,galore] && \
-    pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt
+    git checkout e16f637d079ef5d56321a240ef0547a50c37b708 && \
+    cd /packages/axolotl/ && \
+    MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation --no-cache-dir -e .[flash-attn,mamba-ssm,fused-dense-lib,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/ && \
-    pip install --no-cache-dir -e .
+    pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt -e .
+
+# Add source code for finetuning
 COPY --chown=jovyan:users . /tmp_home/jovyan/llm-finetune/
 ENV JUPYTER_APP_LAUNCHER_PATH=/home/jovyan/llm-finetune/.jp_app_launcher/

plugins/axolotl_truefoundry/pyproject.toml

@@ -9,9 +9,9 @@ description = "TrueFoundry plugin for Axolotl"
 requires-python = ">=3.8.1,<4.0"
 dependencies = [
     "transformers>=4.0.0,<5",
-    "truefoundry==0.4.4",
+    "truefoundry>=0.4.4,<0.5.0",
     "pynvml>=11.0.0,<12",
-    "torch>=2.3.0,<3",
+    "torch>=2.3.0,<2.4.0",
     "pydantic>=2.0.0,<3",
 ]
 

requirements.txt

@@ -7,5 +7,5 @@ rich>=13.0.0,<14
 snowflake-connector-python[pandas]==3.12.3
 torch==2.3.1+cu121
 torchao==0.6.1+cu121
-truefoundry==0.4.4
+truefoundry>=0.4.4,<0.5.0
 unsloth @ git+https://github.com/unslothai/unsloth@38663b01f5dd0e610b12475bd95b144303cff539