Adafactor fused backward pass and optimizer step, lowers SDXL (@ 1024 resolution) VRAM usage to BF16(10GB)/FP32(16.4GB)

[2kpr]

Implemented a 'fused' backward pass and optimizer step when using Adafactor which massively lowers the VRAM used when training. If you want some more information on 'fusing' see here: https://pytorch.org/tutorials/intermediate/optimizer_step_in_backward_tutorial.html

For instance with SDXL (@ 1024 resolution), the VRAM used when training in BF16 drops to 10GB of VRAM, and in FP32 drops to 16.4GB of VRAM, allowing SDXL FP32 'normal' training on a 24GB GPU (which allows a batch size of up to about 16 for BF16 training and 6 for FP32 training when training on a 24GB GPU). Secondarily it allows SDXL 'normal' BF16 training on smaller 12GB/16GB GPUs, etc.

I only modified the sdxl_train.py so far so people could test out this PR and then make any changes necessary before extending this 'fused backward pass' to the other training py files in this repo, etc.

Feel free to edit accordingly, I just placed things where they seemed most logical according to the existing code in this repo.

[FurkanGozukara]

wow nice. yes i use it with onetrainer and 10.3 GB VRAM required for SDXL training

even published a video

Full Stable Diffusion SD & XL Fine Tuning Tutorial With OneTrainer On Windows & Cloud - Zero To Hero

[araleza]

Interesting! Do you know if this is a pure memory saving, or if it changes the gradients at all? Although even if it just saves memory with no consequence, that would be great, especially if Stable Diffusion 3 is round the corner.

I'm interested to see if fp32 training of SDXL has any benefits over bf16. At the very least, it'll stop the text encoder getting quantized to bf16 as the first step, which makes people's faces look a bit weird for a few iterations til the model rebalances itself to deal with the lost precision. But that's just a temporary effect. I hacked my own code to leave the text encoder as fp32 already, even when using full bf16 training to avoid this, but I wouldn't have to do that anymore with this update of yours.

@FurkanGozukara, did you notice any quality improvements to the final results, or just reduced memory?

[2kpr]

Interesting! Do you know if this is a pure memory saving, or if it changes the gradients at all? Although even if it just saves memory with no consequence, that would be great, especially if Stable Diffusion 3 is round the corner.

Pure memory, see: https://pytorch.org/tutorials/intermediate/optimizer_step_in_backward_tutorial.html

I'm interested to see if fp32 training of SDXL has any benefits over bf16. At the very least, it'll stop the text encoder getting quantized to bf16 as the first step, which makes people's faces look a bit weird for a few iterations til the model rebalances itself to deal with the lost precision. But that's just a temporary effect. I hacked my own code to leave the text encoder as fp32 already, even when using full bf16 training to avoid this, but I wouldn't have to do that anymore with this update of yours.

Overall the FP32 is much better.

There are known training issues with BF16, and techniques like 'stochastic rounding' that 'attempt to fix' the problems with training in BF16 by shifting the results nearer to that of having trained in FP32 instead. See here: https://storage.googleapis.com/deepmind-media/research/language-research/Training%20Gopher.pdf , check out pages 53 (starting at setion C.2) and 54 , and also this paper if you want to read about BF16 and 'stochastic rounding' a bit more: https://arxiv.org/pdf/2010.06192.pdf

In a future PR to this repo the 'stochastic rounding' can be added, which should allow those that can 'only' train in BF16 (due to VRAM limits, which this PR now opens the window up of letting even 12GB GPUs train 'normal' SDXL) to have their results nearer to the results attained when training in FP32, which I have tested on my own end and the results are much better than 'normal' BF16 training, but still a bit shy of 'actual' FP32 trainings (which are now possible on 24GB GPUs via this PR with the 'fused' backward pass and optimizer step).

[araleza]

Alright, you've convinced me. 😊 I'll grab your branch and see if I can notice any better results on some of my existing image sets.

[FurkanGozukara]

Interesting! Do you know if this is a pure memory saving, or if it changes the gradients at all? Although even if it just saves memory with no consequence, that would be great, especially if Stable Diffusion 3 is round the corner.

I'm interested to see if fp32 training of SDXL has any benefits over bf16. At the very least, it'll stop the text encoder getting quantized to bf16 as the first step, which makes people's faces look a bit weird for a few iterations til the model rebalances itself to deal with the lost precision. But that's just a temporary effect. I hacked my own code to leave the text encoder as fp32 already, even when using full bf16 training to avoid this, but I wouldn't have to do that anymore with this update of yours.

@FurkanGozukara, did you notice any quality improvements to the final results, or just reduced memory?

no changes only reduced speed

[FurkanGozukara]

@2kpr fp32 cooked my model with same LR of BF16

have you made any changes? i weren't able to research it yet

[kohya-ss]

Thank you for this. This looks really nice! I will review and merge this sooner.

[araleza]

I've been running tests for a few hours now, and I'm seeing great improvements! Furkan is correct that the iterations go significantly more slowly, and so does the generation of sample images. But, I feel I'm also seeing good gains in image quality and learning ability. Fine-grained details that were being missed out on are being learned this time, which they weren't when I used bf16. And the images seem to stay 'in focus' better, without smudging so often.

This is a recommended update.

[2kpr]

@2kpr fp32 cooked my model with same LR of BF16

have you made any changes? i weren't able to research it yet

Not that I'm aware of.

If you consider what is happening you might see why it would 'appear to cook' your model at that same learning rate as BF16, it is not that your model is being 'cooked' it's more like you are seeing a 'real learning' happening since FP32 is retaining all of the model update information whereas when you train with BF16 you are losing a lot of information due to the 'limited precision' of BF16.

BF16 has the 'range' to match FP32, but BF16 has way less 'precision' than FP32 and therefore does not 'properly update the model' on each step when training in BF16, that is what 'stochastic rounding' tries to fix for instance, see here: https://storage.googleapis.com/deepmind-media/research/language-research/Training%20Gopher.pdf , check out pages 53 (starting at setion C.2) and 54.

So, it isn't that this code or FP32 'cooked' your model, it is more apt to say that BF16 'undertrained' / 'undercooked' your model at the same LR and that BF16 was 'losing a lot of information during model updates' due to it's very limited 'precision'.

Makes sense?

[2kpr]

Thank you for this. This looks really nice! I will review and merge this sooner.

NP, thanks, awesome, feel free to move things around to more 'fit in' with your existing codebase :)

[2kpr]

I've been running tests for a few hours now, and I'm seeing great improvements! Furkan is correct that the iterations go significantly more slowly, and so does the generation of sample images. But, I feel I'm also seeing good gains in image quality and learning ability. Fine-grained details that were being missed out on are being learned this time, which they weren't when I used bf16. And the images seem to stay 'in focus' better, without smudging so often.

This is a recommended update.

Awesome, thanks for the feedback :)

Well, yes, that is just what happens when training in FP32 vs BF16 unfortunately, goes a bit slower re it/s, but the fact that the model is not losing information in updates is key, see here: #1259 (comment)

[araleza]

I'm finding I can set the LR significantly lower when using fp32 training. Previously, I found that cranking the LR up actually got me better generalization, as lower LRs seemed to get stuck in local minimas. But I guess it's just that the weight updates to get out of these (apparent) minimas were too small for bf16 to represent.

[2kpr]

I'm finding I can set the LR significantly lower when using fp32 training. Previously, I found that cranking the LR up actually got me better generalization, as lower LRs seemed to get stuck in local minimas. But I guess it's just that the weight updates to get out of these (apparent) minimas were too small for bf16 to represent.

Cool, yea, that is a great way to put it :)

[CCpt5]

@2kpr fp32 cooked my model with same LR of BF16

have you made any changes? i weren't able to research it yet

You need to go w/ a much slower LR if using stochastic rounding or fp32 as numbers that were being rounded to 0 the prior method now have small values. I usually use 1e-5 for dreambooth training. In my handful of tests w/ stochastic and a couple w/ fp32 1e-6 is where I need to keep it to avoid cooking the model as you say.

Run samples every 100 steps or so the and you'll see it converges much more quickly.

[araleza]

Hi @2kpr, I think there is a bug somewhere in this implementation. My training runs start okay, but if I get to the first autosave point, and then stop the run and restart it from that autosave (instead of sdxl base), I get this error message:

Traceback (most recent call last):
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/./sdxl_train.py", line 835, in
train(args)
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/./sdxl_train.py", line 635, in train
accelerator.backward(loss)
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 1905, in backward
loss.backward(**kwargs)
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/venv/lib/python3.10/site-packages/torch/_tensor.py", line 522, in backward
torch.autograd.backward(
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/venv/lib/python3.10/site-packages/torch/autograd/init.py", line 266, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/venv/lib/python3.10/site-packages/torch/autograd/function.py", line 289, in apply
return user_fn(self, *args)
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/venv/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 319, in backward
torch.autograd.backward(outputs_with_grad, args_with_grad)
File "/home/ara/m.2/Dev/sdxl/sd-scripts_fp32/venv/lib/python3.10/site-packages/torch/autograd/init.py", line 266, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [1280, 1280]], which is output 0 of AsStridedBackward0, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

I'm not saving out in fp16 or bf16 format, I switched to saving out in fp32.

[araleza]

@2kpr - an update on the previous bug report: I only get that error message if I have --train_text_encoder set in the command line on the second run. That parameter works when starting from base, but not from an autosave checkpoint.

[stepfunction83]

I'm getting the following error upon attempting to begin a training run:

Traceback (most recent call last):
  File "/home/myname/Stable-Diffusion/kohya_ss/sd-scripts/sdxl_train.py", line 835, in <module>
    train(args)
  File "/home/myname/Stable-Diffusion/kohya_ss/sd-scripts/sdxl_train.py", line 635, in train
    accelerator.backward(loss)
  File "/home/myname/Stable-Diffusion/kohya_ss/venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 1903, in backward
    self.scaler.scale(loss).backward(**kwargs)
  File "/home/myname/Stable-Diffusion/kohya_ss/venv/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/home/myname/Stable-Diffusion/kohya_ss/venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/myname/Stable-Diffusion/kohya_ss/sd-scripts/sdxl_train.py", line 441, in __grad_hook
    accelerator.clip_grad_norm_(tensor, args.max_grad_norm)
  File "/home/myname/Stable-Diffusion/kohya_ss/venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 2040, in clip_grad_norm_
    self.unscale_gradients()
  File "/home/myname/Stable-Diffusion/kohya_ss/venv/lib/python3.10/site-packages/accelerate/accelerator.py", line 2003, in unscale_gradients
    self.scaler.unscale_(opt)
  File "/home/myname/Stable-Diffusion/kohya_ss/venv/lib/python3.10/site-packages/torch/cuda/amp/grad_scaler.py", line 296, in unscale_
    raise RuntimeError(
RuntimeError: unscale_() has already been called on this optimizer since the last update().

Update:
Commenting lines 440 and 441 in sdxl_train.py allowed the run to begin:

                        # if accelerator.sync_gradients and args.max_grad_norm != 0.0:
                        #     accelerator.clip_grad_norm_(tensor, args.max_grad_norm)

It seems the current max_grad_norm implementation is not compatible

[sdbds]

Looks great, is it possible to implement it on other optimizers? For example adamW

[stepfunction83]

My understanding of creating stuff from scratch in PyTorch is minimal, so just out of curiosity, why does this require implementing adafactor from scratch instead of using the base implementation and applying the additional steps as described in the documentation?

I guess effectively, why is the change to sdxl_train.py not sufficient here?

[stepfunction83]

After looking through the code further and experimenting to answer my above question, I believe that while the current version of the code is functional, the logic should be made a little more generic to allow for additional fused versions of other optimizers to be incrementally added in addition to adafactor.

This would likely involve:

• Changing the adafactor_fused.py to optimizer_fused.py with the expectation of including additional optimizer patches in it
• Changing sdxl_train.py to reference the more generic optimizer_fused.py
• Changing train_util.py to have a more generic list of accepted fused optimizers which can be used

[stepfunction83]

Since I can't for the life of me figure out how to add a commit to this PR, I created an additional PR on 2kpr's repo with the AdamW optimizer added and the above changes made to make the logic more generic. I couldn't test the logic as it blew past my VRAM when I attempted to run it. It very well might be an issue with my implementation of AdamW, which I copied from the Transformers implementation

PR Created: 2kpr#1

[IdiotSandwichTheThird]

With this PR, Utilizing Adafactor and --fused_backward_pass, --mixed_precision="no" uses less VRAM than --mixed_precision="bf16", which I feel extremely suspicious about. Is anyone else having this happen?

[CCpt5]

With this PR, Utilizing Adafactor and --fused_backward_pass, --mixed_precision="no" uses less VRAM than --mixed_precision="bf16", which I feel extremely suspicious about. Is anyone else having this happen?

Have you tried w/o that mixed_precision line completely? If so your VRAM usage was the same as "no"?

[IdiotSandwichTheThird]

With this PR, Utilizing Adafactor and --fused_backward_pass, --mixed_precision="no" uses less VRAM than --mixed_precision="bf16", which I feel extremely suspicious about. Is anyone else having this happen?

Have you tried w/o that mixed_precision line completely? If so your VRAM usage was the same as "no"?

Here's what that looks like, no appreciable difference between "no" and removing the argument entirely.

[CCpt5]

Here's what that looks like, no appreciable difference between "no" and removing the argument entirely.

Cool, and for clarification, are you running the OPs PR or the PR created by stepfunction83 shared in the comments (2kpr#1)

Thanks!

[IdiotSandwichTheThird]

Here's what that looks like, no appreciable difference between "no" and removing the argument entirely.

Cool, and for clarification, are you running the OPs PR or the PR created by stepfunction83 shared in the comments (2kpr#1)

Thanks!

Running with the OP PR. For the run itself these are the arguments if anyone can't reproduce it:

accelerate launch --num_cpu_threads_per_process 8 sdxl_train.py `
    --logging_dir="logs" --log_prefix="AdaFused" `
    --max_data_loader_n_workers=2 --persistent_data_loader_workers `
    --caption_extension=".txt" `
    --save_precision="bf16" `
    --xformers `
    --save_model_as=safetensors `
    --output_dir="./out" --pretrained_model_name_or_path="./models/model.safetensors" `
    --output_name="AdaFused" `
    --learning_rate=5e-6 `
    --max_train_epochs=20 --save_every_n_epochs=5 `
    --resolution=1024 `
    --train_batch_size=1 `
    --optimizer_type="AdaFactor" `
    --lr_scheduler="constant_with_warmup" --lr_warmup_steps=13 `
    --seed=1337 `
    --optimizer_args "scale_parameter=False" "relative_step=False" "warmup_init=False" --fused_backward_pass `
    --cache_latents --cache_latents_to_disk --vae_batch_size=2 `
    --enable_bucket --gradient_checkpointing --dataset_config=./dataset_config.toml

[stepfunction83]

I can also confirm that I experienced the same thing when I was testing and forgot to turn off full FP16.

[IdiotSandwichTheThird]

With this PR, Utilizing Adafactor and --fused_backward_pass, --mixed_precision="no" uses less VRAM than --mixed_precision="bf16", which I feel extremely suspicious about. Is anyone else having this happen?

Have you tried w/o that mixed_precision line completely? If so your VRAM usage was the same as "no"?

Here's what that looks like, no appreciable difference between "no" and removing the argument entirely.

Upon further inspection, the same behavior occurs in onetrainer with an independent fused adafactor implementation, so I guess that's just how it is ¯\_(ツ)_/¯

[Nerogar]

Upon further inspection, the same behavior occurs in onetrainer with an independent fused adafactor implementation, so I guess that's just how it is ¯_(ツ)_/¯

Not really that independent I think. Looking at the code, it looks like most of it was directly copied from my implementation.

I can also explain why it's done this way, instead of the more generic implementation from here: I didn't want to create a new optimizer for every single parameter. Some optimizers (like prodigy) have some global state that's shared between all parameters. This wouldn't even work with the generic implementation. Also, I feel like there would be a lot of overhead, and saving/loading the optimizer state would also be harder.

@2kpr Copying my code is fine. I made it open source for a reason. But at lease mention it somewhere 😉

btw: This implementation is missing the GradScaler check. Without it, mixed precision training with fp32 weights and fp16 autocasting will degrade a lot.

[2kpr]

@2kpr Copying my code is fine. I made it open source for a reason. But at lease mention it somewhere 😉

To be fair, I didn't copy the code, for had I done so wouldn't the gradscaler been copied? :)

I took the general files / code we had been sharing between ourselves in PM when I first sent my first two stochastic rounding implementations to you, after which you made your very nice and succinct SR implementation. I then took those files and in fact did look at your code in general how you implemented the fused backpass and saw you used the pytorch hook for such which I didn't know existed, after which I read the pytorch article on fusing the backward pass, and then from 'my files' hand wrote what is in this PR, and considering there is only one pytorch hook function, and we both copied/used/broke up the step method from Adafactor, of course the PR 'can look copied', but I assure you it wasn't.

I also had to of course 'try' to integrate it all within the framework of sd-scripts, add the parameters, Japanese translations, and for instance add "if parameter.requires_grad:" since it was erroring otherwise, etc, aka putting in that work that shows it wasn't just some 'copy / paste' operation.

Did I take inspiration from your code, sure, you were after all the first one to mention the fused backpass on discord, but in a similar vein, aka in a similar 'code inspiration sharing between you and me', I sent you my first couple stochastic rounding implementations via PM that you played with and then you developed your own 'simplified' and more 'performant' versions thereof.

That all said, I will fully acknowledge that it was you in particular that recalled on the Stable Cascade discord and let everyone know about the potential benefits of stochastic rounding as a concept due to the old paper you referenced (of which I read, made my own SR implementation attempts, sent them to you, aka what I wrote above, etc), but that you were also the one to suggest the fusing of the backpass as a concept that could be used to great measure in addition to SR, etc :)

So if you take any offense still to my not mentioning of you in this PR, I'm sorry, I hope the above clears things up.

[IdiotSandwichTheThird]

btw: This implementation is missing the GradScaler check. Without it, mixed precision training with fp32 weights and fp16 autocasting will degrade a lot.

Curious. Would this affect mixed precision training in bf16 as well?

[Nerogar]

btw: This implementation is missing the GradScaler check. Without it, mixed precision training with fp32 weights and fp16 autocasting will degrade a lot.

Curious. Would this affect mixed precision training in bf16 as well?

No. The GradScaler only supports (and is only needed) when fp16 is involved.

[kohya-ss]

Thank you for this, and sorry for the delay. This works really well!

I will do some more tests and merge into dev sooner.

[kohya-ss]

I'd like to add 'parameter groups' feature on #1319 to balance the memory savings and complexity of fused backward pass with the availability of an arbitrary optimizer.

Any comments would be appreciated.

[gesen2egee]

In a future PR to this repo the 'stochastic rounding' can be added, which should allow those that can 'only' train in BF16 (due to VRAM limits, which this PR now opens the window up of letting even 12GB GPUs train 'normal' SDXL) to have their results nearer to the results attained when training in FP32, which I have tested on my own end and the results are much better than 'normal' BF16 training, but still a bit shy of 'actual' FP32 trainings (which are now possible on 24GB GPUs via this PR with the 'fused' backward pass and optimizer step).

Hi, I'm particularly interested in the stochastic rounding aspect. Will there be more PRs adding this feature in the future, or has the current PR already implemented it?

Thank you for your work.