Update src/axolotl/utils/models.py

Co-authored-by: NanoCode012 <kevinvong@rocketmail.com>
2023-05-24 23:31:12 -04:00
parent ffd1043607
commit 1cf21daf51
1 changed files with 5 additions and 1 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -362,7 +362,11 @@ def load_lora(model, cfg):
        PeftModel,
    )

-    bits = 4 if cfg.load_in_4bits else 8 if cfg.load_in_8bits else None
+    bits = None
+    if cfg.cfg.load_in_4bits: 
+        bits = 4
+    elif cfg.cfg.load_in_8bits: 
+        bits = 8
    linear_names = find_all_linear_names(bits, model)
    logging.info(f"found linear modules: {repr(linear_names)}")
    lora_target_modules = cfg.lora_target_modules + linear_names