support group offloading under auto offloading

[yiyixuxu]

related to #14346

with this PR, we now support the use case where a group offloaded model now takes part in auto offloading but places itself:

• it can still make room by moving other models aside (its pre_forward consults the strategy as before)
• it is never chosen as the thing to move, since moving it does nothing
• the manager no longer pretends to offload it

Either order works — group offload before or after enable_auto_cpu_offload.

Also here, because it is what makes the combination useful:

• added support for custon offload strategy, i.e. enable_auto_cpu_offload(..., offload_strategy=...) and set_offload_strategy(...),
• A warning when group offloading meets the default AutoOffloadStrategy, it is not meaningful to use AutoOffloadStrategy with group offloading because the default strategy make its decisions from model memory footprints — a number that does not describe a model holding one group at a time.

Demo

"""Peak memory of Z-Image-Turbo under four placement strategies.

  1. on gpu           everything resident, no offloading
  2. auto             ComponentsManager auto offloading, default strategy
  3. custom           auto offloading, but a strategy that only ever moves the text encoder
  4. custom + group   the same, with the transformer group offloaded

`AutoOffloadStrategy` decides from model memory footprints, which say nothing useful about a group
offloaded model — it holds one group at a time, not its whole weight. A strategy that decides from the
workflow instead ("the text encoder is done, move it") does not have that problem, which is what 3 and
4 show.

  python groupoffload_demo.py
  python groupoffload_demo.py --modes custom+group --steps 8
"""

import argparse
import gc

import torch

from diffusers import ModularPipeline
from diffusers.modular_pipelines import ComponentsManager


REPO = "Tongyi-MAI/Z-Image-Turbo"
GB = 1024**3


class OffloadTextEncoderOnly:
    """Move the text encoder off the device whenever something else needs to load, and nothing else.

    Same signature as `AutoOffloadStrategy`: `(hooks, model_id, model, execution_device) -> hooks`.
    """

    def __call__(self, hooks, model_id, model, execution_device):
        return [hook for hook in hooks if "text_encoder" in hook.model_id]


def run(mode: str, steps: int):
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    manager = None if mode == "on gpu" else ComponentsManager()
    pipe = ModularPipeline.from_pretrained(REPO, components_manager=manager)
    pipe.load_components(dtype=torch.bfloat16)

    if mode == "on gpu":
        pipe.to("cuda")
    else:
        if mode == "custom+group":
            pipe.transformer.enable_group_offload(
                onload_device=torch.device("cuda"),
                offload_device=torch.device("cpu"),
                offload_type="block_level",
                num_blocks_per_group=1,
                use_stream=True,
            )
        manager.enable_auto_cpu_offload(device="cuda")
        if mode.startswith("custom"):
            manager.set_offload_strategy(OffloadTextEncoderOnly())

    image = pipe(
        prompt="a red fox in a snowy forest",
        num_inference_steps=steps,
        generator=torch.Generator("cpu").manual_seed(0),
        output_type="pt",
        output="images",
    )
    peak = torch.cuda.max_memory_allocated() / GB

    del pipe, manager
    gc.collect()
    torch.cuda.empty_cache()
    return peak, tuple(image.shape)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--steps", type=int, default=4)
    p.add_argument("--modes", default="on gpu,auto,custom,custom+group")
    a = p.parse_args()

    print(f"{REPO}  {torch.cuda.get_device_name(0)}  steps={a.steps}\n")
    for mode in a.modes.split(","):
        peak, shape = run(mode, a.steps)
        print(f"{mode:14s} peak={peak:6.2f}GB  out={shape}", flush=True)

Tongyi-MAI/Z-Image-Turbo (11.5GB transformer, 3.7GB text encoder in bf16) on one H100, 4 steps:

on gpu         peak= 21.68GB
auto           peak= 21.69GB
custom         peak= 14.15GB
custom+group   peak=  7.77GB

The whole custom strategy is:

class OffloadTextEncoderOnly:
    def __call__(self, hooks, model_id, model, execution_device):
        return [hook for hook in hooks if "text_encoder" in hook.model_id]

auto matching on gpu is the point: on an 80GB card the default strategy never fires, because
nothing is ever short of memory. Deciding from the workflow ("the text encoder is done") buys 7.5GB,
and group offloading the transformer buys another 6.4GB — 64% off the resident case.

Not in this PR

• No tests. will add after [modular] improve auto offload #14304
• Auto group offloading on OOM. After [modular] improve auto offload #14304, the plan is for OOM recovery to apply group
  offloading itself when everything else is already offloaded. The guards here hold for that: when the
  manager group offloads a model it has already hooked, the hook stops moving it from that point on, so no extra bookkeeping is needed.

[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.