Den Spitzenspeicher sicherer Transformer-Berechnung schrittweise senken

Hintergrund

Mit Secure Neural Network Inference kann ein Client ein Modell nutzen, das ein Server hält, ohne dass eine Seite die Daten der anderen sieht: Die Eingabe des Clients und die Gewichte des Servers bleiben verborgen. Die Systeme rechnen auf Secret Shares, Schlüsseln für Function Secret Sharing oder homomorphen Chiffraten. Diese Darstellungen sind viel größer als die Klartextwerte, sodass der Speicher zur begrenzenden Ressource werden kann. Im unveränderten System SHARK hält jede der beiden Online-Parteien für BERT-base ihre gesamte Preprocessing-Datei im Speicher, zusammen etwa 108 GB. Das passt nicht auf einen Cluster-Knoten mit 95 GB.

In meiner Masterarbeit, Optimizing the Memory Footprint of Secure Neural Network Inference for Transformer Architectures (Universität Münster), habe ich den Spitzenspeicher dieser Forschungssysteme gemessen, ihn in die Objekte zerlegt, die ihn belegen, und ihn mit einer Änderung pro Schritt gesenkt. Das Begleit-Repository liegt unter github.com/Jwrede/snni-memory-reduction (Tag thesis-submission).

Was reduziert wurde

Das Repository enthält 15 Linien über neun Systeme und deren Dealer-, GPU- und ViT-Varianten. Sie decken vier Paradigmen ab: Function Secret Sharing (SHARK, LLAMA, SIGMA), arithmetisches Secret Sharing (SHAFT, PUMA), hybride Verfahren aus HE und MPC (BumbleBee, BOLT) und vollhomomorphe Verschlüsselung (MOAI, ARION). Jede Linie rechnet BERT-base mit Sequenzlänge 128, zwei Transfer-Linien rechnen ViT-Base/16 auf SHAFT. Alle Läufe liefen auf dem PALMA-Cluster der Universität Münster.

Zwei Beispiele aus den veröffentlichten STEPS.md-Dateien:

Das Verfahren

Jeder Schritt ändert eine Sache und wird auf einer festen Maschine gemessen. Die Regel für die Wahl des nächsten Schritts ist im README festgelegt:

1. Die Spitze messen und im Moment der Spitze in die Objekte zerlegen, die sie belegen.
2. Die Objekte nach residenter Größe ordnen und das größte nehmen, für das es eine kostenlose Korrektur gibt, also eine, deren Laufzeitänderung innerhalb der gemessenen Streuung bleibt.
3. Neu messen und wiederholen. Gibt es keine kostenlose Korrektur mehr, folgen kostenpflichtige, geordnet nach eingespartem Speicher pro Einheit Laufzeit.

Die Korrekturen fallen in vier Arten: ein totes Objekt freigeben, von einem doppelt gehaltenen Objekt eine Kopie behalten, ein Objekt beim Lesen neu laden und ein Objekt erst bei seiner ersten Verwendung erzeugen. Eine Änderung, die Genauigkeit, Sicherheit oder Portabilität opfert, ist nie ein Schritt; sie steht getrennt in der off-path.md des jeweiligen Systems.

Spitze und Objekttabelle stammen aus demselben Lauf. Ein nativer Recorder protokolliert jede Allokation ab 64 KiB mit ihrer Aufrufstelle und liest ihre residenten Seiten aus /proc/<pid>/pagemap. Für jede Aufnahme wird der Prozess mit SIGSTOP angehalten, sodass Objekte und Mappings einen einzigen Zeitpunkt beschreiben. Jeder Schritt läuft dreimal, veröffentlicht wird das Replikat mit dem Median. Ein Korrektheitsgate vergleicht die Ausgabe jedes Schritts mit der Ausgabe der Baseline auf einer deklarierten Stufe, von byte-identisch bis gleich innerhalb einer numerischen Toleranz. Schritte, die Arbeit neu aufteilen, ändern die Protokoll-Zufälligkeit und damit die Ausgabe; sie sind gesondert markiert und stützen sich auf ein unverändertes Transkript und einen Klartext-Test.

Welche Schritte am meisten brachten

In SHARK liest die Baseline die gesamte Preprocessing-Datei (das korrelierte Schlüsselmaterial des Offline-Dealers) vor Beginn der Berechnung in einen Puffer. Der Schritt s3_key_stream liest die Datei sequenziell, während sie verbraucht wird. Dieser eine Schritt senkte die Spitze um 52,53 GiB, von 53,68 GiB auf 1,14 GiB, bei gemessenen 5,1 % mehr Laufzeit; damit gilt er als kostenpflichtig. Zwei weitere Schritte brachten die Spitze auf 0,45 GiB, gegenüber einem hergeleiteten Ziel von 217,4 MB für Schlüsselmaterial und Puffer eines Trunkierungs-Batches.

In BumbleBee war der größte kostenlose Schritt b1_dot_encode_chunk, der Matrizen in Scheiben kodiert, statt einen ganzen Operanden auf einmal. Er senkte die Spitze um 6,51 GiB bei einer Laufzeitänderung von 1,0 %. Die späteren Schritte b5 bis b8 senken die Zahl paralleler Oblivious-Transfer-Instanzen von 16 auf 1. Jede Instanz hält eigene Puffer. Diese Schritte senkten die Spitze daher von 6,44 GiB auf 1,39 GiB, und die Laufzeit stieg von 388 s auf 1342 s. steps.csv führt alle vier als kostenpflichtig.

In beiden Linien kamen die großen Reduktionen von Objekten, die vollständig gehalten wurden, obwohl die Berechnung sie stückweise nutzte: die Schlüsseldatei, kodierte Matrizen und der Protokollzustand pro Instanz.

Die Zahlen nachrechnen

Jede Zeile einer steps.csv und jede peak-objects.md wird aus den Rohlogs erzeugt, die daneben liegen. Diese Befehle prüfen, dass sich die veröffentlichten Dateien neu erzeugen lassen und dass jeder Schritt dem Verfahren folgt:

python3 tools/make_steps.py systems/<system> --check
python3 tools/make_impl.py systems/<system> --check
bash tools/test_conformance.sh
bash tools/check_policy.sh

Container-Images gehören nicht zum Repository. Die BUILD.md jedes Schritts nennt den festgelegten Upstream-Commit und die Build-Befehle, und die RUN_META jedes Laufs hält den Image-Digest fest.

Code, Daten und Dokumentation liegen unter github.com/Jwrede/snni-memory-reduction.