Technische Samenvatting: Subnetwork Data Parallelism (SDP)

Probleemstelling

De snelle schaling van diepe neurale netwerken heeft geleid tot ongekende geheugeneisen op versnellers, wat vaak gedistribueerde trainingsstrategieën noodzakelijk maakt. Hoewel Data Parallelism (DDP) breed wordt gebruikt, repliceert het het volledige model op elke GPU, wat leidt tot een hoge geheugenoverhead en dure gradiënt-synchronisatie. Daarentegen splitst Model Parallelism (bijv. Pipeline Parallelism, Tensor Parallelism) het model over apparaten om grotere architecturen te laten passen, maar dit introduceert aanzienlijke communicatiekosten voor activaties en lijdt vaak onder pipeline-bubbles en load-imbalans.

Huidige oplossingen zoals Fully Sharded Data Parallelism (FSDP) en ZeRO verminderen het geheugen per apparaat door parameters, gradiënten en optimizer-states te sharden, maar ze brengen nog steeds aanzienlijke communicatie-overhead met zich mee, met name tijdens de gradiënt-synchronisatie. Er blijft een open uitdaging om de activatie-bandbreedte en het geheugengebruik per node te verminderen zonder de convergentiekwaliteit op te offeren of hoogwaardige interconnects te vereisen.

Methodologie: Subnetwork Data Parallelism (SDP)

De auteurs stellen Subnetwork Data Parallelism (SDP) voor, een gedistribueerd trainingsframework dat een model partitioneert in gestructureerde subnetwerken die getraind worden over workers zonder activaties uit te wisselen. In tegen tegenstelling tot pipeline-benaderingen die berekening sequentieel splitsen, wijst SDP elke worker een structureel volledig deel van het model toe (een subnetwork) dat een volledige pad van input naar loss behoudt, wat onafhankelijke gradiëntberekening mogelijk maakt.

Kernmechanismen

1. Subnetwork Constructie:

   • Block-Level (B-SDP): Verwijdert volledige architecturale blokken (bijv. Transformer-blokken of ResNet basic blocks).
   • Neuron/Channel-Level (N-SDP): Verwijdert specifieke neuronen in fully connected lagen of kanalen in convolutionele lagen, waarbij consistentie tussen aangrenzende lagen wordt gewaarborgd.
   • Workers synchroniseren overlappende parameters via stapsgewijze averaging.

2. Masking Regimes:
   Het paper onderzoekt twee complementaire masking-strategieën om te definiëren hoe subnetwerken worden benut tijdens training:

   • Forward-Masking: Het model wordt geëvalueerd met gemaskeerde parameters ($m_{fwd} = m$). Dit verwijdert parameters, activaties en gradiënten voor de gemaskeerde componenten, wat aanzienlijke geheugen- en rekenbesparingen oplevert. Echter, de gradiënten reflecteren de gemaskeerde forward pass.
   • Backward-Masking: De forward pass gebruikt het volledige model ($m_{fwd} = m_{uni}$), terwijl sparsity wordt toegepast tijdens de backpropagation en aggregatie ($m_{bwd} = m$). Dit behoudt onbevooroordeelde gradiënt-schattingen (theoretisch onderbouwd) en biedt een principieel baseline, hoewel het primair geheugen bespaart op gradiënten en optimizer-states in plaats van op activaties.

3. Theoretische Basis:
   De auteurs bieden een convergentieanalyse voor het backward-masked regime in $L$-smooth settings. Ze demonstreren dat de convergentiesnelheid wordt bepaald door de "masking error", gedefinieerd als de afstand tussen de uniforme masking en de specifieke mask-distributie. De analyse koppelt de convergentiekwaliteit aan de spectrale kloof (spectral gap) van de mask-graph, waarbij ze aantonen dat overlappende parameter-toewijzingen met periodieke averaging partiële synchronisatie behouden.

Belangrijkste Bijdragen

1. Nieuw Paradigma: Introductie van SDP als een geheugenefficiënt gedistribueerd trainingsparadigma dat de training van subnetwerken over nodes verdeelt, waardoor het geheugen per device met 28%–60% wordt verminderd terwijl de prestaties gelijk blijven of verbeteren.
2. Theoretische Garanties: Een theoretische basis die de convergentie van backward-masked SDP koppelt aan een spectral-gap conditie op de mask-graph, waarmee bewezen wordt dat onbevooroordeelde gradiënten behouden blijven onder specifieke masking-condities.
3. Composability: Demonstratie dat SDP orthogonaal is aan bestaande geheugenbesparende technieken. Het verwijdert fysiek modelcomponenten, waardoor parameters, gradiënten, optimizer-states en activaties gelijktijdig worden verkleind. Dit maakt het mogelijk dat SDP schoon componeert met FSDP (ZeRO-3) en Activation Checkpointing (AC).
4. Generaliseerbaarheid: Validatie van SDP over diverse architecturen, inclusief grote taalmodellen (LLaMA) en beeldclassificatiemodellen (ResNet, Swin Transformer), wat aantoont dat het een "drop-in" alternatief is voor standaard data parallelism.

Experimentele Resultaten

Large Language Models (LLaMA)

Experimenten werden uitgevoerd op FineWeb met 134M, 500M en 1B parameter LLaMA-modellen.

• Geheugenefficiëntie: Op de 1B schaal verminderde B-SDP met 69% actieve componenten ($C=0.69$) het piek GPU-geheugen met 28% vergeleken met DDP, terwijl een gelijke of lagere validatie loss werd behaald. Op de 500M schaal verminderde B-SDP met 50% actieve componenten het geheugen met 40%.
• Downstream Performance: Op vijf downstream benchmarks (ARC-E, BoolQ, HellaSwag, OBQA, SciQ) kwamen alle SDP-varianten binnen ~1 punt op de macro averages overeen met DDP, waarbij sommige varianten (B-SDP) lichte verbeteringen vertoonden op specifieke taken.
• Composability: Het combineren van B-SDP met FSDP en Activation Checkpointing bereikte een piek geheugengebruik van 9.7 GB per device voor het 1B model, een 8else reductie vergeleken met DDP (70.9 GB) en strikt lager dan FSDP+AC alleen (11.0 GB).
• Communicatie: SDP verminderde het communicatievolume per stap met ~28% (het verplaatsen van 2.253 MB versus 3.143 MB voor DDP). In inter-node settings verminderde de combinatie van B-SDP en FSDP de wall-clock trainingstijd met 19.9% vergeleken met DDP.

Beeldclassificatie

Experimenten werden uitgevoerd op CIFAR-10 en CIFAR-100 met ResNet-18, WideResNet-18 en Swin-Tiny.

• ResNet-18: SDP-configuraties behaalden competitieve of hogere nauwkeurigheid dan DDP met slechts 40% van het DDP-geheugenbudget. N-SDP (neuron-level) vertoonde een sterk regularisatie-effect en presteerde beter dan DDP bij 64% geheugengebruik.
• Swin Transformer: B-SDP kwam overeen met de DDP-nauwkeurigheid op CIFAR-10 met 68% van het geheugen en verbeterde de nauwkeurigheid met 2% op CIFAR-100 met een reductie van 32% in geheugen.
• Sparsity Limieten: Hoewel forward masking (B-SDP/N-SDP) goed presteerde bij matige sparsity, bleek Backward-Masking (Bb-SDP) superieur bij extreme sparsity (bijv. $C=3/8$), waardoor modelinstorting werd voorkomen waar andere varianten faalden.

Betekenis en Claims

Het paper claimt dat SDP een praktische "derde as" van parallellisme biedt, complementair aan data- en modelparallelisme. De primaire betekenis ligt in het vermogen om:

1. Geheugen te ontkoppelen van Bandbreedte: Door de noodzaak om activaties tussen workers uit te wisselen te elimineren, verzacht SDP de bandbreedte-bottlenecks die vaak pipeline- en tensorparallelisme beperken.
2. FLOP-Matched Scaling mogelijk te maken: SDP staat het trainen van grotere modellen of langere sequenties toe onder hetzelfde hardwarebudget. De lichte reductie in de per-iteratie convergentiesnelheid wordt gecompenseerd door het vermogen om het aantal iteraties te verhogen (FLOP-matched) dankzij de verminderde per-stap compute en communicatie.
3. Een Drop-in Oplossing te bieden: SDP functioneert als een drop-in vervanging voor standaard data parallelism, waarbij geen wijzigingen nodig zijn in de onderliggende gedistribueerde infrastructuur, behalve de masking-logica.

De auteurs benadrukken dat SDP bijzonder waardevol is in scenario's waarin geheugen en bandbreedte de primaire beperkingen zijn, wat een schaalbaar pad biedt voor het trainen van grote modellen op geheugenbeperkte versnellers zonder de complexiteit van complexe pipeline-scheduling of de communicatie-overhead van volledige sharding.