Training Stable Cascade on cheap and old GPUs

I am liking Stable Cascade so far. I was concerned when I heard it was based on wuerstchen, which in my experience produced disappointing results, but SAI said it performed better than their SDXL, so I gave it a shot and it looks quite good. In general, its similar in quality to SDXL now, but it seems to have less of an issue scaling up than SDXL does. I have no problem generating 2048x2048 images on my machine, and they look great.

It does have some weaknesses though. Not much community development has been done yet, diffusers barely supports it, and loading LoRAs is not easy. Training is more complicated, and the official models only support bf16 or f32, which means you need either a newer GPU or a ton of VRAM (which probably means you have a newer GPU anyway). You can't even run inference in f16 mode normally. There are tiny versions of the models with less than a third of the parameters that can be trained with much less VRAM, but who whats to run "lite" versions - lame.

So without LoRAs and without nice community fine-tuned models and without easy training on the full models, Stable Cascade seems a lot less attractive.

What can we do?

Turns out there are a couple things we can do to get inference working in f16 mode. I searched around and found KBlueLeaf on huggingface had inference in f16 working by manually scaling down the weights that cause f16 overflows. There is a repo with code (although its very ugly) and models with the scaling applied. Either using the model or the code allows inference in f16 mode.

The other option is to just do most of the inference in f16, but do the ResBlocks in f32. This works, and the modification is a lot simpler (a single line of code), but of course this results in significantly higher VRAM usage (although still reasonable).

Great, but this doesn't solve the problem, we need to do training.

KBlueLeaf's method doesn't work for training. I tried starting from the KBlueLeaf f16 model, but the losses go to NaN after 15-20 steps. I tried applying the scaling logic to the training scripts, and this looked like it was working at first, but the losses still go to NaN after 50-80 steps.

I tried training in f32 mode, but my GPUs have nowhere near the VRAM needed for that.

I tried integrating xformers into the training scripts to see if that would allow me to get under my VRAM budget, but everything I tried said my GPU wasn't supported.

I tried just doing the ResBlocks in f32 and the rest in f16, but this still exceeds my VRAM.

I tried removing all references to the previewer from the training code, but the previewer model is tiny, and removing it doesn't help much.

Then, I implemented gradient checkpointing on the ResBlocks, along with training them in f32 and the rest of the model in f16. This prevents the losses from going to NaN, and the training script runs and outputs checkpoints! Although the gradient checkpointing adds 10 hours to training time (10000 steps), long training is better than no training.

However, these checkpoints are trash. They produce garbage, noise, goofy colors. In the early checkpoints, sometimes you could see something reminiscent of the correct image, and sometimes very many inference steps could make it look okay-ish, but whats the point of training if it makes the results worse?

I figure that training with ResBlocks in f32 allows them to exceed the range of f16 or even bf16, so if we saved the weights as f32 and did inference in f32 it would work, but I was saving my weights as f16.

ugh... very disappointing...

Then I remembered LoRA. I could train a LoRA with this setup and it wouldn't be a problem because LoRAs only apply to the attention layers and not the ResBlocks. But, so far, loading LoRAs is not easy, and my diffusers based inference scripts use slightly different names than the training scripts I was using, so I would have a hard time writing the code to load the resulting LoRAs.

But... what if I combined these? What if I train the whole model, but knowing that I can get results without training the ResBlocks, I just don't train the ResBlocks?

So that's what I did. I implemented gradient checkpointing, cast the ResBlocks to f32, and don't calculate losses for ResBlocks at all, and then train the rest of the model as f16.

Surprisingly, this not only works, but also cuts down on VRAM usage and training time significantly! Training time is reduced to 20-25% of what it would take to train the full model in f16, and I'm only using ~10GB VRAM total. That's low enough to enable training on lots of cheap consumer GPUs.

But how does it perform?

TBD... Seems like it's learning something along the lines of my training data, but I have a long way to go before I can say if I am happy with the results or not.