SD 3.5 Medium sur M1 : optimisation à 11,6 Go

Optimisation de la mémoire pour Stable Diffusion 3.5 Medium sur Apple Silicon M1

Le modèle Stable Diffusion 3.5 Medium, avec ses 2,5 milliards de paramètres, peut consommer jusqu’à 21 Go de mémoire unifiée sur les processeurs Apple Silicon en raison de l’encodeur T5-XXL (10,75 Go en fp16). Une charge optimisée réduit le pic d’utilisation à seulement 11,6 Go, permettant une exécution fluide sur les Mac M1 dotés de 16 Go de RAM. Ce guide explique le problème et présente une stratégie de chargement étape par étape.

Causes principales des pics de mémoire :

• Duplication des poids : le modèle est chargé sur le CPU, puis copié vers MPS.
• enable_model_cpu_offload() est inefficace — le CPU et MPS s’affrontent pour la mémoire unifiée.
• Les hooks Accelerate avec device_map maintiennent des références fortes, empêchant la récupération automatique de mémoire.

Code standard :

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    torch_dtype=torch.float16, variant="fp16")
pipe.to("mps")

Pic de mémoire : 21,4 Go, indépendamment de num_inference_steps.

Solution : Chargement séquentiel

Divisez le processus en phases : encodage du prompt (T5), déchargement, puis génération (DiT + VAE).

1. Charger avec device_map="balanced"

Chargez directement sur MPS sans passer par le CPU :

pipe = StableDiffusion3Pipeline.from_pretrained(
    path, torch_dtype=torch.float16, device_map="balanced")

Pic du T5 : ~11 Go.

2. Suppression correcte des hooks

from accelerate.hooks import remove_hook_from_submodules

for attr in list(vars(pipe)):
    comp = getattr(pipe, attr, None)
    if isinstance(comp, torch.nn.Module):
        remove_hook_from_submodules(comp)
        setattr(pipe, attr, None)

del pipe
gc.collect()
torch.mps.synchronize()
torch.mps.empty_cache()

Sans cette étape, le T5 reste en mémoire (10,75 Go).

3. Workflow complet séquentiel

Phase 1 : Encodeur

enc_pipe = StableDiffusion3Pipeline.from_pretrained(path,
    transformer=None, vae=None,
    text_encoder=None, text_encoder_2=None,
    tokenizer=None, tokenizer_2=None,
    variant="fp16", torch_dtype=torch.float16,
    device_map="balanced")

prompt_embeds, neg_embeds, pooled, neg_pooled = enc_pipe.encode_prompt(
    prompt="un guerrier avec une épée",
    prompt_2="un guerrier avec une épée",
    prompt_3="un guerrier avec une épée",
    device="mps", num_images_per_prompt=1)

Déchargez après cette étape.

Phase 2 : Générateur

gen_pipe = StableDiffusion3Pipeline.from_pretrained(path,
    text_encoder=None, text_encoder_2=None, text_encoder_3=None,
    tokenizer=None, tokenizer_2=None, tokenizer_3=None,
    variant="fp16", torch_dtype=torch.float16)
gen_pipe.to("mps")

image = gen_pipe(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=neg_embeds,
    pooled_prompt_embeds=pooled,
    negative_pooled_prompt_embeds=neg_pooled,
    num_inference_steps=40).images[0]

.to("mps") réutilise le cache du T5, réduisant la charge du DiT de 11,5 Go à 6 Go.

Comparaison de l'utilisation de la mémoire

| Étape | Avant optimisation | Après optimisation |

|-------------------|--------------------|--------------------|

| Chargement T5 | 15,3 Go | 11,5 Go |

| T5 après décharge | 10,75 Go | 0,01 Go |

| DiT + diffusion | 8,9 Go | 8,9 Go |

| Pic maximal | 21,4 Go | 11,6 Go |

Performances sur matériel

• M1 16 Go : 4 min 54 s (T5 : 1 min 6 s, DiT : 3 min 34 s), pic de 11–13 Go.
• M3 Pro 36 Go : 1 min 47 s (T5 : 14 s, DiT : 1 min 27 s), pic de 9–11 Go.

Résolution : 512x512, 40 étapes d’inférence.

Points clés :

• Le T5-XXL est le principal consommateur de mémoire (~71 % du total).
• device_map="balanced" combiné à la suppression des hooks est essentiel pour le déchargement.
• Le chargement séquentiel réduit le pic mémoire de 45 %.
• Fonctionne de manière fiable sur les Mac M1 16 Go sans échange disque.
• .to("mps") sur le DiT permet une réutilisation efficace du cache.

— Editorial Team