Dreambooth: reduce VRAM usage

[gleb-akhmerov]

This allowed me to run Dreambooth on a 3060 with VRAM usage of 11124MiB using the following parameters:

--mixed_precision="fp16" \
--use_8bit_adam \
--enable_xformers_memory_efficient_attention \
--gradient_accumulation_steps=1 \
--gradient_checkpointing \
--resolution=512 \
--train_batch_size=1 \
--train_text_encoder

Also, I was able to use batches of size up to 5 images, with VRAM usage of 11694MiB.

--resolution=768 is possible too, with VRAM usage of 11450MiB.
--resolution=1024 uses 11598MiB, but it's pretty slow (about 4s/it).

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

cc @patil-suraj

[patil-suraj]

Thanks a lot for the PR! It saves quite a bit of memory. But according to torch docs, it changes certain behaviour. So I think instead of setting it to True by default we could make it a argument that the user can control, something like --optimizer_set_to_none, which will be false by default. Also cc @pcuenca @patrickvonplaten

[patrickvonplaten]

Agree with @patil-suraj here, @gleb-akhmerov could we maybe add an arguments --set_grads_to_none ? (maybe a bit better than optimizer_set_to_none

[gleb-akhmerov]

@patil-suraj @patrickvonplaten Good point! I've changed the code to include the argument. Also, I've updated the readme to allow people to find this argument more easily. What do you think?

[patil-suraj]

Looks great, thanks!