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

Added adafactor implementation that handles stochastic rounding of update and accumulation

Browse Source
This commit is contained in:
Jaret Burkett
2024-10-30 05:25:57 -06:00
parent e72b59a8e9
commit 58f9d01c2b
6 changed files with 458 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
jobs/process/BaseSDTrainProcess.py
View File

@@ -1750,7 +1750,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
with torch.no_grad():
# torch.cuda.empty_cache()
if self.train_config.optimizer.lower().startswith('dadaptation') or \
# if optimizer has get_lrs method, then use it
if hasattr(optimizer, 'get_learning_rates'):
learning_rate = optimizer.get_learning_rates()[0]
elif self.train_config.optimizer.lower().startswith('dadaptation') or \
self.train_config.optimizer.lower().startswith('prodigy'):
learning_rate = (
optimizer.param_groups[0]["d"] *