Fix axolotl bug with LlamaRotaryEmbedding

Dockerfile

@@ -9,7 +9,7 @@ RUN mkdir -p /packages && \
     cd /packages && \
     git clone https://github.com/OpenAccess-AI-Collective/axolotl && \
     cd axolotl/ && \
-    git checkout 7803f0934f71c503589d4097598f13129105700c
+    git checkout 40a88e8c4a2f32b63df0fe2079f7acfe73329273
 RUN cd /packages/axolotl/ && \
     MAX_JOBS=1 NVCC_APPEND_FLAGS="--verbose --threads 1" pip install -v -U --no-build-isolation -e .[deepspeed,flash-attn,mamba-ssm,fused-dense-lib] && \
     pip uninstall -y mlflow tfy-mlflow-client && \

Dockerfile-notebook

@@ -21,7 +21,7 @@ USER jovyan
 RUN cd /packages && \
     git clone https://github.com/OpenAccess-AI-Collective/axolotl && \
     cd axolotl/ && \
-    git checkout 43265208299242e3bc32690e22efadef79365c9d
+    git checkout 40a88e8c4a2f32b63df0fe2079f7acfe73329273
 RUN cd /packages/axolotl/ && \
     MAX_JOBS=1 NVCC_APPEND_FLAGS="--threads 1" pip install -v -U --no-build-isolation -e .[deepspeed,flash-attn,mamba-ssm,fused-dense-lib] && \
     pip uninstall -y mlflow tfy-mlflow-client && \

config-base.yaml

@@ -63,6 +63,8 @@ flash_attn_cross_entropy: True
 flash_attn_rms_norm: True
 gradient_accumulation_steps: 4
 gradient_checkpointing: True
+gradient_checkpointing_kwargs:
+  use_reentrant: True
 learning_rate: 0.00001
 load_best_model_at_end: True
 load_in_8bit: False