Add support for AdEMAMix8bit optimizer

bmaltais · Sep 29, 2024 · 34db05e · 34db05e
1 parent 79a4a3a
commit 34db05e
Show file tree

Hide file tree

Showing 7 changed files with 7 additions and 6 deletions.
diff --git a/kohya_gui/class_basic_training.py b/kohya_gui/class_basic_training.py
@@ -196,6 +196,8 @@ def init_lr_and_optimizer_controls(self) -> None:
                     "AdamWScheduleFree",
                     "AdamW8bit",
                     "Adafactor",
+                    "bitsandbytes.optim.AdEMAMix8bit",
+                    "bitsandbytes.optim.PagedAdEMAMix8bit",
                     "DAdaptation",
                     "DAdaptAdaGrad",
                     "DAdaptAdam",

diff --git a/requirements.txt b/requirements.txt
@@ -34,6 +34,5 @@ toml==0.10.2
 transformers==4.44.2
 voluptuous==0.13.1
 wandb==0.18.0
-scipy==1.11.4
 # for kohya_ss library
 -e ./sd-scripts # no_verify leave this to specify not checking this a verification stage
diff --git a/requirements_linux.txt b/requirements_linux.txt
@@ -1,5 +1,5 @@
 torch==2.4.0+cu124 torchvision==0.19.0+cu124 --extra-index-url https://download.pytorch.org/whl/cu124
-bitsandbytes==0.43.3
+bitsandbytes==0.44.0
 tensorboard==2.15.2 tensorflow==2.15.0.post1
 onnxruntime-gpu==1.17.1
 xformers==0.0.27.post2

diff --git a/requirements_linux_docker.txt b/requirements_linux_docker.txt
@@ -1,4 +1,4 @@
 xformers>=0.0.20
-bitsandbytes==0.43.3
+bitsandbytes==0.44.0
 accelerate==0.33.0
 tensorboard
diff --git a/requirements_runpod.txt b/requirements_runpod.txt
@@ -1,5 +1,5 @@
 torch==2.4.0+cu124 torchvision==0.19.0+cu124 --extra-index-url https://download.pytorch.org/whl/cu124 # no_verify leave this to specify not checking this a verification stage
-bitsandbytes==0.43.3
+bitsandbytes==0.44.0
 tensorboard==2.14.1 tensorflow==2.14.0 wheel
 tensorrt
 onnxruntime-gpu==1.17.1

diff --git a/requirements_windows.txt b/requirements_windows.txt
@@ -1,4 +1,4 @@
-bitsandbytes==0.43.3
+bitsandbytes==0.44.0
 tensorboard
 tensorflow>=2.16.1
 onnxruntime-gpu==1.17.1

diff --git a/sd-scripts b/sd-scripts