Support new optimizer Schedule free

[sdbds]

From facebookresearch,no need set any schedule
https://github.com/facebookresearch/schedule_free

For SGD, a learning rate 10x-50x larger than classical rates seems to be a good starting point.
For AdamW, learnings rates in the range 1x-10x larger than with schedule based approaches seem to work.

Do not use lr_schedule_warmup and with it in optimizer args
warmup_steps
warmup_steps (int): Enables a linear learning rate warmup (default 0).

other optimizer args for this

r (float): Use polynomial weighting in the average 
            with power r (default 0).
weight_lr_power (float): During warmup, the weights in the average will
            be equal to lr raised to this power. Set to 0 for no weighting
            (default 2.0)

Dependences

pip install schedulefree

Need lastest accelerate
pip install git+https://github.com/huggingface/accelerate.git@main#egg=accelerate

Fix TYPO in library/train_util.py line 3090
grandient to gradient

[kohya-ss]

Thank you for this! I'll merge it into dev or main once accelerate supports this feature in the release version!

[feffy380]

This implementation is not correct. optimizer.train() needs to be called somewhere before the training loop, but this PR blindly inserts it after every unet.train(), which is not correct (especially in train_network.py where we're not even training the unet). It's most obvious in the textual inversion scripts where optimizer.train() is never called at all.

I think you also need to call optimizer.eval() before saving the model because calling train() and eval() modifies the weights. Based on the example script, we need to save the eval() version of the weights.

And the "support" added by accelerate is just a convenience function that passes the train()/eval() calls to the internal optimizer object. You can call optimizer.optimizer.train() directly instead of waiting for accelerate to update.

[gesen2egee]

Edit
I see the problem now.
facebookresearch/schedule_free#5

===

Currently running will result in such an error.
accelerate 0.30.0.dev0
schedulefree 1.2.1

steps: 0%| | 0/3000 [00:00<?, ?it/s]
epoch 1/9
warnings.warn(
Traceback (most recent call last):
File "D:\SDXL\sd-scripts\train_network.py", line 1128, in
trainer.train(args)
File "D:\SDXL\sd-scripts\train_network.py", line 921, in train
optimizer.step()
File "D:\SDXL\sd-scripts\venv\lib\site-packages\accelerate\optimizer.py", line 148, in step
self.scaler.step(self.optimizer, closure)
File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\cuda\amp\grad_scaler.py", line 416, in step
retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)
File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\cuda\amp\grad_scaler.py", line 315, in _maybe_opt_step
retval = optimizer.step(*args, **kwargs)
File "D:\SDXL\sd-scripts\venv\lib\site-packages\accelerate\optimizer.py", line 203, in patched_step
return method(*args, **kwargs)
File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\optim\lr_scheduler.py", line 68, in wrapper
return wrapped(*args, **kwargs)
File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\optim\optimizer.py", line 373, in wrapper
out = func(*args, **kwargs)
File "D:\SDXL\sd-scripts\venv\lib\site-packages\schedulefree\adamw_schedulefree.py", line 122, in step
ckp1 = weight/weight_sum
ZeroDivisionError: float division by zero
steps: 0%| | 0/3000 [00:36<?, ?it/s]

[sdbds]

This implementation is not correct. optimizer.train() needs to be called somewhere before the training loop, but this PR blindly inserts it after every unet.train(), which is not correct (especially in train_network.py where we're not even training the unet). It's most obvious in the textual inversion scripts where optimizer.train() is never called at all.

I think you also need to call optimizer.eval() before saving the model because calling train() and eval() modifies the weights. Based on the example script, we need to save the eval() version of the weights.

And the "support" added by accelerate is just a convenience function that passes the train()/eval() calls to the internal optimizer object. You can call optimizer.optimizer.train() directly instead of waiting for accelerate to update.

Thank you for reviewing！
I did make some mistakes, but I'd like to say what I think.
1、I change sdxl_train.py first, it use training_models to contain unet and textencoder if we choose.
And loop i had insert is in every epoch start and before train loop start, it is no problem.
optimizer.train() just switch optimizer status when train or eval.(and default use train model)

(OK，For greater clarity, I will put it before loading train_model loop just once)

2、i actually forget to add eval before save models.
i will add it in before every save models and every sample images

3、i missed TI...ill add it soon

4、I know convenience function about accelerate because that updates is my request, it can keep completion about code.

[sdbds]

Edit I see the problem now. facebookresearch/schedule_free#5

===

Currently running will result in such an error. accelerate 0.30.0.dev0 schedulefree 1.2.1

steps: 0%| | 0/3000 [00:00<?, ?it/s] epoch 1/9 warnings.warn( Traceback (most recent call last): File "D:\SDXL\sd-scripts\train_network.py", line 1128, in trainer.train(args) File "D:\SDXL\sd-scripts\train_network.py", line 921, in train optimizer.step() File "D:\SDXL\sd-scripts\venv\lib\site-packages\accelerate\optimizer.py", line 148, in step self.scaler.step(self.optimizer, closure) File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\cuda\amp\grad_scaler.py", line 416, in step retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs) File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\cuda\amp\grad_scaler.py", line 315, in _maybe_opt_step retval = optimizer.step(*args, **kwargs) File "D:\SDXL\sd-scripts\venv\lib\site-packages\accelerate\optimizer.py", line 203, in patched_step return method(*args, **kwargs) File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\optim\lr_scheduler.py", line 68, in wrapper return wrapped(*args, **kwargs) File "D:\SDXL\sd-scripts\venv\lib\site-packages\torch\optim\optimizer.py", line 373, in wrapper out = func(*args, **kwargs) File "D:\SDXL\sd-scripts\venv\lib\site-packages\schedulefree\adamw_schedulefree.py", line 122, in step ckp1 = weight/weight_sum ZeroDivisionError: float division by zero steps: 0%| | 0/3000 [00:36<?, ?it/s]

set lr_scheduler=null and not use lr_warmup will fix this

[sdbds]

I add eval() in train loop because args.save_every_n_steps so maybe anysteps save in loop.
i add it before save and after train.

at the time i had to move train() to train loop before train start for eval() in loop maybe refresh train() status

Add Ti part.

[sdbds]

Thank you for this! I'll merge it into dev or main once accelerate supports this feature in the release version!

@kohya-ss

accelerate release newest version to support schedule free,and this pr had checked and fixed some mistakes now.

https://github.com/huggingface/accelerate/releases/tag/v0.29.2

[kohya-ss]

Thank you for update! I will review and merge as soon as I have time. I will also consider whether the code can be made common.

[gesen2egee]

args.optimizer_type.lower().endswith("scheduleFree"):
I'm not sure if this is a problem.

[sdbds]

args.optimizer_type.lower().endswith("scheduleFree"): I'm not sure if this is a problem.

It doesn't really matter, but I modified it for aesthetic purposes

[feffy380]

It does matter because optimizer_type.lower() is all lowercase and will never match scheduleFree

[rockerBOO]

When it goes to do sample_images it should also set the optimizer to eval as well? I see big spikes in the loss after generating samples.

[sdbds]

When it goes to do sample_images it should also set the optimizer to eval as well? I see big spikes in the loss after generating samples.

Actually eval() switching is before saving, sampler is processed after saving.
And train() is before the start of each training step.
I guess it change because from eval() sampler to train().

[rockerBOO]

Actually eval() switching is before saving, sampler is processed after saving. And train() is before the start of each training step. I guess it change because from eval() sampler to train().

Ahh my bad, I see that now. Must be something odd with my logging.

[kohya-ss]

I merged this into scheduler-free-opt branch. Sorry for the delay.

I also made some modifications to simplify the code. I'll do some more testing and plan to merge it into the dev branch.