Confusion whether autocast(...) should be used or not?

[patrickvonplaten]

Copy-paste the very nice comment/question from PrinceRay7 from here: #993 (comment)

I'm quite confused. In Diffusers docs, it is strongly disencouraged to use torch.autocast, however, in accelerate scripts, when training with the flag fp16, it is still implemented by using autocast. So i really wonder what should we do when training diffusion models. It seems like accelerate is very necessary and now it conflicts so I'm confused now. Looking forward to your reply, thanks!

[patrickvonplaten]

That's a very good question and @PrinceRay7 is referring to the Diffusers maintainers actively discouraging to use torch.autocast (see warning here) while using it our accelerate example scripts:

diffusers/examples/dreambooth/train_dreambooth.py

Line 88 in c7da8fd

with torch.autocast("cuda"):

So let's try to be more precise here!

1.) torch.autocast(...) should never be used for a trained model during inference. To be more precise, whenever a checkpoint is trained/finetuned and you wish to run it in inference-only, torch.autocast should never be used.

2.) torch.autocast(...) should never be used for the training loop itself. If you look closely into https://github.com/huggingface/diffusers/examples/dreambooth/train_dreambooth.py you see that torch.autocast is only used during validation & evaluation but never for the actually training.

3.) Finally, the only time we have to use torch.autocast is when evaluating models where parts are down-casted to fp16 and parts or kept in fp32 which applies to dreambooth. The rationale here is the following:

• The unet is trained and therefore weights are kept in fp32
• The vae decoder is not trained and therefore can be casted to fp16
• => now this creates a problem when using this combination of vae and unet for validation inference. To solve this we could upcast the vae to fp32 and everything would work fine, however then we increase the memory unnecessarily. We cannot downcast the unet to fp16 as this way we would loose information during training and would make training with and without validation different. So the a good solution (while never advertised for inference) is to use autocast here.

Does this make sense?

[PrinceRay7]

That does make sense! Thank you for your reply, It is elaborate and clear.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.