ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-29 02:31:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added support for AdEMAMix8bit

Browse Source
This commit is contained in:
Jaret Burkett
2024-09-28 14:33:51 -06:00
parent a508caad1d
commit 69aa92bce5
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
toolkit/optimizer.py
View File

@@ -46,6 +46,8 @@ def get_optimizer(
if lower_type == "adam8bit": if lower_type == "adam8bit":
return bitsandbytes.optim.Adam8bit(params, lr=learning_rate, eps=1e-6, **optimizer_params) 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": elif lower_type == "adamw8bit":
return bitsandbytes.optim.AdamW8bit(params, lr=learning_rate, eps=1e-6, **optimizer_params) return bitsandbytes.optim.AdamW8bit(params, lr=learning_rate, eps=1e-6, **optimizer_params)
elif lower_type == "lion8bit": elif lower_type == "lion8bit":