From db2a3586f3757123a7951dfeca374061cc384f01 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Mon, 31 Jul 2023 16:31:53 -0400
Subject: [PATCH] add peft install back since it doesn't get installed by
 setup.py (#331)

---
 docker/Dockerfile           |  1 +
 src/axolotl/utils/models.py | 18 +++++++++++-------
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 590cff664..b5198e1d9 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -11,6 +11,7 @@ RUN apt-get update && \
 
 WORKDIR /workspace
 
+RUN pip3 install --force-reinstall "peft @ git+https://github.com/huggingface/peft.git@main"
 RUN git clone --depth=1 https://github.com/OpenAccess-AI-Collective/axolotl.git
 # If AXOLOTL_EXTRAS is set, append it in brackets
 RUN cd axolotl && \
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
index afe1632ea..fa6a7c1c3 100644
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -147,13 +147,17 @@ def load_model(
         LOG.exception(err)
         raise err
 
-    try:
-        from peft import prepare_model_for_kbit_training
-    except ImportError:
-        # For backward compatibility
-        from peft import (
-            prepare_model_for_int8_training as prepare_model_for_kbit_training,
-        )
+    if not cfg.gptq and (
+        (cfg.adapter == "lora" and load_in_8bit)
+        or (cfg.adapter == "qlora" and cfg.load_in_4bit)
+    ):
+        try:
+            from peft import prepare_model_for_kbit_training
+        except ImportError:
+            # For backward compatibility
+            from peft import (
+                prepare_model_for_int8_training as prepare_model_for_kbit_training,
+            )
 
     model_kwargs = {}
     if cfg.model_revision: