Diffusers documentation

Сокращение использования памяти

Diffusers

Начало работы

Diffusers Установка Быстрый старт Интерфейс командной строки

Вывод

Использование конвейеров диффузии

Адаптеры и персонализация

Оптимизация и масштабирование

Базовая производительность

Скорость и память

Ускорение вывода Кэширование Бэкенды внимания Пакетный вывод Сокращение использования памяти Удалённый вывод Компиляция и выгрузка квантованных моделей

Распределённый вывод

Квантование

Ускорение для конкретного оборудования

Методы сообщества

Модульные Diffusers

Обучение и дообучение

Ресурсы

Внести вклад

API

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Сокращение использования памяти

Современные диффузионные модели, такие как Flux и Wan, содержат миллиарды параметров, что требует значительного объема памяти на вашем оборудовании для инференса. Это создает трудности, так как у распространенных GPU часто недостаточно памяти. Чтобы преодолеть ограничения по памяти, вы можете использовать несколько GPU (если они доступны), выгружать некоторые компоненты пайплайна на CPU и применять другие методы.

В этом руководстве показано, как уменьшить потребление памяти.

Имейте в виду, что эти методы может потребоваться адаптировать в зависимости от модели. Например, диффузионная модель на основе трансформера может извлекать из этих оптимизаций памяти не такую же выгоду, как модель на основе UNet.

Несколько GPU

Если у вас есть доступ более чем к одному GPU, существует несколько вариантов для эффективной загрузки и распределения большой модели по вашему оборудованию. Эти функции поддерживаются библиотекой Accelerate, поэтому убедитесь, что она установлена.

pip install -U accelerate

Шардированные контрольные точки

Загрузка больших контрольных точек, разбитых на несколько фрагментов, полезна, поскольку фрагменты загружаются по одному. Это позволяет поддерживать низкое потребление памяти, требуя достаточно памяти только для размера модели и размера самого большого фрагмента. Мы рекомендуем использовать разбиение на фрагменты, если контрольная точка в формате fp32 превышает 5 ГБ. Размер фрагмента по умолчанию составляет 5 ГБ.

Разбейте контрольную точку на части в save_pretrained() с помощью параметра max_shard_size.

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")

Теперь вы можете использовать шунтированную контрольную точку вместо обычной, чтобы сэкономить память.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

unet = AutoModel.from_pretrained(
    "username/sdxl-unet-sharded", dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    unet=unet,
    dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

Размещение на устройствах

Размещение на устройствах — это экспериментальная функция, и API может измениться. На данный момент поддерживается только стратегия balanced. В будущем мы планируем поддержать дополнительные стратегии отображения.

Параметр device_map управляет тем, как компоненты модели в конвейере или слои отдельной модели распределяются по устройствам.

уровень конвейера

уровень модели

При проектировании собственного device_map он должен представлять собой словарь, в котором ключами являются имена конкретных модулей или слоёв модели, а значениями — идентификаторы устройств (целое число для GPU, cpu для CPU и disk для диска).

Вызовите hf_device_map для модели, чтобы увидеть, как распределены её слои, а затем спроектируйте собственную схему.

print(transformer.hf_device_map)
{'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 'cpu', 'single_transformer_blocks.11': 'cpu', 'single_transformer_blocks.12': 'cpu', 'single_transformer_blocks.13': 'cpu', 'single_transformer_blocks.14': 'cpu', 'single_transformer_blocks.15': 'cpu', 'single_transformer_blocks.16': 'cpu', 'single_transformer_blocks.17': 'cpu', 'single_transformer_blocks.18': 'cpu', 'single_transformer_blocks.19': 'cpu', 'single_transformer_blocks.20': 'cpu', 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'}

Например, device_map ниже размещает single_transformer_blocks.10 через single_transformer_blocks.20 на втором GPU (1).

import torch
from diffusers import AutoModel

device_map = {
    'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 1, 'single_transformer_blocks.11': 1, 'single_transformer_blocks.12': 1, 'single_transformer_blocks.13': 1, 'single_transformer_blocks.14': 1, 'single_transformer_blocks.15': 1, 'single_transformer_blocks.16': 1, 'single_transformer_blocks.17': 1, 'single_transformer_blocks.18': 1, 'single_transformer_blocks.19': 1, 'single_transformer_blocks.20': 1, 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'
}

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map=device_map,
    dtype=torch.bfloat16
)

Передайте словарь, сопоставляющий максимальное использование памяти с каждым устройством, чтобы установить ограничение. Если устройство отсутствует в max_memory, оно игнорируется, и компоненты конвейера не распределяются на него.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

max_memory = {0:"1GB", 1:"1GB"}
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="balanced",
    max_memory=max_memory
)

Diffusers по умолчанию использует максимальный объём памяти всех устройств, но если модели не помещаются в память GPU, необходимо использовать один GPU и выгружать компоненты на CPU с помощью описанных ниже методов.

enable_model_cpu_offload() работает только на одном GPU, но очень большая модель может не поместиться на нём
enable_sequential_cpu_offload() может работать, но крайне медленно и поддерживает только один GPU

Используйте метод reset_device_map() для сброса device_map. Это необходимо, если вы хотите использовать методы, такие как .to(), enable_sequential_cpu_offload() и enable_model_cpu_offload(), на пайплайне, для которого была настроена карта устройств.

pipeline.reset_device_map()

Слайсинг VAE

Срезание VAE экономит память, разбивая большие батчи входных данных на одиночные батчи и обрабатывая их по отдельности. Этот метод наиболее эффективен при генерации более одного изображения за раз.

Например, если вы генерируете 4 изображения одновременно, декодирование увеличит пиковую память активаций в 4 раза. Срезание VAE снижает это значение, декодируя по одному изображению за раз, а не все 4 изображения одновременно.

Вызовите enable_slicing(), чтобы включить слайсинг VAE. При декодировании батчей из нескольких изображений можно ожидать небольшое повышение производительности, а для батчей из одного изображения влияние на производительность отсутствует.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_slicing()
pipeline(["An astronaut riding a horse on Mars"]*32).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

Классы AutoencoderKLWan и AsymmetricAutoencoderKL не поддерживают нарезку на слайсы.

Плиточная обработка VAE

Тайлинг VAE экономит память, разбивая изображение на более мелкие перекрывающиеся плитки вместо обработки всего изображения сразу. Это также снижает пиковое потребление памяти, поскольку GPU обрабатывает только одну плитку за раз.

Вызовите enable_tiling(), чтобы включить тайлинг VAE. Сгенерированное изображение может иметь небольшие различия в тоне от тайла к тайлу, так как они декодируются отдельно, но заметных швов между тайлами быть не должно. Тайлинг отключается для разрешений ниже заранее заданного (но настраиваемого) предела. Например, для VAE в StableDiffusionPipeline этот предел составляет 512x512.

import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_tiling()

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, image=init_image, strength=0.5).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan и AsymmetricAutoencoderKL не поддерживают тайлинг.

Чтобы полностью исключить использование памяти VAE на вашем GPU, кодируйте и декодируйте данные на удалённом эндпоинте с помощью Remote inference.

Выгрузка

Стратегии выгрузки перемещают слои или модели, не используемые в данный момент, на CPU, чтобы избежать роста потребления памяти GPU. Эти стратегии можно комбинировать с квантованием и torch.compile, чтобы сбалансировать скорость вывода и использование памяти.

См. руководство Компиляция и выгрузка квантованных моделей для получения подробной информации.

Размещение на CPU

CPU offloading выборочно перемещает веса с GPU на CPU. Когда компонент требуется, он передаётся на GPU, а когда не требуется — перемещается обратно на CPU. Этот метод работает с подмодулями, а не с целыми моделями. Он экономит память, избегая хранения всей модели на GPU.

CPU offloading значительно снижает потребление памяти, но при этом крайне медленный, поскольку подмодули многократно передаются между устройствами. Из-за низкой скорости этот подход часто оказывается непрактичным.

Не переносите конвейер на CUDA до вызова enable_sequential_cpu_offload(), иначе экономия памяти будет минимальной (подробнее см. в этой заявке). Это операция с состоянием, которая устанавливает хуки на модель.

Вызовите enable_sequential_cpu_offload(), чтобы включить его в конвейере.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_sequential_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

Выгрузка модели

Model offloading перемещает целые модели на GPU, а не выборочно переносит отдельные слои или компоненты модели. Одна из основных моделей конвейера, обычно текстовый энкодер, UNet и VAE, размещается на GPU, тогда как остальные компоненты удерживаются на CPU. Компоненты, такие как UNet, которые выполняются несколько раз, остаются на GPU до тех пор, пока их работа полностью не завершится и они больше не понадобятся. Это устраняет накладные расходы на коммуникацию, связанные с CPU offloading, и делает model offloading более быстрым вариантом. Компромисс состоит в том, что экономия памяти будет не такой значительной.

Имейте в виду, что если модели переиспользуются вне конвейера после установки хуков (см. Удаление хуков для получения подробной информации), вам необходимо запускать весь конвейер и модели в ожидаемом порядке, чтобы корректно выгружать их из памяти. Это операция с состоянием, которая устанавливает хуки на модель.

Вызовите enable_model_cpu_offload(), чтобы включить его в конвейере.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

enable_model_cpu_offload() также помогает, когда вы используете метод encode_prompt() самостоятельно для генерации скрытых состояний текстовых энкодеров.

Групповая разгрузка

Group offloading перемещает группы внутренних слоёв (torch.nn.ModuleList или torch.nn.Sequential) на CPU. Он использует меньше памяти, чем model offloading, и работает быстрее, чем CPU offloading, поскольку снижает накладные расходы на коммуникацию.

Групповая разгрузка может не работать со всеми моделями, если реализация прямого прохода содержит приведение типов входных данных к устройству, зависящее от весов, поскольку это может конфликтовать с механизмом приведения типов к устройству, используемым в групповой разгрузке.

Включите групповую разгрузку, установив параметр offload_type в значение block_level или leaf_level.

block_level выгружает группы слоёв на основе параметра num_blocks_per_group. Например, если num_blocks_per_group=2 для модели с 40 слоями, то за один раз загружаются и выгружаются 2 слоя (всего 20 операций загрузки/выгрузки). Это значительно снижает требования к памяти.
leaf_level выгружает отдельные слои на нижнем уровне и эквивалентен CPU offloading. Однако его можно ускорить, используя потоки, без потери скорости инференса.

Групповая разгрузка поддерживается для целых конвейеров или отдельных моделей. Применение групповой разгрузки ко всему конвейеру — самый простой вариант, тогда как выборочное применение к отдельным моделям даёт пользователям больше гибкости для использования различных техник разгрузки для разных моделей.

pipeline

модель

CUDA-поток

Параметр use_stream можно активировать для CUDA-устройств, поддерживающих потоки асинхронной передачи данных, чтобы сократить общее время выполнения по сравнению с выгрузкой на CPU. Он перекрывает передачу данных и вычисления, используя предварительную загрузку слоёв. Следующий для выполнения слой загружается на GPU, пока текущий слой ещё выполняется. Это может значительно увеличить потребление памяти CPU, поэтому убедитесь, что у вас есть вдвое больше памяти, чем размер модели.

Установите record_stream=True, чтобы получить больший прирост скорости ценой незначительного увеличения потребления памяти. Подробнее см. в документации torch.Tensor.record_stream.

При использовании use_stream=True для VAE с включённым тайлингом убедитесь, что перед инференсом выполняется пробный прямой проход (можно и с пробными входными данными), чтобы избежать ошибок несовпадения устройств. Это может не работать во всех реализациях, поэтому при возникновении проблем не стесняйтесь открыть issue.

Если вы используете разгрузку групп на уровне блоков block_level с включённым use_stream, параметр num_blocks_per_group должен быть установлен в значение 1, иначе будет выдано предупреждение.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)

Параметр low_cpu_mem_usage можно установить в True, чтобы снизить потребление оперативной памяти процессора при использовании потоков во время разгрузки групп. Это наиболее эффективно для разгрузки уровня leaf_level и в случаях, когда оперативная память процессора является узким местом. Экономия памяти достигается за счёт создания закреплённых тензоров на лету, а не их предварительного закрепления. Однако это может увеличить общее время выполнения.

Разгрузка на диск

Групповая разгрузка может потреблять значительный объём системной памяти в зависимости от размера модели. На системах с ограниченной памятью попробуйте использовать диск как вторичную память для групповой разгрузки.

Укажите аргумент offload_to_disk_path в enable_group_offload() или apply_group_offloading(), чтобы выгрузить модель на диск.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", offload_to_disk_path="path/to/disk")

apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2, offload_to_disk_path="path/to/disk")

См. эти две таблицы, чтобы сравнить компромиссы между скоростью и потреблением памяти.

Послойное приведение типов

Сочетайте послойное приведение типов с group offloading для ещё большей экономии памяти.

Слойное приведение типов хранит веса в формате с меньшим разрядом (например, torch.float8_e4m3fn и torch.float8_e5m2), чтобы экономить память, и приводит их к более высокой точности, такой как torch.float16 или torch.bfloat16, для вычислений. Некоторые слои (веса, связанные с нормализацией и модуляцией) пропускаются, так как их хранение в fp8 может ухудшить качество генерации.

Послойное приведение типов может не работать со всеми моделями, если реализация forward содержит внутреннее приведение типов весов. Текущая реализация послойного приведения типов предполагает, что прямой проход не зависит от точности весов, а типы входных данных всегда указаны в compute_dtype (см. здесь пример несовместимой реализации).

Послойное приведение типов (casting) также может не работать в пользовательских реализациях моделей с слоями PEFT. Существуют некоторые проверки, но они не прошли обширного тестирования и не гарантируют корректную работу во всех случаях.

Вызовите enable_layerwise_casting(), чтобы задать типы данных для хранения и вычислений.

import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b",
    transformer=transformer,
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
prompt = (
    "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
    "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
    "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
    "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
    "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
    "atmosphere of this unique musical performance."
)
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
export_to_video(video, "output.mp4", fps=8)

Метод apply_layerwise_casting() также можно использовать, если требуется больший контроль и гибкость. Его можно применить частично к слоям модели, вызвав его для конкретных внутренних модулей. Используйте параметры skip_modules_pattern или skip_modules_classes, чтобы указать модули, которые нужно пропустить, например, слои нормализации и модуляции.

import torch
from diffusers import CogVideoXTransformer3DModel
from diffusers.hooks import apply_layerwise_casting

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)

# skip the normalization layer
apply_layerwise_casting(
    transformer,
    storage_dtype=torch.float8_e4m3fn,
    compute_dtype=torch.bfloat16,
    skip_modules_classes=["norm"],
    non_blocking=True,
)

torch.channels_last

torch.channels_last меняет способ хранения тензоров с (batch size, channels, height, width) на (batch size, height, width, channels). Это приводит тензоры в соответствие с тем, как оборудование последовательно обращается к тензорам, хранящимся в памяти, и исключает необходимость перескакивать по памяти для доступа к значениям пикселей.

Не все операторы в настоящее время поддерживают формат channels-last, что может привести к ухудшению производительности, но всё же стоит попробовать.

print(pipeline.unet.conv_out.state_dict()["weight"].stride())  # (2880, 9, 3, 1)
pipeline.unet.to(memory_format=torch.channels_last)  # in-place operation
print(
    pipeline.unet.conv_out.state_dict()["weight"].stride()
)  # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works

Внимание с эффективной работой с памятью

Diffusers поддерживает несколько энергоэффективных бэкендов внимания (FlashAttention, xFormers, SageAttention и другие) с помощью set_attention_backend(). Ознакомьтесь с руководством Attention backends, чтобы узнать, как переключаться между ними.

Update on GitHub

←Пакетный вывод Удалённый вывод→