[core] feat: support group offloading at the pipeline level

[sayakpaul]

What does this PR do?

onload_device=torch.device("cuda")
offload_device=torch.device("cpu")
apply_group_offloading(
text_encoder,
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)
apply_group_offloading(
vae,
onload_device=onload_device,
offload_device=offload_device,
offload_type="block_level",
offload_type="leaf_level",
use_stream=True
)

to

pipe.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True
)

Of course, if users still want to apply different offloading techniques to different model-level components, they can easily choose to do so. But IMO, enable_group_offload() is an easier entrypoint.

We can allow users to pass mappings like we do for quant_mapping in PipelineQuantizationConfig in the future.

Will request for reviews after CI.

TODOs

• Docs (cc: @stevhliu)

[sayakpaul]

I think it's okay to expose this as an argument as opposed to how we do model_cpu_offload_seq, for example:

diffusers/src/diffusers/pipelines/flux/pipeline_flux.py

Line 179 in 764b624

model_cpu_offload_seq="text_encoder->text_encoder_2->image_encoder->transformer->vae"

This is because model CPU offloading relies on a sequence for device management. I don't think we have that constraint in the case of group offloading.

[HuggingFaceDocBuilderDev]

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.

[stevhliu]

Really neat how simplified this has become! 🎉

[a-r-r-o-w]

Nice! We could reduce LoC by creating a kwargs dict and passing into both branches, but definitely not a blocker

Tests look good but maybe the two could be combined into one to reduce total run count