From e8e71cf20d31396accc91ec6ee83f80bb08df02f Mon Sep 17 00:00:00 2001
From: youngmae <youngmae>
Date: Fri, 13 Dec 2024 15:14:41 -0800
Subject: [PATCH 1/3] add 8bit optimizer support & documentation

---
 README.md                            | 10 ++++++++++
 setup.py                             |  1 +
 stable_audio_tools/training/utils.py | 17 ++++++++++++-----
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index a8f5ce3f..9241b41a 100644
--- a/README.md
+++ b/README.md
@@ -149,6 +149,16 @@ The following properties are defined in the top level of the model configuration
 - `training`
   - The training configuration for the model, varies based on `model_type`. Provides parameters for training as well as demos.
 
+
+### Optimizer config
+The model config file defines all of the information needed to load a model for training or inference. It also contains the training configuration needed to fine-tune a model or train from scratch.
+
+- `backend`
+  - The type of optimizer library being used, currently limited to one of `"bnb", "default"`.
+- `type`
+  - Optimizer name to use. If using bnb, enabled the use of `"AdamW8bit"` and other 8bit optimizers.
+
+
 ## Dataset config
 `stable-audio-tools` currently supports two kinds of data sources: local directories of audio files, and WebDataset datasets stored in Amazon S3. More information can be found in [the dataset config documentation](docs/datasets.md)
 
diff --git a/setup.py b/setup.py
index 7e7470d3..929bfce5 100644
--- a/setup.py
+++ b/setup.py
@@ -11,6 +11,7 @@
         'aeiou==0.0.20',
         'alias-free-torch==0.0.6',
         'auraloss==0.4.0',
+        'bitsandbytes==0.35.0',
         'descript-audio-codec==1.0.0',
         'einops==0.7.0',
         'einops-exts==0.0.4',
diff --git a/stable_audio_tools/training/utils.py b/stable_audio_tools/training/utils.py
index 38a3fccc..48f60be7 100644
--- a/stable_audio_tools/training/utils.py
+++ b/stable_audio_tools/training/utils.py
@@ -84,13 +84,20 @@ def create_optimizer_from_config(optimizer_config, parameters):
     """
 
     optimizer_type = optimizer_config["type"]
+    optimizer_backend = optimizer_config.get("backend", "")
 
-    if optimizer_type == "FusedAdam":
-        from deepspeed.ops.adam import FusedAdam
-        optimizer = FusedAdam(parameters, **optimizer_config["config"])
-    else:
-        optimizer_fn = getattr(torch.optim, optimizer_type)
+    if optimizer_backend == "bnb":
+        import bitsandbytes as bnb
+        optimizer_fn = getattr(bnb.optim, optimizer_type)
         optimizer = optimizer_fn(parameters, **optimizer_config["config"])
+    else:
+        if optimizer_type == "FusedAdam":
+            from deepspeed.ops.adam import FusedAdam
+            optimizer = FusedAdam(parameters, **optimizer_config["config"])
+        else:
+            optimizer_fn = getattr(torch.optim, optimizer_type)
+            optimizer = optimizer_fn(parameters, **optimizer_config["config"])
+            
     return optimizer
 
 def create_scheduler_from_config(scheduler_config, optimizer):

From 32f2930fed4bc83499d6eda11e1d36bfe3907e79 Mon Sep 17 00:00:00 2001
From: youngmae <young.maeng@gmail.com>
Date: Fri, 13 Dec 2024 15:16:23 -0800
Subject: [PATCH 2/3] Update README.md

optimizer config readme
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 9241b41a..9feb0b6d 100644
--- a/README.md
+++ b/README.md
@@ -151,7 +151,7 @@ The following properties are defined in the top level of the model configuration
 
 
 ### Optimizer config
-The model config file defines all of the information needed to load a model for training or inference. It also contains the training configuration needed to fine-tune a model or train from scratch.
+The optimizer config file allows for use of different optimizer implementations, including those that allow for fine tuning with 24GB VRAM.
 
 - `backend`
   - The type of optimizer library being used, currently limited to one of `"bnb", "default"`.

From 5a790887ed0d4750798de4e828f391a367db6425 Mon Sep 17 00:00:00 2001
From: youngmae <young.maeng@gmail.com>
Date: Fri, 13 Dec 2024 15:17:22 -0800
Subject: [PATCH 3/3] Update README.md

updated location
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 9feb0b6d..938bf464 100644
--- a/README.md
+++ b/README.md
@@ -151,7 +151,7 @@ The following properties are defined in the top level of the model configuration
 
 
 ### Optimizer config
-The optimizer config file allows for use of different optimizer implementations, including those that allow for fine tuning with 24GB VRAM.
+The optimizer config, inside of the training subsection of the model config, allows for use of different optimizer implementations, including those that allow for fine tuning with 24GB VRAM.
 
 - `backend`
   - The type of optimizer library being used, currently limited to one of `"bnb", "default"`.