LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs

Stel je voor dat je een robot probeert te leren de wereld te begrijpen door hem een enorme, verwarde bol wol te laten zien. Elke knoop in de wol is een persoon, een plaats of een ding, en de draden die hen verbinden zijn hun relaties. Dit is hoe computers leren over complexe netwerken zoals sociale media-vrienden, wetenschappelijke citaties of aanbevelingssystemen. Dit vakgebied wordt Graph Neural Networks (GNN's) genoemd. De robot moet naar een knoop kijken, zien met wie die verbonden is, naar de verbindingen van die knopen kijken, enzovoort, om het hele plaatje te begrijpen.

Het probleem is dat naarmate de bol wol groter wordt, het onmogelijk wordt om de hele bol op het bureau van de robot te passen. Als de robot probeert de hele bol in zijn handen (zijn geheugen) te houden om deze te bestuderen, raakt hij de ruimte kwijt. Daarom snijden wetenschappers de wol meestal in kleinere, hanteerbare stukken die "mini-batches" worden genoemd om ze één voor één te bestuderen. Maar hier is de crux: als de wol in een groot magazijn ver weg wordt bewaard, en de robot in een klein kantoor zit, moet de robot telkens weer heen en weer rennen naar het magazijn om deze stukken te pakken. De tijd die wordt besteed aan het heen en weer rennen (het versturen van gegevens over het netwerk) duurt vaak langer dan de tijd die nodig is om de knopen daadwerkelijk te bestuderen. Dit is de "communicatie-bottleneck" die alles vertraagt.

Dit is waar een nieuw idee genaamd LGNNIC om de hoek komt kijken. De onderzoekers stelden een simpele vraag: wat als we de wol niet alleen in het magazijn zouden bewaren, maar ook een slimme schaar en een kleine, supersnelle assistent direct naast de wol zouden geven? In plaats van dat de robot helemaal naar het magazijn rent om een enorme brok wol te pakken en deze vervolgens op maat te knippen, zou de assistent in het magazijn het knippen en inkrimpen kunnen doen voordat de robot er zelfs maar om vraagt.

Het artikel, getiteld "LGNNIC: Acceleration of Large-Scale GNN Training using SmartNICs", onderzoekt precies dit. Het team heeft een systeem gebouwd waarbij het "magazijn" (een afgelegen geheugenknooppunt) beschikt over een speciale, intelligente netwerkkaart genaamd een SmartNIC (specifiek een NVIDIA BlueField-2). Deze SmartNIC fungeert als die slimme assistent. Hij zit direct naast de data en voert twee magische trucs uit voordat hij iets naar de hoofdcomputer (het "trainingsknooppunt" met de krachtige GPU) stuurt:

1. Neighbor Sampling (Buur-sampling): In plaats van de robot een enorme, rommelige brok van de grafiek te sturen, knipt de SmartNIC de onnodige verbindingen weg en houdt alleen de meest relevante buren over die de robot moet bestuderen. Dit verandert een gigantisch, zwaar bos wol in een klein, net pakketje.
2. Quantization (Kwantisering): De SmartNIC verkleint ook de omvang van de informatie binnen het pakketje. Het zet de data om van een zwaar, hoog-precisie formaat (32-bit floating-point) naar een lichter, iets minder precies formaat (16-bit floating-point). Denk aan het comprimeren van een high-definition video naar een kleiner bestand dat er nog steeds geweldig uitziet, maar de helft minder ruimte inneemt.

De onderzoekers testten deze opstelling met echte datasets zoals Reddit (een enorm forum) en netwerken van academische papers. Ze ontdekten dat door de SmartNIC het zware werk te laten doen van het knippen en inkrimpen van de data, ze de hoeveelheid informatie die over het netwerk moest reizen, drastisch konden verminderen.

De resultaten waren indrukwekkend. Wanneer ze een standaard, tragere methode gebruikten om data te verplaatsen (zoals het versturen van e-mails via een reguliere internetverbinding, wat ze "Sockets" noemen), maakte de voor-verkleining en het inkrimpen door de SmartNIC de trainingsprocedure tot wel 62,4 keer sneller voor sommige taken. Zelfs met een snellere, directere verbindingsmethode (genaamd "DOCA-DMA"), zagen ze nog steeds snelheidsverbeteringen van wel 17,5 keer. De "inkrimpingstruek" (kwantisering) voegde er nog meer snelheid aan toe, waardoor transfers tot 3,6 keer sneller waren met de standaardmethode en 1,3 keer sneller met de directe methode.

Cruciaal is dat het artikel opmerkt dat hoewel de SmartNIC zelf langzamer is bij het knippen vergeleken met een superkrachtige hoofdcomputer, de tijd die wordt bespaard door niet een enorme, ongesneden bundel over het netwerk te moeten slepen, de moeite waard is. De "assistent" in het magazijn is langzamer in het werk, maar het bespaart de "loper" (het netwerk) van het dragen van een zware last. De onderzoekers hebben ook gecontroleerd of dit inkrimpen de antwoorden van de robot niet fout maakte; de nauwkeurigheid van de resultaten bleef bijna exact hetzelfde, met slechts minuscule, verwaarloosbare veranderingen.

Kortom, het artikel suggereert dat door een deel van het werk naar de rand van het netwerk te verplaatsen—precies daar waar de data zich bevindt—we kunnen voorkomen dat de robot tijd verspilt aan het heen en weer rennen. Het is een slimme manier om de training van gigantische, complexe AI veel sneller te maken zonder dat er grotere, duurdere computers gebouwd hoeven te worden, simpelweg door slimmer te zijn over hoe we de data verplaatsen.