Réduire l'utilisation de la mémoire

Les modèles de diffusion modernes tels que Flux et Wan comptent des milliards de paramètres, ce qui requiert une grande quantité de mémoire sur votre matériel pour l’inférence. C’est un défi, car les GPU courants ne disposent souvent pas de mémoire suffisante. Pour contourner ces limitations de mémoire, vous pouvez utiliser plusieurs GPU (si disponibles), décharger certains composants du pipeline sur le CPU, et plus encore.

Ce guide vous montre comment réduire votre consommation de mémoire.

Gardez à l’esprit que ces techniques peuvent nécessiter des ajustements en fonction du modèle. Par exemple, un modèle de diffusion basé sur un transformer ne tirera pas forcément autant parti de ces optimisations de mémoire qu’un modèle basé sur un UNet.

Plusieurs GPU

Si vous disposez de plusieurs GPU, il existe plusieurs options pour charger et distribuer efficacement un grand modèle sur votre matériel. Ces fonctionnalités sont prises en charge par la bibliothèque Accelerate, assurez-vous donc qu'elle est installée au préalable.

pip install -U accelerate

Points de contrôle fractionnés

Charger de grands points de contrôle en plusieurs fragments est utile car les fragments sont chargés un par un. Cela maintient l'utilisation de la mémoire au minimum, ne nécessitant que suffisamment de mémoire pour la taille du modèle et la taille du plus grand fragment. Nous recommandons de fragmenter lorsque le point de contrôle fp32 dépasse 5 Go. La taille de fragment par défaut est de 5 Go.

Fragmenter un point de contrôle dans save_pretrained() avec le paramètre 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")

Vous pouvez désormais utiliser le point de contrôle fragmenté, au lieu du point de contrôle standard, pour économiser de la mémoire.

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"

Placement sur l'appareil

Le placement sur les appareils est une fonctionnalité expérimentale et l'API est susceptible d'évoluer. Seule la stratégie balanced est prise en charge pour le moment. Nous prévoyons de prendre en charge d'autres stratégies de correspondance à l'avenir.

Le paramètre device_map contrôle la manière dont les composants du modèle dans un pipeline ou les couches d'un modèle individuel sont répartis sur les différents appareils.

niveau pipeline

niveau du modèle

Lors de la conception de votre propre device_map, il doit s’agir d’un dictionnaire associant le nom spécifique d’un module ou d’une couche du modèle à un identifiant de périphérique (un entier pour les GPU, cpu pour les CPU et disk pour le disque).

Appel de hf_device_map sur un modèle pour voir comment les couches du modèle sont distribuées, puis concevez la vôtre.

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'}

Par exemple, le device_map ci-dessous place single_transformer_blocks.10 à single_transformer_blocks.20 sur une seconde 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
)

Passez un dictionnaire associant l’utilisation mémoire maximale à chaque appareil pour imposer une limite. Si un appareil n’est pas présent dans max_memory, il est ignoré et les composants du pipeline ne seront pas répartis dessus.

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 utilise par défaut la mémoire maximale de tous les appareils, mais si les modèles ne tiennent pas sur les GPU, vous devrez utiliser un seul GPU et décharger sur le CPU à l’aide des méthodes ci-dessous.

enable_model_cpu_offload() ne fonctionne que sur un seul GPU, mais un modèle très volumineux peut ne pas y tenir
enable_sequential_cpu_offload() peut fonctionner, mais il est extrêmement lent et limité à une seule GPU

Utilisez la méthode reset_device_map() pour réinitialiser le device_map. C'est nécessaire si vous souhaitez utiliser des méthodes comme .to(), enable_sequential_cpu_offload() et enable_model_cpu_offload() sur un pipeline qui a été mappé sur un appareil.

pipeline.reset_device_map()

Découpage du VAE

Le découpage du VAE économise de la mémoire en divisant de grands lots d'entrées en un seul lot de données et en les traitant séparément. Cette méthode fonctionne le mieux lors de la génération de plusieurs images à la fois.

Par exemple, si vous générez 4 images en même temps, le décodage augmenterait la mémoire d’activation de pointe par 4. Le tranchage du VAE réduit cette augmentation en ne décodant qu’une image à la fois au lieu des 4 images simultanément.

Appel de enable_slicing() pour activer le VAE par tranches. Vous pouvez vous attendre à une légère amélioration des performances lors du décodage de lots d'images multiples et à aucun impact sur les performances pour les lots d'images uniques.

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")

Les classes AutoencoderKLWan et AsymmetricAutoencoderKL ne prennent pas en charge le tranchage.

Tiling du VAE

Le découpage en tuiles du VAE économise de la mémoire en divisant une image en tuiles plus petites et chevauchantes, plutôt qu’en traitant l’image entière d’un seul coup. Cela réduit également l’utilisation mémoire de pointe, car le GPU ne traite qu’une tuile à la fois.

Appel de enable_tiling() pour activer le découpage en tuiles du VAE. L’image générée peut présenter de légères variations de tonalité d’une tuile à l’autre, car elles sont décodées séparément, mais il ne devrait y avoir aucune couture visible entre les tuiles. Le découpage est désactivé pour les résolutions inférieures à une limite prédéfinie (mais configurable). Par exemple, cette limite est de 512x512 pour le VAE de StableDiffusionPipeline.

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 et AsymmetricAutoencoderKL ne prennent pas en charge le découpage en tuiles.

Pour éviter d'utiliser la mémoire du VAE sur votre GPU, effectuez l'encodage et le décodage sur un point de distanciation distant avec Remote inference.

Offloading

Les stratégies de déchargement déplacent les couches ou modèles non actifs vers le CPU afin d'éviter d'augmenter la mémoire GPU. Ces stratégies peuvent être combinées avec la quantification et torch.compile pour équilibrer la vitesse d'inférence et l'utilisation de la mémoire.

Consultez le guide Compilation et offloading des modèles quantifiés pour plus de détails.

Déchargement sur CPU

Le déchargement sur CPU (CPU offloading) déplace sélectivement les poids de la GPU vers la CPU. Lorsqu'un composant est requis, il est transféré vers la GPU ; lorsqu'il ne l'est pas, il est déplacé vers la CPU. Cette méthode fonctionne sur des sous-modules plutôt que sur des modèles entiers. Elle économise de la mémoire en évitant de stocker le modèle entier sur la GPU.

Le déchargement sur CPU réduit considérablement l'utilisation de la mémoire, mais il est également extrêmement lent car les sous-modules sont transférés de nombreuses fois entre les appareils. Il peut souvent s'avérer impraticable en raison de sa lenteur.

Ne déplacez pas le pipeline sur CUDA avant d’appeler enable_sequential_cpu_offload(), sinon la quantité de mémoire économisée sera minime (reportez-vous à cette issue pour plus de détails). Il s’agit d’une opération avec état qui installe des hooks sur le modèle.

Appel de enable_sequential_cpu_offload() pour l’activer sur un pipeline.

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")

Déchargement du modèle

Le déchargement de modèle (model offloading) déplace des modèles entiers vers le GPU, au lieu de déplacer sélectivement certaines couches ou composants du modèle. L’un des principaux modèles du pipeline, généralement l’encodeur de texte, le UNet et le VAE, est placé sur le GPU, tandis que les autres composants sont conservés sur le CPU. Les composants comme le UNet, qui s’exécutent plusieurs fois, restent sur le GPU jusqu’à ce qu’ils soient totalement terminés et non plus nécessaires. Cela élimine la surcharge de communication liée au déchargement sur CPU (CPU offloading) et rend le déchargement de modèle une alternative plus rapide. Le compromis est que les économies de mémoire ne seront pas aussi importantes.

Gardez à l’esprit que si les modèles sont réutilisés en dehors du pipeline après l’installation des hooks (voir Removing Hooks pour plus de détails), vous devez exécuter l’ensemble du pipeline et les modèles dans l’ordre attendu afin de les décharger correctement. Il s’agit d’une opération avec état qui installe des hooks sur le modèle.

Appel de enable_model_cpu_offload() pour l’activer sur un pipeline.

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() est également utile lorsque vous utilisez la méthode encode_prompt() seule pour générer l’état caché des encodeurs de texte.

Group offloading

Le group offloading déplace des groupes de couches internes (torch.nn.ModuleList ou torch.nn.Sequential) vers le CPU. Il utilise moins de mémoire que le model offloading et est plus rapide que le CPU offloading car il réduit la surcharge de communication.

Le déchargement par groupe peut ne pas fonctionner avec tous les modèles si l’implémentation du calcul avant (forward) contient un cast de dispositif des entrées dépendant des poids, car cela peut entrer en conflit avec le mécanisme de cast de dispositif du déchargement par groupe.

Activez le déchargement par groupe en configurant le paramètre offload_type sur block_level ou leaf_level.

block_level décharge des groupes de couches en fonction du paramètre num_blocks_per_group. Par exemple, si num_blocks_per_group=2 sur un modèle de 40 couches, 2 couches sont chargées et déchargées à la fois (20 chargements/déchargements au total). Cela réduit considérablement les besoins en mémoire.
leaf_level décharge les couches individuelles au niveau le plus bas et est équivalent à la décharge CPU. Mais il peut être accéléré si vous utilisez des flux sans sacrifier la vitesse d’inférence.

Le group offloading est pris en charge pour des pipelines entiers ou pour des modèles individuels. Appliquer le group offloading à l’ensemble du pipeline est l’option la plus simple, tandis que son application sélective à des modèles individuels offre aux utilisateurs une plus grande flexibilité pour utiliser différentes techniques d’offloading pour différents modèles.

pipeline

model

Flux CUDA

Le paramètre use_stream peut être activé sur les appareils CUDA qui prennent en charge les flux de transfert de données asynchrones afin de réduire le temps d’exécution global par rapport au déchargement sur CPU. Il superpose le transfert de données et le calcul en utilisant la prélecture des couches. La couche suivante à exécuter est chargée sur le GPU pendant que la couche actuelle est encore en cours d’exécution. Cela peut augmenter considérablement la mémoire CPU, assurez-vous donc de disposer de deux fois la taille du modèle en mémoire.

Définissez record_stream=True pour un gain de vitesse plus important, au prix d’une légère augmentation de l’utilisation de la mémoire. Consultez la documentation de torch.Tensor.record_stream pour en savoir plus.

Lorsque use_stream=True est utilisé sur des VAE avec le tiling activé, assurez-vous d'effectuer un passage avant factice (possible également avec des entrées factices) avant l'inférence afin d'éviter les erreurs de correspondance de périphérique. Cela peut ne pas fonctionner sur toutes les implémentations, n'hésitez donc pas à ouvrir un ticket si vous rencontrez des problèmes.

Si vous utilisez le déchargement de groupes au niveau des blocs (block_level) avec use_stream activé, le paramètre num_blocks_per_group doit être défini sur 1, sinon un avertissement sera émis.

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

Le paramètre low_cpu_mem_usage peut être défini sur True pour réduire l’utilisation de la mémoire CPU lors de l’utilisation de flux pendant le déchargement par groupes. C’est particulièrement adapté au déchargement leaf_level et lorsque la mémoire CPU est un goulot d’étranglement. La mémoire est économisée en créant des tenseurs épinglés à la volée plutôt qu’en les épinglant au préalable. Toutefois, cela peut augmenter le temps d’exécution global.

Déchargement sur disque

Le group offloading peut consommer une quantité importante de mémoire système en fonction de la taille du modèle. Sur les systèmes à mémoire limitée, essayez le group offloading sur le disque en tant que mémoire secondaire.

Définissez l'argument offload_to_disk_path dans enable_group_offload() ou apply_group_offloading() pour décharger le modèle sur le disque.

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")

Consultez ces deux tableaux pour comparer les compromis entre vitesse et mémoire.

Casting par couche

Combinez le group offloading avec le casting par couche pour des économies de mémoire encore plus importantes.

Le casting par couche stocke les poids dans un format de données plus compact (par exemple, torch.float8_e4m3fn et torch.float8_e5m2) afin de réduire la mémoire utilisée, puis les convertit en une précision supérieure comme torch.float16 ou torch.bfloat16 pour le calcul. Certaines couches (les poids liés à la normalisation et à la modulation) sont ignorées, car les stocker en fp8 peut dégrader la qualité de la génération.

Le castage par couche peut ne pas fonctionner avec tous les modèles si l’implémentation du passage avant (forward) contient un castage interne des poids. L’implémentation actuelle du castage par couche suppose que le passage avant est indépendant de la précision des poids et que les types de données d’entrée sont toujours spécifiés dans compute_dtype (voir ici pour une implémentation incompatible).

Le castage par couche peut également échouer sur des implémentations de modélisation personnalisées avec des couches PEFT. Des vérifications sont disponibles, mais elles ne sont pas testées de manière exhaustive ni garanties de fonctionner dans tous les cas.

Appel enable_layerwise_casting() pour définir les types de données de stockage et de calcul.

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)

La méthode apply_layerwise_casting() peut également être utilisée si vous avez besoin de plus de contrôle et de flexibilité. Elle peut être appliquée partiellement aux couches du modèle en l'appelant sur des modules internes spécifiques. Utilisez les paramètres skip_modules_pattern ou skip_modules_classes pour spécifier les modules à éviter, tels que les couches de normalisation et de modulation.

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 inverse la façon dont les tenseurs sont stockés, passant de (batch size, channels, height, width) à (batch size, height, width, channels). Cela aligne les tenseurs sur la manière dont le matériel accède séquentiellement aux tenseurs stockés en mémoire et évite de sauter d’un endroit à l’autre en mémoire pour accéder aux valeurs des pixels.

Tous les opérateurs ne prennent pas encore en charge le format channels-last et peuvent entraîner des performances dégradées, mais il vaut toujours la peine d’essayer.

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

Attention efficace en mémoire

Diffusers prend en charge plusieurs backends d'attention efficaces en mémoire (FlashAttention, xFormers, SageAttention, et plus encore) via set_attention_backend(). Consultez le guide Backends d'attention pour apprendre à basculer entre eux.

Update on GitHub