Speicheroptimierte Diffusionsmodelle: Neue Technologien verbessern KI-Effizienz

from optimum.quanto import freeze, qfloat8, quantize
from diffusers import PixArtSigmaPipeline
import torch

pipeline = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
).to("cuda")

quantize(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)

image = pipeline("ghibli style, a fantasy landscape with castles").images[0]

Batch Size  Quantization  Memory (GB)  Latency (Seconds)
1           None          12.086       1.2
1           FP8           11.547       1.540
4           None          12.087       4.482
4           FP8           11.548       5.109

quantize(pipeline.text_encoder, weights=qfloat8)
freeze(pipeline.text_encoder)

Batch Size  Quantization  Quantize TE  Memory (GB)  Latency (Seconds)
1           FP8           False        11.547       1.540
1           FP8           True         5.363        1.601
4           FP8           False        11.548       5.109
4           FP8           True         5.364        5.141

Allgemeinheit der Beobachtungen

Das Quantisieren des Text-Encoders zusammen mit dem Diffusions-Backbone funktioniert im Allgemeinen für die von uns getesteten Modelle. Stable Diffusion 3 ist ein Sonderfall, da es drei verschiedene Text-Encoder verwendet. Wir haben festgestellt, dass das Quantisieren des zweiten Text-Encoders nicht gut funktioniert. Wir empfehlen daher die folgenden Alternativen:

- Nur den ersten Text-Encoder (CLIPTextModelWithProjection) quantisieren - Nur den dritten Text-Encoder (T5EncoderModel) quantisieren - Den ersten und dritten Text-Encoder quantisieren

Die folgende Tabelle gibt eine Vorstellung von den erwarteten Speichereinsparungen für verschiedene Kombinationen der Text-Encoder-Quantisierung (der Diffusions-Transformer ist in allen Fällen quantisiert):

Batch Size  Quantization  Quantize TE 1  Quantize TE 2  Quantize TE 3  Memory (GB)  Latency (Seconds)
1           FP8           1             1             1             8.200        2.858
1           FP8           0             0             1             8.294        2.781
1           FP8           1             1             0             14.384       2.833
1           FP8           0             1             0             14.475       2.818
1           FP8           1             0             0             14.384       2.730
1           FP8           0             1             1             8.325        2.875
1           FP8           1             0             1             8.204        2.789
1           None          -             -             -             16.403       2.118

Zusätzliche Erkenntnisse

Die Verwendung von bfloat16 kann für unterstützte GPU-Architekturen wie H100 oder 4090 schneller sein. Die folgende Tabelle zeigt einige Zahlen für PixArt, die auf unserer H100-Referenzhardware gemessen wurden:

Batch Size  Precision  Quantization  Memory (GB)  Latency (Seconds)  Quantize TE
1           FP16       INT8          5.363        1.538              True
1           BF16       INT8          5.364        1.454              True
1           FP16       FP8           5.363        1.601              True
1           BF16       FP8           5.363        1.495              True

Das Potenzial von qint8

Wir haben festgestellt, dass die Quantisierung mit qint8 (anstelle von fp8) im Allgemeinen besser in Bezug auf die Inferenzlatenz ist. Dieser Effekt wird deutlicher, wenn wir die QKV-Projektionen horizontal verschmelzen (indem wir fuse_qkv_projections() in Diffusers aufrufen), wodurch die Dimensionen der int8-Kerne verdickt werden, um die Berechnung zu beschleunigen. Hier einige Belege für PixArt:

Batch Size  Quantization  Memory (GB)  Latency (Seconds)  Quantize TE  QKV Projection
1           INT8          5.363        1.538              True         False
1           INT8          5.536        1.504              True         True
4           INT8          5.365        5.129              True         False
4           INT8          5.538        4.989              True         True

Wie sieht es mit INT4 aus?

Wir haben zusätzlich mit qint4 experimentiert, wenn wir bfloat16 verwenden. Dies gilt nur für bfloat16 auf H100, da andere Konfigurationen noch nicht unterstützt werden. Mit qint4 können wir mehr Verbesserungen beim Speicherverbrauch erwarten, jedoch auf Kosten einer erhöhten Inferenzlatenz. Dies ist zu erwarten, da es keine native Hardwareunterstützung für int4-Berechnungen gibt – die Gewichte werden mit 4 Bit übertragen, aber die Berechnung erfolgt weiterhin in bfloat16. Die folgende Tabelle zeigt unsere Ergebnisse für PixArt-Sigma:

Batch Size  Quantize TE  Memory (GB)  Latency (Seconds)
1           No           9.380        7.431
1           Yes          3.058        7.604

Beachten Sie jedoch, dass aufgrund der aggressiven Diskretisierung von INT4 die Endergebnisse einen Schlag erleiden können. Aus diesem Grund lassen wir bei Transformer-basierten Modellen im Allgemeinen die letzte Projektionsebene aus der Quantisierung heraus. In Quanto tun wir dies, indem wir:

quantize(pipeline.transformer, weights=qint4, exclude="proj_out")
freeze(pipeline.transformer)

„proj_out“ entspricht der letzten Schicht in pipeline.transformer. Die folgende Tabelle zeigt Ergebnisse für verschiedene Einstellungen:

Quantize TE: No, Layer exclusion: None
Quantize TE: No, Layer exclusion: „proj_out“
Quantize TE: Yes, Layer exclusion: None
Quantize TE: Yes, Layer exclusion: „proj_out“

Bonus - Speichern und Laden von Diffusers-Modellen in Quanto

Quantisierte Diffusers-Modelle können gespeichert und geladen werden:

from diffusers import PixArtTransformer2DModel
from optimum.quanto import QuantizedPixArtTransformer2DModel, qfloat8

model = PixArtTransformer2DModel.from_pretrained("PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", subfolder="transformer")
qmodel = QuantizedPixArtTransformer2DModel.quantize(model, weights=qfloat8)
qmodel.save_pretrained("pixart-sigma-fp8")

Der resultierende Checkpoint ist 587MB groß, anstelle der ursprünglichen 2.44GB. Wir können ihn dann laden:

from optimum.quanto import QuantizedPixArtTransformer2DModel
import torch

transformer = QuantizedPixArtTransformer2DModel.from_pretrained("pixart-sigma-fp8") 
transformer.to(device="cuda", dtype=torch.float16)

Und in einer Diffusions-Pipeline verwenden:

from diffusers import DiffusionPipeline
import torch

pipe = DiffusionPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", 
    transformer=None,
    torch_dtype=torch.float16,
).to("cuda")
pipe.transformer = transformer

prompt = "A small cactus with a happy face in the Sahara desert."
image = pipe(prompt).images[0]

In Zukunft können wir erwarten, den Transformer direkt bei der Initialisierung der Pipeline zu übergeben, sodass dies funktioniert:

pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", 
    transformer=transformer,
    torch_dtype=torch.float16,
).to("cuda")

Schlussfolgerung

In diesem Artikel haben wir gezeigt, wie man Transformer-Modelle aus Diffusers quantisieren und deren Speicherverbrauch optimieren kann. Die Effekte der Quantisierung werden deutlicher, wenn wir zusätzlich die beteiligten Text-Encoder quantisieren. Wir hoffen, dass Sie einige der Workflows auf Ihre Projekte anwenden und davon profitieren können.

Quellen

1. https://arxiv.org/abs/2401.04339
2. https://huggingface.co/blog/lora
3. https://huggingface.co/docs/diffusers/main/en/optimization/memory
4. https://www.photoroom.com/inside-photoroom/stable-diffusion-100-percent-faster-with-memory-efficient-attention
5. https://arxiv.org/abs/2210.11794
6. https://github.com/huggingface/diffusers/issues/1892
7. https://pytorch.org/blog/accelerated-diffusers-pt-20/
8. https://github.com/THUDM/Inf-DiT
9. https://x.com/risingsayak?lang=kn
10. https://www.researchgate.net/publication/371923418_Diffuser_Efficient_Transformers_with_Multi-Hop_Attention_Diffusion_for_Long_Sequences