Update base image spec to axolotl 0.4.0

src/common.py

@@ -3,24 +3,18 @@
 
 APP_NAME = "example-axolotl"
 
-# Latest image hash of winglian/axolotl:main-py3.10-cu118-2.0.1 (2023-12-11)
+# Axolotl image hash corresponding to 0.4.0 release
 AXOLOTL_REGISTRY_SHA = (
-    "5c19a5154fd522225953b9c3f6206750f4191e0e92ee424f02963f7963ada698"
+    "af4d878e9fbc90c7ba30fa78ce4d6d95b1ccba398ab944efbd322d7c0d6313c8"
 )
-# Need to patch transformers to an older version to avoid checkpointing errors.
-TRANSFORMERS_SHA = "5324bf9c07c318015eccc5fba370a81368a8df28"
 
 axolotl_image = (
     Image.from_registry(f"winglian/axolotl@sha256:{AXOLOTL_REGISTRY_SHA}")
     .run_commands(
         "git clone https://github.com/OpenAccess-AI-Collective/axolotl /root/axolotl",
-        "cd /root/axolotl && git checkout a581e9f8f66e14c22ec914ee792dd4fe073e62f6",
+        "cd /root/axolotl && git checkout v0.4.0",
     )
     .pip_install("huggingface_hub==0.19.4", "hf-transfer==0.1.4")
-    .pip_install(
-        f"transformers @ git+https://github.com/huggingface/transformers.git@{TRANSFORMERS_SHA}",
-        "--force-reinstall",
-    )
     .env(dict(HUGGINGFACE_HUB_CACHE="/pretrained", HF_HUB_ENABLE_HF_TRANSFER="1"))
 )
 
@@ -29,8 +23,6 @@
 ).pip_install(
     "vllm==0.2.5",
     "torch==2.1.2",
-    "torchvision==0.16.2",
-    "torchaudio==2.1.2",
 )
 
 stub = Stub(APP_NAME)  # , secrets=[Secret.from_name("huggingface")])