SD 3.5 Medium na M1: optymalizacja do 11.6 GB

Optymalizacja pamięci dla Stable Diffusion 3.5 Medium na Apple Silicon M1

Model zawiera T5-XXL (5,5 mld parametrów, 10,75 GB w formacie fp16), DiT-transformator (2,1 mld, 4,18 GB) oraz VAE (0,16 GB). Łączna waga wynosi ok. 15 GB, ale uruchomienie pipe.to("mps") po from_pretrained powoduje szczytowe zużycie ponad 21 GB.

Przyczyny szczytu:

• Duplikacja wag: ładowanie do CPU, kopiowanie do MPS.
• enable_model_cpu_offload() jest nieefektywne — CPU i MPS dzielą pamięć unified.
• Hooki z accelerate z device_map utrzymują silne odniesienia, blokując GC.

Standardowy kod:

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

Szczyt: 21,4 GB niezależnie od liczby num_inference_steps.

Rozwiązanie: sekwencyjne ładowanie

Podział na fazy: kodowanie promptu (T5), wyładowanie, generacja (DiT + VAE).

1. Ładowanie z device_map="balanced"

Bezpośrednie załadowanie do MPS bez kopiowania do CPU:

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

Pik T5: ~11 GB.

2. Poprawne wyładowanie hooków

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

Bez tego T5 pozostaje w pamięci (10,75 GB).

3. Pełna sekwencja

Faza 1: Enkoder

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="a warrior with a sword",
    prompt_2="a warrior with a sword",
    prompt_3="a warrior with a sword",
    device="mps", num_images_per_prompt=1)

Wyładowanie według kroku 2.

Faza 2: Generator

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") ponownie wykorzystuje cache T5, zmniejszając zużycie DiT z 11,5 do 6 GB.

Porównanie zużycia pamięci

| Etap | Przed optymalizacją | Po optymalizacji |

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

| Ładowanie T5 | 15,3 GB | 11,5 GB |

| T5 po wyładowaniu| 10,75 GB | 0,01 GB |

| DiT + dyfuzja | 8,9 GB | 8,9 GB |

| Szczyt | 21,4 GB | 11,6 GB |

Wydajność na sprzęcie

• M1 16 GB: 4:54 min (T5: 1:06, DiT: 3:34), 11–13 GB.
• M3 Pro 36 GB: 1:47 min (T5: 0:14, DiT: 1:27), 9–11 GB.

Rozdzielczość 512x512, 40 kroków.

Co ważne:

• T5-XXL to główny zużywacz pamięci (71%).
• device_map="balanced" + usunięcie hooków krytyczne do wyładowania.
• Sekwencyjne ładowanie obniża szczyt o 45%.
• Działa na M1 16 GB bez swapowania.
• .to("mps") dla DiT rozwiązuje ponowne wykorzystanie cache.

— Editorial Team