Group offloading with cuda stream prefetching

[a-r-r-o-w]

See #10503 for more details.

Code

import argparse
import gc
import os
import time

# os.environ["TORCH_LOGS"] = "+recompiles_verbose,guards"
os.environ["TORCH_LOGS"] = "+recompiles_verbose"

import types
types.MethodType

import torch
import torch._dynamo
import torch._dynamo.utils
from diffusers import LTXPipeline
from diffusers.utils import export_to_video
from diffusers.hooks import apply_group_offloading

torch._dynamo.config.cache_size_limit = 4

def cleanup():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()


def benchmark_fn(f, *args, **kwargs):
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    start.record()
    output = f(*args, **kwargs)
    end.record()
    torch.cuda.synchronize()
    elapsed_time = round(start.elapsed_time(end) / 1000, 3)

    return elapsed_time, output


def normal():
    pipe = LTXPipeline.from_pretrained("a-r-r-o-w/LTX-Video-diffusers", torch_dtype=torch.bfloat16)
    pipe.to("cuda")
    
    cleanup()
    print(f"Model memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")

    prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
    negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"

    t1 = time.time()
    video = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=768,
        height=512,
        num_frames=161,
        num_inference_steps=50,
    ).frames[0]
    torch.cuda.synchronize()
    t2 = time.time()
    export_to_video(video, "output.mp4", fps=24)
    
    print(f"Inference time: {t2 - t1:.2f} s")
    print(f"Inference memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")


def model_cpu_offload(compile: bool = False):
    pipe = LTXPipeline.from_pretrained("a-r-r-o-w/LTX-Video-diffusers", torch_dtype=torch.bfloat16)
    pipe.enable_model_cpu_offload()

    if compile:
        pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs")
    
    cleanup()
    print(f"Model memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")

    prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
    negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"

    t1 = time.time()
    video = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=768,
        height=512,
        num_frames=161,
        num_inference_steps=50,
    ).frames[0]
    torch.cuda.synchronize()
    t2 = time.time()
    export_to_video(video, "output.mp4", fps=24)
    
    print(f"Inference time: {t2 - t1:.2f} s")
    print(f"Inference memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")


def sequential_cpu_offload(compile: bool = False):
    pipe = LTXPipeline.from_pretrained("a-r-r-o-w/LTX-Video-diffusers", torch_dtype=torch.bfloat16)
    pipe.enable_sequential_cpu_offload()

    if compile:
        pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs")
    
    cleanup()
    print(f"Model memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")

    prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
    negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"

    t1 = time.time()
    video = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=768,
        height=512,
        num_frames=161,
        num_inference_steps=50,
    ).frames[0]
    torch.cuda.synchronize()
    t2 = time.time()
    export_to_video(video, "output.mp4", fps=24)
    
    print(f"Inference time: {t2 - t1:.2f} s")
    print(f"Inference memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")


def group_offloading(offload_group_patterns: str = "diffusers_block", num_blocks_per_group: int = 4, compile: bool = False):
    pipe = LTXPipeline.from_pretrained("a-r-r-o-w/LTX-Video-diffusers", torch_dtype=torch.bfloat16)
    pipe.text_encoder.to("cuda")
    pipe.vae.to("cuda")
    apply_group_offloading(
        module=pipe.transformer,
        offload_group_patterns=offload_group_patterns,
        num_blocks_per_group=num_blocks_per_group,
        offload_device=torch.device("cpu"),
        onload_device=torch.device("cuda"),
        force_offload=True,
        non_blocking=True,
        # cuda_stream=True,
    )
    if compile:
        pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs")
        # pipe.transformer = torch.compile(pipe.transformer, backend="eager")
    
    cleanup()
    print(f"Model memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")

    prompt = "A woman with long brown hair and light skin smiles at another woman with long blonde hair. The woman with brown hair wears a black jacket and has a small, barely noticeable mole on her right cheek. The camera angle is a close-up, focused on the woman with brown hair's face. The lighting is warm and natural, likely from the setting sun, casting a soft glow on the scene. The scene appears to be real-life footage"
    negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted"

    t1 = time.time()
    video = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=768,
        height=512,
        num_frames=161,
        num_inference_steps=50,
    ).frames[0]
    torch.cuda.synchronize()
    t2 = time.time()
    export_to_video(video, "output.mp4", fps=24)

    compile_times = torch._dynamo.utils.compile_times()
    print(f"Compile times: {compile_times}")
    
    print(f"Inference time: {t2 - t1:.2f} s")
    print(f"Inference memory: {torch.cuda.max_memory_reserved() / 1024 ** 3:.2f} GB")


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--offload_type", type=str, default="normal", choices=["normal", "model", "sequential", "group"])
    parser.add_argument("--offload_group_patterns", type=str, default="diffusers_block")
    parser.add_argument("--num_blocks_per_group", type=int, default=4)
    parser.add_argument("--compile", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()

    if args.offload_type == "normal":
        normal()
    elif args.offload_type == "model":
        model_cpu_offload(args.compile)
    elif args.offload_type == "sequential":
        sequential_cpu_offload(args.compile)
    else:
        group_offloading(args.offload_group_patterns, args.num_blocks_per_group, args.compile)

Results for:

python3 test.py --offload_type group --offload_group_patterns diffusers_block --num_blocks_per_group 4

Without cuda stream prefetching:

Model memory: 9.70 GB
Inference time: 85.36 s
Inference memory: 21.25 GB

With cuda streams:

# 4 block groups
Model memory: 10.26 GB
Inference time: 47.71 s
Inference memory: 21.89 GB

# 1 block groups
Model memory: 9.84 GB
Inference time: 46.15 s
Inference memory: 21.04 

Currently not compatible with torch.compile and triggers several recompilations.

Reading:

• https://pytorch.org/tutorials/intermediate/pinmem_nonblock.html
• https://gist.github.com/gau-nernst/9408e13c32d3c6e7025d92cce6cba140
• https://pytorch.org/docs/main/torch.compiler_troubleshooting.html

cc @gau-nernst because the original ideas for layer prefetching come from his implementation

cc @DN6 @yiyixuxu

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

This is a bit hacky for the moment just so that I could get it running quickly without putting in much thought. Will try to improve soon

[a-r-r-o-w]

To avoid some confusion, the numbers in the PR are not a 1:1 comparison with the group offloading PR.

In that PR, the benchmark only times and notes the memory requirement for transformer forward pass (prompt embeddings are precomputed and latents are decoded separately), but here they are including both prompt embedding and decoding steps. Will do a fair comparison numbers report shortly after completing benchmark script