Technische Samenvatting: Schalen van de verificatie van neurale netwerken met Tensor Parallelisme en Fully Sharded Data Parallelism

1. Probleemstelling

Formele verificatie van neurale netwerken streeft ernaar wiskundig te bewijzen dat een netwerk voldoet aan veiligheidseigenschappen voor alle inputs binnen een gespecificeerd domein. Hoewel algoritmen voor grensbepaling (bound-propagation) zoals IBP, CROWN en $\alpha$-CROWN de state-of-the-art zijn voor deze taak, wordt hun praktische toepassing ernstig beperkt door GPU-geheugen. Standaardimplementaties vereisen dat gewichtsmatrices en lineaire-relaxatie $A$-matrices volledig op één accelerator aanwezig zijn. Voor grote netwerken (bijv. een enkele verborgen laag met breedte $h=262.144$) kunnen de gewichtsmatrices alleen al meer dan 4 GB in float32 beslaan, waardoor er onvoldoende geheugen overblijft voor de $A$-matrices, batching, of intermediaire tensoren die nodig zijn voor verificatie. Dit artikel adresseert de noodzaak om de single-GPU geheugenbottleneck te verwijderen door parameter-sharding strategieën uit grootschalige modeltraining aan te passen aan het verificatieframework.

2. Methodologie

De auteurs passen twee parallellismetechnieken toe—Tensor Parallelisme (TP) en Fully Sharded Data Parallelism (FSDP)—op het auto LiRPA / $\alpha, \beta$-CROWN verificatieframework.

2.1 Tensor Parallelisme (TP)

TP splitst matrixvermenigvuldigingen binnen een enkele laag over $P$ GPU's.

• Implementatie: De auteurs hebben aangepaste BoundLinearTP Col en BoundLinearTP Row operatoren geïmplementeerd. Gewichten worden geshard langs de output- (kolom) of inputdimensie (rij).
• Mechanisme: In een Column–Row paar wordt de berekening verdeeld, waarbij één AllReduce operatie nodig is om het resultaat te herstellen.
• Beperking: Wanneer ReLU-activaties optreden binnen een gesharde zone (tussen de Column en Row lagen), vereist de CROWN backward pass intermediaire grenzen om de hellingen van de relaxatie ($\lambda_i$) te bepalen. Het exact berekenen hiervan zou bij elke stap inter-GPU communicatie vereisen. Om de lokaliteit te behouden, vervangen de auteurs de complexere methode door de simpelere Interval Bound Propagation (IBP) voor deze intermediaire grenzen. Dit introduceert een "wrapping effect", waardoor de nauwkeurigheid van de grenzen (bound tightness) afneemt naarmate de diepte van de gesharde zone toeneemt.

2.2 Fully Sharded Data Parallelism (FSDP)

FSDP versnippert (shards) parametertensoren gelijkmatig over $P$ GPU's, maar reconstrueert ze voor de berekening.

• Implementatie: Geïmplementeerd met ongeveer 100 regels Python, vervangt FSDP BoundParams tensoren door shards. Voordat een laagberekening plaatsvindt, assembleert een AllGather de volledige gewichtsmatrix; direct daarna wordt de volledige tensor vrijgegeven, zodat alleen de shard behouden blijft.
• Mechanisme: Deze aanpak zorgt ervoor dat er op elk gegeven moment maximaal één volledige gewichtsmatrix in het GPU-geheugen aanwezig is. Cruciaal is dat FSDP de volgorde of het type rekenkundige operaties niet verandert; het verandert enkel de opslaglocatie van de gewichten.
• Correctheid: Omdat de rekenkunde identiek blijft aan de single-GPU baseline, produceert FSDP bitwise-identieke grenzen ($\Delta = 0$).
• Integratie: FSDP integreert naadloos met volledige verificatie ($\beta$-CROWN + Branch-and-Bound) en convolutionele lagen (BoundConv).

2.3 Experimentele Opstelling

Experimenten werden uitgevoerd op NVIDIA A40 GPU's met behulp van PyTorch 2.x en de auto LiRPA bibliotheek. De auteurs introduceerden een bab.force_synchronous flag om eerlijke vergelijkingen tussen single-GPU en gedistribueerde modi te garanderen door automatische batch-resizing en pruning in de single-GPU path uit te schakelen, waardoor identieke workloads worden gewaarborgd.

3. Belangrijkste Bijdragen

1. TP Implementatie voor Verificatie: De auteurs hebben TP voor auto LiRPA geïmplementeerd met aangepaste operatoren, waarbij zij ongeveer $2\times$ reductie in piekgeheugen bereikten bij $P=2$. Zij hebben de trade-off formeel geanalyseerd en bevestigd dat hoewel de soundheid behouden blijft, de nauwkeurigheid van de grenzen afneemt met het aantal gesharde zones door de IBP-substitutie.
2. FSDP Implementatie: Zij hebben FSDP geïmplementeerd voor grensbepaling, wat een reductie van 80–90% in baseline geheugen en 34–39% in piekgeheugen oplevert op brede MLPs. Cruciaal is dat deze methode bitwise-identieke grenzen produceert ten opzichte van de single-GPU baseline.
3. Integratie met Volledige Verificatie: FSDD is succesvol geïntegreerd met $\beta$-CROWN en Branch-and-Bound (BaB), wat de eerste volledige verificatie (UNSAT resultaat) van een CIFAR-100 ResNet-large model onder FSDP mogelijk maakte.
4. Identificatie van Bottlenecks: Door uitgebreide experimenten hebben de auteurs geïdentificeerd dat in $\alpha$-CROWN + BaB modus de primaire geheugenbottleneck niet de gewichtsmatrices zijn, maar de per-neuron alpha tensoren (ReLU hellingen).

4. Resultaten

4.1 Geheugenreductie

• TP: Bereikte een piekgeheugen reductiefactor van $\approx 1.985$ (bijna de theoretische $P=2$) voor een brede MLP ($h=262,144$).
• FSDP:
  • Baseline Geheugen: Verminderd met 80–90% over diverse modelgroottes, aangezien alleen shards worden gealloceerd vóór de grensbepaling.
  • Piekgeheugen: Verminderd met 34–39% op brede MLPs ($h \ge 4096$). De besparingen zijn lager dan de baseline besparingen omdat de niet-gesharde $A$-matrices (die schalen met de breedte in het kwadraat) in het geheugen blijven.

4.2 Grens-nauwkeurigheid en Soundheid

• TP:
  • Modellen met één gesharde zone (geen intermediaire ReLU) bereikten overeenstemming op machineprecisie ($\sim 10^{-8}$) met de single-GPU referentie.
  • Modellen met meerdere gesharde zones vertoonden significante afwijking (bijv. $\sim 750$ eenheden voor een 3-zone model) door IBP-substitutie, hoewel de soundheid (éénzijdige overschatting) behouden bleef.
• FSDP:
  • Bereikte exact nul afwijking ($\Delta lb = \Delta ub = 0.00$) over alle geteste modellen (MLPs en ONNX benchmarks) voor zowel IBP als CROWN.
  • Verifieerde succesvol een Transformer model (ViT) en convolutionele ResNets met bitwise-identieke resultaten vergeleken met de baseline.

4.3 Volledige Verificatie (BaB)

In $\beta$-CROWN + BaB modus wordt het geheugengebruik gedomineerd door alpha tensoren (bijv. 1,5 GB voor alpha tensoren versus 1,5 MB voor gewichten in een specifiek MNIST netwerk).

• Het sharden van gewichten via FSDP leverde verwaarloosbare piekgeheugenbesparingen op in deze modus, omdat de alpha tensoren niet werden geshard.
• Echter, FSDP maakte de verificatie van een ResNet-large model (VNN-COMP 2024) mogelijk die anders onhaalbaar of moeilijk te beheren zou zijn, waarmee een UNSAT eigenschap werd bevestigd zonder BaB branching.

5. Betekenis en Beperkingen

Het artikel demonstreert dat parameter-sharding strategieën uit de training met een bescheiden engineering-inspanning kunnen worden overgedragen naar formele verificatie.

• Betekenis: FSDP biedt een pad om grotere netwerken te verifiëren met exacte grenzen, terwijl TP de hoogste geheugenreductie biedt ten koste van de nauwkeurigheid van de grenzen voor diepe netwerken. Het werk vestigt dat FSDP compatibel is met de complexe eisen van volledige verificatie (BaB) en diverse architecturen (CNN's, Transformers).
• Beperkingen:
  • Schaalbaarheid: Experimenten waren beperkt tot $P=2$. De auteurs merken op dat naarmate $P$ toeneemt, de niet-gesharde $A$-matrices waarschijnlijk de dominante kostenpost zullen worden, wat leidt tot verminderde meeropbrengsten.
  • TP Nauwkeurigheid: De degradatie van de nauwkeurigheid van de grenzen in TP voor diepe netwerken beperkt de bruikbaarheid voor precieze verificatie van diepe architecturen.
  • Verschuiving van Bottleneck: De experimenten laten zien dat toekomstig multi-GPU verificatieonderzoek zich moet richten op het sharden van per-neuron alpha tensoren, niet alleen gewichten, om significante schaling in volledige verificatiemodi te bereiken.

De auteurs concluderen dat hoewel FSDP de gewichtsopslag-bottleneck oplost, de alpha tensor-bottleneck de kritieke uitdaging blijft voor het schalen van volledige verificatie naar grotere modellen.