Skip to content

this edit makes it so that ai toolkit can now accept user-defined optimizers such as schedule free and paged optimizers. - update optimizer.py - #446

Closed
Epispasm wants to merge 1 commit into
ostris:mainfrom
Epispasm:patch-2
Closed

Epispasm wants to merge 1 commit into
ostris:mainfrom
Epispasm:patch-2

Conversation

@Epispasm

@Epispasm Epispasm commented Oct 1, 2025 •

Copy link
Copy Markdown

this edit makes it so that ai toolkit can now accept user-defined optimizers.

see changes to jobs/process/BaseSDTrainProcess.py which is needed in order for this to work. It simply edits a redundant lower() function so that user input case can be preserved for the dynamic import.

this will still allow using the same optimizers you had hard-coded so that it doesn't break current configs.

I have run several short tests with SDXL models to make sure it is training as expected and seems to work fully as expected.

Usage:
The user enters the full module and class name of the optimizer they want to use.
Examples:

  • "prodigyplus.prodigy_plus_schedulefree.ProdigyPlusScheduleFree" (must be installed in the venv first - pip install prodigy-plus-schedule-free)
  • "bitsandbytes.optim.PagedAdamW" - this is already installed with bitsandbytes. a great optimizer for getting around OOM errors. It uses paging to offload some of the load to the CPU without losing much if any training speed and no real quality loss. IMO it is faster than block swapping.
  • "bitsandbytes.optim.PagedAdamW8bit"
  • "bitsandbytes.optim.PagedLion"

Also would probably work just fine with other schedule free optimizers.

see changes to jobs/process/BaseSDTrainProcess.py which is needed in order for this to work.  It simply edits a redundant lower() function so that user input case can be preserved for the dynamic import.

this edit makes it so that ai toolkit can now accept user-defined optimizers.  

It will still use the same optimizers you had hard-coded so that it doesn't break current configs.
@Epispasm Epispasm changed the title Update optimizer.py this edit makes it so that ai toolkit can now accept user-defined optimizers. - update optimizer.py Oct 3, 2025
@Epispasm Epispasm changed the title this edit makes it so that ai toolkit can now accept user-defined optimizers. - update optimizer.py this edit makes it so that ai toolkit can now accept user-defined optimizers such as schedule free and paged optimizers. - update optimizer.py Oct 3, 2025
@Epispasm Epispasm closed this Dec 23, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant