[training] add an offload utility that can be used as a context manager.

Open sayakpaul opened this issue 8 months ago • 3 comments

What does this PR do?

Useful for handling offloading in a simple and easy manner.

@DN6 I am including changes to examples/**/*.py to also trigger the fast PR tests on GPU. Hope, that's okay?

Cc: @a-r-r-o-w for visibility.

Jun 23 '25 08:06 sayakpaul

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Jun 23 '25 08:06 HuggingFaceDocBuilderDev

nice @sayakpaul, thanks! should we also add a test to [test_dreambooth_lora_hidream.py]?(https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/test_dreambooth_lora_hidream.py)

Jun 23 '25 10:06 linoytsaban

should we also add a test to [test_dreambooth_lora_hidream.py]?

I think it should be added as a utility test under tests/others. Currently, it's a bit non-trivial to test in terms of how the testing suite is designed for the scripts.

Jun 23 '25 10:06 sayakpaul

@DN6 okay regarding?

I am including changes to examples/**/*.py to also trigger the fast PR tests on GPU. Hope, that's okay?

Jul 04 '25 04:07 sayakpaul

@linoytsaban two things:

Would you be able to propagate / introduce the offloading related changes to the other important training scripts?
In https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_flux_kontext.py, we introduced the ability to precompute text embeddings even when each image has a custom prompt. We could propagate that to the rest of the Flux training scripts?

Jul 16 '25 08:07 sayakpaul