【Intel Arc A770】【Diffusers】「enable_model_cpu_offload」と「enable_sequential_cpu_offload」を使う

はじめに

「enable_model_cpu_offload」と「enable_sequential_cpu_offload」はDiffusersにおけるVRAM使用量削減のための手段です。

Intel Arcでは使えないと思っていたのですが「(device="xpu")」を付けると使えるようになりました。

今回はstable-diffusion-3.5-mediumで試しています。

PC環境

Ubuntu 24.04
Intel Arc A770 (VRAM 16GB)
Python 3.12

実行

Intel Arcのサポートを開始したnativeなPyTorchとIntel Extension for PyTorchの両方で実行してその速度も比較してみました。

native PyTorch

torch==2.6.0.dev20241222+xpu

enable_model_cpu_offload

torch.xpu.max_memory_allocated: 12.90 GB
time: 259.08 sec

enable_sequential_cpu_offload

torch.xpu.max_memory_allocated: 8.69 GB
time: 309.09 sec

Intel Extension for PyTorch

torch==2.5.1+cxx11.abi

enable_model_cpu_offload

torch.xpu.max_memory_allocated: 10.18 GB
time: 97.65 sec

enable_sequential_cpu_offload

torch.xpu.max_memory_allocated: 4.10 GB
time: 147.62 sec

結果

圧倒的に「Intel Extension for PyTorch」の方が高速です。

以前と同様の結果でした。
touch-sp.hatenablog.com

使用したPythonスクリプト

import torch
from diffusers import StableDiffusion3Pipeline
from decorator import time_monitor

@time_monitor
def main():
    pipe = StableDiffusion3Pipeline.from_pretrained(
        "stable-diffusion-3.5-medium",
        torch_dtype=torch.bfloat16
    )
    pipe.enable_model_cpu_offload(device="xpu")
    #pipe.enable_sequential_cpu_offload(device="xpu")

    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=40,
        guidance_scale=4.5,
        generator=torch.manual_seed(42)
    ).images[0]

    image.save("result.jpg")

    print(f"torch.xpu.max_memory_allocated: {torch.xpu.max_memory_allocated()/ 1024**3:.2f} GB")

if __name__ == "__main__":
    main()

タイム計測のために以下の「decorator.py」を用意しました。

import functools
import time

def time_monitor(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"time: {(end_time - start_time):.2f} sec")
        return result
    return wrapper

関連記事

touch-sp.hatenablog.com