Fix Group offloading behaviour when using streams

[a-r-r-o-w]

Fixes #11041, #11086

Requires #11094 to be merged first (also see it for more details).

code

import torch
from diffusers import AutoencoderKLWan, WanPipeline, UniPCMultistepScheduler
from diffusers.hooks import apply_group_offloading, ModelHook, HookRegistry
from diffusers.utils import export_to_video
from diffusers.utils.logging import set_verbosity_debug

set_verbosity_debug()

class NanDetectorHook(ModelHook):
    def __init__(self, name):
        super().__init__()
        self.name = name
    
    def post_forward(self, module, output):
        if torch.is_tensor(output):
            if torch.isnan(output).any():
                print(f"NaN detected in {self.name}")
                # breakpoint()
        elif isinstance(output, (list, tuple)):
            for i, o in enumerate(output):
                if torch.is_tensor(o) and torch.isnan(o).any():
                    print(f"NaN detected in {self.name}[{i}]")
                    # breakpoint()
        return output

# Available models: Wan-AI/Wan2.1-T2V-14B, Wan-AI/Wan2.1-T2V-1.3B
model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
scheduler = UniPCMultistepScheduler(num_train_timesteps=1000, prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0)
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)

onload_device = torch.device("cuda")
offload_device = torch.device("cpu")

pipe.text_encoder.to(onload_device)
pipe.vae.to(onload_device)
apply_group_offloading(
    pipe.transformer,
    onload_device=onload_device,
    offload_device=offload_device,
    offload_type="leaf_level",
    use_stream=True,
)

for name, module in pipe.transformer.named_modules():
    registry = HookRegistry.check_if_exists_or_initialize(module)
    hook = NanDetectorHook(name)
    registry.register_hook(hook, "nan_detector")

prompt = "A cat walks on the grass, realistic"
negative_prompt = negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"

output = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=480,
    width=832,
    num_frames=49,
    guidance_scale=5.0,
    num_inference_steps=30,
    generator=torch.Generator().manual_seed(42),
).frames[0]
export_to_video(output, "output3.mp4", fps=15)

🐛 explanation:

• When using streams, we need to find out the layer execution order for prefetching weights on separate cuda stream
• To find the layer execution order, we use LazyPrefetchGroupOffloadingHook, which is active during the first forward pass only
• When using streams, we always set non_blocking parameter to True
• For the first forward pass, every module group needs to onload itself
• So, if non_blocking is True and every module is onloading itself (for the first forward pass), we need to either cuda sync after weight transfer, or do the weight transfer in blocking manner
• If we don't wait for weight transfer to complete and computation kernel launches, the result computed is garbage
• In this PR, we make sure to do the first forward pass in blocking manner, and only do non-blocking alternate-stream transfers once we know the prefetching order.

As a result, using group offloading alongside streams is slightly slower (depends on model size). For Wan, there seems to be an extra 3-5s overhead due to first forward pass being done in blocking manner.

Revealed with good old absolute rockstar pytorch profiler!

Note: this appears to be random and not reproducible on all GPUs for all height/width/num_frames and for all models 🫠

@Passenger12138 LMK if this fixes the issue you were facing.

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

[a-r-r-o-w]

Failing tests are unrelated