Repository navigation
Implement attention splicing for textual inversion #757
Description
Is your feature request related to a problem? Please describe.
Textual inversion requires a lot of (V)RAM and could possibly benefit from attention splicing.
Describe the solution you'd like
Implement attention splicing for textual inversion
Describe alternatives you've considered
I tried to use unet.set_attention_slice(unet.config.attention_head_dim // 2), but it didn't seem to do much. I am not sure if I need to enable it at more places or if it just is not efficient for the training pipeline.
Activity
patil-suraj commentedon Oct 7, 2022
Hi @allo- ! the attention slicing is mostly meant to be used for inference, it won't help much for training.
Also the textual inversion training fits on 16GB GPU and we are working on couple of other things to reduce it a bit further.
cf #687
allo- commentedon Oct 7, 2022
With batch size 1 it fits into 12 GB for me, but when the GPU is used for other things as well, it becomes complicated. A workaround is to learn smaller sizes. 480 leaves more room in GPU RAM and still yields acceptable results. But every optimization is welcome!
Next thing I wanted to try is, if fp16 trained embeddings work on fp32 models. I guess it works (better than the other way round) as the fp16 embeddings can be represented in larger floating point types as well and one may only miss a better minimum when training with less precision?
patil-suraj commentedon Oct 7, 2022
fp16embeddings should work fine with f32 model.allo- commentedon Oct 7, 2022
After disabling it again, it seems it made at least some difference to enable attention slicing on unet.
At least when I commented the line I got a OOM crash on the next run. It could of course be coincidence but the last runs before it was fine (while still being near the limit of my VRAM). When it saves some memory, it could be useful to (optionally) enable it.
patil-suraj commentedon Oct 10, 2022
Once #687 is merged it should help save a lot of VRAM, we can revisit this if that saving is not enough and then optionally enable this.