Memory-Efficient Diffusion Transformers: Wie Quanto und Diffusers die Künstliche Intelligenz Revolutionieren

Einführung

In den letzten Monaten haben wir einen Anstieg der Nutzung von Transformer-basierten Diffusions-Backbones für die hochauflösende Text-zu-Bild (T2I) Generierung erlebt. Diese Modelle nutzen die Transformer-Architektur als Baustein für den Diffusionsprozess, anstelle der zuvor vorherrschenden UNet-Architektur. Dank der Natur der Transformer zeigen diese Backbones eine gute Skalierbarkeit, mit Modellen, die von 0,6 Milliarden bis zu 8 Milliarden Parametern reichen.

Herausforderungen und Lösungen

Mit zunehmender Größe der Modelle steigen auch die Speicheranforderungen. Das Problem verschärft sich, da eine Diffusions-Pipeline normalerweise aus mehreren Komponenten besteht: einem Text-Encoder, einem Diffusions-Backbone und einem Bild-Decoder. Moderne Diffusions-Pipelines verwenden oft mehrere Text-Encoder – beispielsweise gibt es bei Stable Diffusion 3 drei davon. Es benötigt 18,765 GB GPU-Speicher, um SD3-Inferenz mit FP16-Präzision auszuführen.

Diese hohen Speicheranforderungen erschweren die Nutzung dieser Modelle mit Verbraucher-GPUs, was die Akzeptanz verlangsamt und Experimente erschwert. In diesem Artikel zeigen wir, wie die Speichereffizienz von Transformer-basierten Diffusions-Pipelines durch den Einsatz der Quantisierungs-Utilities von Quanto aus der Diffusers-Bibliothek verbessert werden kann.

Quantisierung einer Diffusions-Pipeline mit Quanto

Die Quantisierung eines Modells mit Quanto ist unkompliziert. Hier ein Beispielcode:

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)

Wir rufen quantize() auf dem zu quantisierenden Modul auf und spezifizieren, was wir quantisieren möchten. Im obigen Fall quantisieren wir nur die Parameter und belassen die Aktivierungen unverändert. Wir quantisieren auf den FP8-Datentyp. Schließlich rufen wir freeze() auf, um die originalen Parameter durch die quantisierten Parameter zu ersetzen.

Wir können diese Pipeline dann normal aufrufen:

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

Wir stellen fest, dass die Verwendung von FP8 zu Speichereinsparungen führt, jedoch mit leicht höherer Latenz und nahezu keiner Qualitätsverschlechterung:

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

Wir können den Text-Encoder auf die gleiche Weise quantisieren:

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

Das Quantisieren sowohl des Text-Encoders als auch des Diffusions-Backbones führt zu deutlich größeren Speicherverbesserungen:

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

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