Added Lion optimizer

1 files changed, 27 insertions, 11 deletions

diff --git a/train_ti.py b/train_ti.py
index 3aa1027..507d710 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -290,9 +290,10 @@ def parse_args():
         default=0.9999
     )
     parser.add_argument(
-        "--use_8bit_adam",
-        action="store_true",
-        help="Whether or not to use 8-bit Adam from bitsandbytes."
+        "--optimizer",
+        type=str,
+        default="lion",
+        help='Optimizer to use ["adam", "adam8bit", "lion"]'
     )
     parser.add_argument(
         "--adam_beta1",
@@ -564,15 +565,34 @@ def main():
         args.learning_rate = 1e-5
         args.lr_scheduler = "exponential_growth"
 
-    if args.use_8bit_adam:
+    if args.optimizer == 'adam8bit':
         try:
             import bitsandbytes as bnb
         except ImportError:
             raise ImportError("To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`.")
 
-        optimizer_class = bnb.optim.AdamW8bit
+        create_optimizer = partial(
+            bnb.optim.AdamW8bit,
+            betas=(args.adam_beta1, args.adam_beta2),
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            amsgrad=args.adam_amsgrad,
+        )
+    elif args.optimizer == 'adam':
+        create_optimizer = partial(
+            torch.optim.AdamW,
+            betas=(args.adam_beta1, args.adam_beta2),
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            amsgrad=args.adam_amsgrad,
+        )
     else:
-        optimizer_class = torch.optim.AdamW
+        try:
+            from lion_pytorch import Lion
+        except ImportError:
+            raise ImportError("To use Lion, please install the lion_pytorch library: `pip install lion_pytorch`.")
+
+        create_optimizer = partial(Lion, use_triton=True)
 
     checkpoint_output_dir = output_dir/"checkpoints"
 
@@ -658,13 +678,9 @@ def main():
         )
         datamodule.setup()
 
-        optimizer = optimizer_class(
+        optimizer = create_optimizer(
             text_encoder.text_model.embeddings.temp_token_embedding.parameters(),
             lr=args.learning_rate,
-            betas=(args.adam_beta1, args.adam_beta2),
-            weight_decay=args.adam_weight_decay,
-            eps=args.adam_epsilon,
-            amsgrad=args.adam_amsgrad,
         )
 
         lr_scheduler = get_scheduler(