From 247c99436c622e7993ff850c774468ed14b0e5ea Mon Sep 17 00:00:00 2001
From: Michael Waskom <michael@modal.com>
Date: Tue, 6 Feb 2024 14:29:28 -0500
Subject: [PATCH] Update base image spec to axolotl 0.4.0

---
 src/common.py | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/src/common.py b/src/common.py
index bceb01f..78e4048 100644
--- a/src/common.py
+++ b/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")])