Mirrors
/
EveryDream2trainer
mirror of https://github.com/victorchall/EveryDream2trainer.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add bnb lion8bit support

This commit is contained in:
Victor Hall 2023-09-21 13:47:26 -04:00
parent a47d65799f
commit ada6037463
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
optimizer/optimizers.py
View File

@ -307,6 +307,18 @@ class EveryDreamOptimizer():
betas=(betas[0], betas[1]), betas=(betas[0], betas[1]),
weight_decay=weight_decay, 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": elif optimizer_name == "adamw":
opt_class = torch.optim.AdamW opt_class = torch.optim.AdamW
if "dowg" in optimizer_name: if "dowg" in optimizer_name: