AEGIS: Scaling Long-Sequence Homomorphic Encrypted Transformer Inference via Hybrid Parallelism on Multi-GPU Systems

Probleemstelling

Volledig Homomorf Versleuteling (FHE) maakt het mogelijk om inferentie op versleutelde data uit te voeren zonder deze te decoderen, wat essentieel is voor privacy in Machine Learning as a Service (MLaaS). Echter, het uitvoeren van Transformer-modellen (zoals BERT) met lange sequenties onder FHE stuit op twee fundamentele beperkingen op multi-GPU-systemen:

1. Geheugenschaarste: Versleutelde gewichten en activeringen groeien exponentieel met de sequentielengte door de vereiste precisie (modulusketens) en "packing" (het verpakken van tokens in ciphertexts). Zelfs voor gematigde werklasten (bijv. 2048 tokens) overschrijdt het geheugenvolume de capaciteit van één GPU.
2. Communicatie-inefficiëntie: Bestaande parallelle strategieën falen bij versleutelde inferentie:
   • Plaintext Tensor Parallelism: Werkt niet omdat ciphertexts vaste "slot"-layouts hebben; het partitioneren van data leidt tot enorme replicatie van ciphertexts of redundante berekeningen.
   • Limb Parallelism (RNS): Verdeelt de berekening over de restantgetalsystemen (RNS) van de ciphertexts. Hoewel dit geheugen bespaart, veroorzaakt het frequente synchronisatie (AllGather) tijdens sleutelwisselingen (Key Switching) en rotaties. Dit leidt tot een communicatie-overhead die de berekening domineert, waardoor het schalen naar meerdere GPU's zelfs tot vertraging kan leiden (bijv. 2.25x langzamer op 2 GPU's).

De kernuitdaging is dat bestaande systemen ofwel applicatie-gericht zijn (ignoreren van cryptografische beperkingen) of cryptografie-gericht zijn (ignoreren van applicatiestructuur), wat leidt tot inefficiënte communicatiepatronen.

Methodologie: AEGIS

De auteurs presenteren AEGIS (Application-Encryption Guided Inference System), een compiler-runtime co-designed framework dat hybride parallelisme toepast door applicatie- en encryptiedomeinen gezamenlijk te optimaliseren. De aanpak bestaat uit drie hoofdpilaren:

1. Dependency-Aware Workload Partitioning (Afhankelijkheidsbewuste Werklastverdeling):

   • In plaats van data willekeurig te verdelen, analyseert AEGIS de afhankelijkheden op twee niveaus:
     • Modulus-coherent plaatsing: Groepeert polynoomdelen die dezelfde modulusketen delen en die samen betrokken zijn bij sleutelwisselingen op dezelfde GPU. Dit elimineert cross-device synchronisatie tijdens deze zware operaties.
     • Token-coherent plaatsing: Zorgt ervoor dat alle ciphertexts die bijbehoren bij één enkel token (bijv. door packing) op dezelfde GPU blijven, tenzij de applicatie expliciet aggregatie over tokens vereist.
   • Dit voorkomt dat applicatie-gerelateerde reducties (zoals bij Attention of FFN) onnodige collectieve communicatie vereisen.

2. Polynomial-Operator Reordering (Polynoom-Operator Hergroepering):

   • AEGIS herschikt de uitvoeringsvolgorde van lage-niveau instructies om communicatie en berekening te overlappen.
   • Bijvoorbeeld: In plaats van te wachten tot alle rotaties voor een token voltooid zijn voordat communicatie begint, worden rotaties met verschillende offsetten verspreid over de GPU's. Hierdoor kunnen collectieve operaties (zoals ReduceScatter voor Softmax) parallel lopen met de berekening van andere delen van de token.
   • Dit verbergt communicatie-latentie deterministisch zonder runtime-heuristieken.

3. Compiler Pass:

   • AEGIS fungeert als een tussenlaag die een hoog-niveau applicatiegraph omzet in een polynoom-instructiegraph. Het analyseert afhankelijkheden, voegt collectieve communicatieprimitieven in op de juiste momenten en genereert een deterministisch uitvoeringsplan voor multi-GPU-systemen.

Belangrijkste Bijdragen

• AEGIS Framework: Het eerste systeem dat applicatie-semantiek en RNS-encryptiebeperkingen gezamenlijk optimaliseert voor schaalbare, lange-sequentiële Transformer-inferentie op multi-GPU.
• Nieuwe Plaatsingsstrategie: Een hiërarchische plaatsingsmethode die eerst RNS-modulusafhankelijkheden en daarna token-coherentie respecteert, wat leidt tot een drastische reductie in communicatievolume.
• Deterministisch Communicatieverbergen: Een mechanisme om communicatie te overlappen met berekening door de volgorde van polynoomoperaties te herschikken, wat wachttijden elimineert.
• Uitgebreide Evaluatie: Een grondige analyse van schaalbaarheid, geheugenefficiëntie en latentie in vergelijking met state-of-the-art baselines (HEBooster, Cinnamon, Hydra, en plaintext Tensor Parallelism).

Resultaten

De prestaties van AEGIS zijn getest op BERT-Base met inputsequenties tot 2048 tokens op NVIDIA GPU-systemen (RTX A6000 en A100):

• Communicatievermindering:
  • Tot 57,9% minder inter-GPU communicatie in Feed-Forward Networks (FFNs).
  • Tot 81,3% minder communicatie in Self-Attention lagen vergeleken met de beste bestaande ontwerpen.
• Schaalbaarheid en Snelheid:
  • Op 4 GPU's bereikt AEGIS een schaalbaarheidsefficiëntie van 96,62%.
  • Er wordt een end-to-end snelheidswinst van 3,86x behaald ten opzichte van single-GPU uitvoering.
  • Vergelijking met baselines: AEGIS is 1,86x sneller dan direct toepassen van plaintext Tensor Parallelism, 2,92x sneller dan HEBooster (Limb Parallelism) en 1,23x sneller dan Hydra.
• Geheugenefficiëntie:
  • AEGIS reduceert het geheugengebruik per apparaat met 69,1% (op 4 GPU's) ten opzichte van systemen die ciphertexts repliceren.
  • Het systeem kan sequenties van 2048 tokens verwerken zonder Out-Of-Memory (OOM) fouten, terwijl andere methoden (zoals Hydra en Cinnamon) al bij 512 of 1024 tokens vastlopen door geheugenoverbelasting.

Betekenis en Conclusie

Dit paper toont aan dat het schalen van versleutelde Transformer-inferentie niet mogelijk is door simpelweg plaintext-parallelle technieken over te nemen. De oplossing vereist een gecoördineerde optimalisatie over de applicatie- en encryptielagen.

AEGIS bewijst dat door de dataafhankelijkheden van het Transformer-model te koppelen aan de structurele beperkingen van de CKKS-encryptie (zoals RNS en modulusketens), het mogelijk is om communicatie-overhead te minimaliseren en geheugen te besparen. Dit legt de basis voor praktische, privacy-bewuste MLaaS-diensten die lange contexten kunnen verwerken zonder dat de gebruiker zijn data hoeft te ontsleutelen. De resultaten suggereren dat toekomstige FHE-systemen "packing", "plaatsing" en "communicatieplanning" als één integraal probleem moeten benaderen.