Stabile Diffusion ~ Raspberry Pi

Stable Diffusion funktioniert auf Raspberry Pi! 260 MB RAM „halten“ ein großes Modell mit 1 Milliarde Parametern

Stable Diffusion wurde vor 11 Monaten geboren und die Nachricht, dass es auf GPUs der Verbraucherklasse ausgeführt werden kann, hat viele Forscher ermutigt. Darüber hinaus hat Apple offiziell schnell Schluss gemacht und Stable Diffusion in iPhone, iPad und Mac „gestopft“, damit es ausgeführt werden kann. Dies reduziert die Anforderungen von Stable Diffusion an Hardwaregeräte erheblich und macht es nach und nach zu einer „schwarzen Technologie“, die jeder nutzen kann.

Mittlerweile läuft es sogar auf dem Raspberry Pi Zero 2.

Raspberry Pi Zero 2 „Genauso klein. Fünfmal so schnell.」

Was ist das für ein Konzept? Das Ausführen von Stable Diffusion ist nicht trivial, es enthält ein großes Transformer-Modell mit 1 Milliarde Parametern, der empfohlene Mindest-RAM/VRAM beträgt normalerweise 8 GB. Der RPI Zero 2 ist nur ein Mikrocomputer mit 512 MB Speicher.

Das bedeutet, dass die Ausführung von Stable Diffusion auf RPI Zero 2 eine große Herausforderung darstellt. Darüber hinaus haben die Autoren während des Laufs weder den Speicherplatz erhöht noch Zwischenergebnisse auf die Festplatte verlagert.

Im Allgemeinen konzentrieren sich die wichtigsten Frameworks und Bibliotheken für maschinelles Lernen auf die Minimierung der Inferenzlatenz und/oder die Maximierung des Durchsatzes, allerdings auf Kosten der Speichernutzung. Daher beschloss der Autor, eine ultrakleine, hackbare Argumentationsbibliothek zu schreiben, die sich der Minimierung des Speicherverbrauchs widmet.

OnnxStream macht es.

Projektadresse https://github.com/vitoplantamura/OnnxStream

OnnxStream basiert auf der Idee, die Inferenz-Engine von der Komponente zu entkoppeln, die für die Bereitstellung von Modellgewichten verantwortlich ist, einer von WeightsProvider abgeleiteten Klasse. Eine WeightsProvider-Spezialisierung kann jede Art von Laden, Zwischenspeichern und Vorabrufen von Modellparametern implementieren. Beispielsweise könnte ein benutzerdefinierter WeightsProvider entscheiden, Daten direkt von einem HTTP-Server herunterzuladen, ohne etwas auf die Festplatte zu laden oder zu schreiben (aus diesem Grund hat OnnxStream „Stream“ im Namen). Es stehen zwei Standard-WeightsProvider zur Verfügung: DiskNoCache und DiskPrefetch.

Im Vergleich zum Inferenz-Framework OnnxStream von Microsoft benötigt OnnxStream nur 1/55 des Speichers, um den gleichen Effekt zu erzielen, aber die Geschwindigkeit (auf der CPU) ist nur 0,5 bis 2 Mal langsamer als beim ersteren.

Als Nächstes sehen Sie den Effekt der Stable Diffusion auf RPI Zero 2 und die Methode dahinter. Es ist zu beachten, dass die Laufgeschwindigkeit zwar langsamer ist, es sich jedoch um einen neuen Versuch handelt, große Modelle auf kleineren und eingeschränkteren Geräten auszuführen.

Netizens finden dieses Projekt cool

Ausführen einer stabilen Diffusion auf einem Raspberry Pi Zero 2

Der VAE-Decoder ist das einzige Modell in Stable Diffusion, das nicht in RPI Zero 2 RAM mit einfacher oder halber Präzision passt. Dies ist auf Restverbindungen, sehr große Tensoren und Faltungen im Modell zurückzuführen. Die einzige Lösung ist die statische Quantisierung (8 Bit).

Die folgenden Bilder wurden mit OnnxStream mit unterschiedlichen Genauigkeiten des VAE-Decoders aus der Stable Diffusion-Beispielimplementierung generiert, die im Repo des Autors enthalten ist.

Das erste Bild wurde auf dem PC des Autors mit demselben latenten Bild erstellt, das vom RPI Zero 2 erzeugt wurde.

Erzeugungseffekt des VAE-Decoders mit W16A16-Präzision 

 Erzeugungseffekt des VAE-Decoders mit W8A32-Präzision

Das dritte Diagramm wurde von RPI Zero 2 in etwa 3 Stunden erstellt.

Legende: Der Generierungseffekt des VAE-Decoders mit W8A8-Präzision

Funktionen von OnnxStream

• Die Inferenz-Engine ist vom WeightsProvider entkoppelt

• WeightsProvider kann DiskNoCache, DiskPrefetch oder benutzerdefiniert sein

• Aufmerksamkeitsscheibe

• Dynamische Quantisierung (8 Bit vorzeichenlos, asymmetrisch, Perzentil)

• Statische Quantisierung (W8A8 vorzeichenlos, asymmetrisch, Perzentil)

• Quantitative Modelle einfach kalibrieren

• FP16-Unterstützung (mit oder ohne FP16-Betrieb)

• 24 ONNX-Operatoren implementiert (die am häufigsten verwendeten Operatoren)

• Operationen werden sequentiell ausgeführt, aber alle Operatoren sind Multithreading

• Einzelne Implementierungsdatei + Header-Datei

• XNNPACK-Aufrufe werden in die XnnPack-Klasse eingeschlossen (zur zukünftigen Ersetzung).

Und beachten Sie, dass OnnxStream auf XNNPACK angewiesen ist, um bestimmte Grundelemente zu beschleunigen: MatMul, Convolution, elementweises Add/Sub/Mul/Div, Sigmoid und Softmax.

Leistungsvergleich

Stable Diffusion besteht aus drei Modellen: einem Text-Encoder (672 Operationen und 123 Millionen Parameter), einem UNET-Modell (2050 Operationen und 854 Millionen Parameter) und einem VAE-Decoder (276 Operationen und 49 Millionen Parameter).

Unter der Annahme, dass die Stapelgröße gleich 1 ist, sind 10 Schritte erforderlich, um ein vollständiges Bild zu generieren, was 2 Durchläufe des Text-Encoders, 20 (d. h. 2*10) Durchläufe des UNET-Modells und 1 Durchlauf des VAE-Decoders erfordert, um gute Ergebnisse zu erzielen (mit Euler Ancestral-Planer).

Die Tabelle zeigt die unterschiedlichen Inferenzzeiten für die drei Modelle der stabilen Diffusion sowie den Speicherverbrauch (d. h. Peak Working Set Size in Windows oder Maximum Resident Set Size in Linux).

Es kann festgestellt werden, dass im UNET-Modell (mit aktivierter FP16-Arithmetik in OnnxStream bei Ausführung mit FP16-Präzision) der Speicherverbrauch von OnnxStream nur 1/55 des Speicherverbrauchs von OnnxRuntime beträgt, die Geschwindigkeit jedoch nur 0,5-2-mal langsamer ist.

Bei diesem Test sind einige Dinge zu beachten:

• Die erste Ausführung von OnnxRuntime ist eine Aufwärminferenz, da ihre InferenceSession vor der ersten Ausführung erstellt und in allen nachfolgenden Ausführungen wiederverwendet wird. Und OnnxStream hat keine Aufwärmgründe, da sein Design rein „eifrig“ ist (spätere Läufe können jedoch vom Cache der Gewichtsdateien des Betriebssystems profitieren).

• Derzeit unterstützt OnnxStream keine Batchgröße! = 1-Eingabe, was sich von OnnxRuntime unterscheidet, das den gesamten Diffusionsprozess durch die Verwendung von Batchgröße = 2 beim Ausführen des UNET-Modells erheblich beschleunigen kann.

• Beim Testen hatte das Ändern der SessionOptions von OnnxRuntime (z. B. EnableCpuMemArena und ExecutionMode) keine spürbaren Auswirkungen auf die Ergebnisse.

• Die Leistung von OnnxRuntime ist hinsichtlich Speicherverbrauch und Inferenzzeit der von NCNN (einem anderen Framework) sehr ähnlich.

• Getestete Betriebsbedingungen: Windows Server 2019, 16 GB RAM, 8750H CPU (AVX2), 970 EVO Plus SSD, 8 virtuelle Kerne auf VMWare.

Attention Slicing und Quantisierung

Der Einsatz von „Attention Slicing“ und die Verwendung der W8A8-Quantisierung für den VAE-Decoder beim Ausführen des UNET-Modells waren entscheidend, um den Speicherverbrauch des Modells auf ein Niveau zu reduzieren, das für die Ausführung auf RPI Zero 2 geeignet ist.

Während es im Internet viele Informationen über quantisierte neuronale Netze gibt, gibt es nur sehr wenig über „Attention Slicing“.

Die Idee hier ist einfach: Das Ziel besteht darin, die Generierung einer vollständigen Q@K^T-Matrix zu vermeiden, wenn die skalierte Skalarproduktaufmerksamkeit für verschiedene Mehrkopfaufmerksamkeiten im UNET-Modell berechnet wird. Wenn im UNET-Modell die Anzahl der Aufmerksamkeitsköpfe 8 beträgt, beträgt die Form von Q (8,4096,40), während K^T (8,40,4096) beträgt. Daher ist die endgültige Form des ersten MatMul (8,4096,4096), was einem 512-MB-Tensor (FP32-Präzision) entspricht. whaosoft  aiot  http://143ai.com

Die Lösung besteht darin, Q vertikal aufzuteilen und dann Aufmerksamkeitsoperationen normal für jeden Q-Block durchzuführen. Die Form von Q_sliced ​​ist (1,x,40), wobei x (in diesem Fall) 4096 ist, geteilt durch onnxstream::Model::m_attention_fused_ops_parts (Standard ist 2, kann aber angepasst werden).

Dieser einfache Trick reduziert den Gesamtspeicherverbrauch des UNET-Modells von 1,1 GB auf 300 MB, wenn es mit FP32-Präzision ausgeführt wird. Eine effizientere Alternative ist die Verwendung von FlashAttention. FlashAttention erfordert jedoch das Schreiben eines benutzerdefinierten Kernels für jede unterstützte Architektur (AVX, NEON) usw., wobei XnnPack im vom Autor angegebenen Beispiel umgangen wird.

Referenzlink:

https://www.reddit.com/r/MachineLearning/comments/152ago3/p_onnxstream_running_stable_diffusion_in_260mb_of/

https://github.com/vitoplantamura/OnnxStream