prepare ia3 for 8bit

axolotl-ai-cloud · Sep 18, 2023 · 3124670 · 3124670
1 parent 31b9e0c
commit 3124670
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/.pylintrc b/.pylintrc
@@ -12,3 +12,4 @@ generated-members=numpy.*, torch.*
 disable=missing-function-docstring, line-too-long, import-error,
     too-many-arguments, too-many-locals, too-many-statements, too-many-branches, too-few-public-methods,
     too-many-instance-attributes, fixme, import-outside-toplevel, logging-fstring-interpolation,
+    too-many-boolean-expressions,
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -367,8 +367,10 @@ def load_model(
                 module.to(torch.float32)
 
     needs_fa2_dtype = cfg.adapter or cfg.fsdp
-    if (cfg.adapter == "lora" and load_in_8bit) or (
-        cfg.adapter == "qlora" and cfg.load_in_4bit
+    if (
+        (cfg.adapter == "lora" and load_in_8bit)
+        or (cfg.adapter == "qlora" and cfg.load_in_4bit)
+        or (cfg.adapter == "ia3" and load_in_8bit)
     ):
         LOG.info("converting PEFT model w/ prepare_model_for_kbit_training")
         if cfg.gradient_checkpointing: