From ada603746389158b4e068cd0dd44c3df93e02a12 Mon Sep 17 00:00:00 2001
From: Victor Hall <victor.charles.hall@gmail.com>
Date: Thu, 21 Sep 2023 13:47:26 -0400
Subject: [PATCH] add bnb lion8bit support

---
 optimizer/optimizers.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/optimizer/optimizers.py b/optimizer/optimizers.py
index 7eb0fe5..7d1c273 100644
--- a/optimizer/optimizers.py
+++ b/optimizer/optimizers.py
@@ -307,6 +307,18 @@ class EveryDreamOptimizer():
                     betas=(betas[0], betas[1]),
                     weight_decay=weight_decay,
                 )
+            elif optimizer_name == "lion8bit":
+                from bitsandbytes.optim import Lion8bit
+                optimizer = opt_class(
+                    itertools.chain(parameters),
+                    lr=curr_lr,
+                    betas=(betas[0], betas[1]),
+                    weight_decay=weight_decay,
+                    percentile_clipping=100,
+                    min_8bit_size=4096,
+                    d0=d0,
+                    log_every=args.log_step,
+                )
             elif optimizer_name == "adamw":
                 opt_class = torch.optim.AdamW
             if "dowg" in optimizer_name: