Fine tuning SSD-1B at full 1024x1024 resolution... with only 8GB VRAM?!

It's well known by now that SDXL uses a lot of RAM, especially for training. There's even a distilled version of it called SSD-1B that's designed to be smaller and more memory efficient, though most training software doesn't support it. How much more memory efficient can it be though? How about a full fine tune - not just a LORA but all the UNet weights - using Kohya's scripts at full 1024x1024 resolution in a mere 8 GB of VRAM.

To be clear, this requires a slightly modified version of Kohya's code to add support and doing it in 8 GB is a bit of a stunt. Making this fit requires using basically all the memory saving options available, at the expense of losing some features and likely affecting training quality. 12 GB or even 16 GB is probably a more realistic target.

Also, all the instructions in this guide assume you're using Linux. It might be possible to do this on Windows too, but it would require some changes and probably also more VRAM.

First, we need to set up a modified version of Kohya's training scripts using the patcher attached to this post:

git clone https://github.com/kohya-ss/sd-scripts
cd sd-scripts
git checkout v0.8.4
unzip kohya-ssd-1b-fine-tuning.zip
sh ./patch-sd-scripts-for-ssd-1b.sh

You'll obviously also need to install PyTorch and all the other dependencies by following the standard instructions if you haven't done so already, and to have a suitable data set prepared to use. You'll also need a copy of the base model you plan to use in the single file checkpoint format used by programs like ComfyUI; diffusers format likely won't work. This could be SSD-1B, one of the handful of existing fine tunes based on it, or even my experimental model distilled from Pony Diffusion XL.

This also requires some aggressive memory saving options. Batch size is obviously 1 with gradient checkpointing, though there's enough room for gradient accumulation to get a larger effective batch size for better training results. I've been using 32 gradient accumulation steps. As usual we're using the Adafactor optimiser. More unusually, we're using the experimental full BF16 training option which may give worse results but saves a lot of memory. Both the latents and text encoder outputs are cached to disk before starting the main training run, saving even more precious VRAM - the SDXL text encoder is quite big compared to 1.5 and encoding latents at 1024 resolution is memory hungry. This does use quite a bit of disk space and mean it's not possible to use augmentation options like random image cropping and tag shuffling. With more VRAM or a lower resolution some of these compromises can be avoided. A config file with all these options and some other defaults is in the attached ZIP file, though I haven't tested it extensively and can't confirm how effective those specific settings are.

Note that due to rounding limitations with bfloat16 format, you may need to increase the learning rate substantially from what's set here or wait for stochastic rounding to be implemented. This is unfortunately not particularly well documented anywhere.

Run the training script as usual:

python3 ssd_1b_train.py --config_file options_finetune_ssd-1b_in_8gb.toml --train_data_dir your-training-data --in_json your-metadata.json --pretrained_model_name_or_path base-model.safetensors --output_dir ssd-finetune-run --output_name ssd-finetuned

With a bit of luck and some compute time, you should be able to get a suitable fine tune out. Unfortunately, I don't currently have an example checkpoint fine tuned using this configuration or sample images - my fine tuning attempts are incomplete and were done with higher VRAM systems before discovering it was possible to reduce RAM usage this much, and I'm out of time budget and normal budget for now. Have fun experimenting and best of luck!