From 69aa92bce55b1b0cc72bc547df73a72be3dcad0d Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sat, 28 Sep 2024 14:33:51 -0600 Subject: [PATCH] Added support for AdEMAMix8bit --- toolkit/optimizer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/toolkit/optimizer.py b/toolkit/optimizer.py index d2a1e925..4ec38501 100644 --- a/toolkit/optimizer.py +++ b/toolkit/optimizer.py @@ -46,6 +46,8 @@ def get_optimizer( if lower_type == "adam8bit": return bitsandbytes.optim.Adam8bit(params, lr=learning_rate, eps=1e-6, **optimizer_params) + if lower_type == "ademamix8bit": + return bitsandbytes.optim.AdEMAMix8bit(params, lr=learning_rate, eps=1e-6, **optimizer_params) elif lower_type == "adamw8bit": return bitsandbytes.optim.AdamW8bit(params, lr=learning_rate, eps=1e-6, **optimizer_params) elif lower_type == "lion8bit":