diff --git a/Dockerfile b/Dockerfile
index ea912a3..9704bb4 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -9,7 +9,7 @@ RUN mkdir -p /packages && \
     cd /packages && \
     git clone https://github.com/truefoundry/axolotl && \
     cd axolotl/ && \
-    git checkout beb43a96d91273da42ec454d772de144ec5f8d6d
+    git checkout bfe2d6b2e5019a8d4486ac8a7726089928e6b3a6
 RUN cd /packages/axolotl/ && \
     MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation -e .[flash-attn,mamba-ssm,fused-dense-lib] && \
     pip install --no-cache-dir -U -r /tmp/requirements.txt && \
diff --git a/Dockerfile-notebook b/Dockerfile-notebook
index 6daf46b..9030a31 100644
--- a/Dockerfile-notebook
+++ b/Dockerfile-notebook
@@ -21,7 +21,7 @@ USER jovyan
 RUN cd /packages && \
     git clone https://github.com/truefoundry/axolotl && \
     cd axolotl/ && \
-    git checkout beb43a96d91273da42ec454d772de144ec5f8d6d
+    git checkout bfe2d6b2e5019a8d4486ac8a7726089928e6b3a6
 RUN cd /packages/axolotl/ && \
     MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -U --no-build-isolation -e .[flash-attn,mamba-ssm,fused-dense-lib] && \
     pip install --no-cache-dir -U -r /tmp/llm-finetune/notebook-requirements.txt