Tensor Parallellisme
Tensor parallellisme is een techniek waarmee de gigantische rekensommen achter moderne AI-taalmodellen worden opgeknipt en verdeeld over meerdere grafische chips (GPU's, gespecialiseerde processoren die van oorsprong voor games werden ontwikkeld maar uitblinken in de wiskundige berekeningen die AI nodig heeft). In plaats van één GPU het hele werk te laten doen, snijdt tensor parallellisme de berekening letterlijk in stukken, zodat verschillende chips tegelijk aan hetzelfde probleem rekenen.
Een handige vergelijking: stel je een enorme spreadsheet voor met miljarden cellen die allemaal vermenigvuldigd en opgeteld moeten worden. Dat past niet op één beeldscherm en is voor één rekenmachine te traag. Tensor parallellisme knipt die spreadsheet in verticale stroken en geeft elke strook aan een andere rekenmachine. Alle rekenmachines werken tegelijk aan hun eigen strook, en aan het einde worden de deeluitkomsten weer samengevoegd tot één antwoord. Dat is precies wat er gebeurt bij het trainen en gebruiken van de allergrootste taalmodellen zoals GPT-achtige systemen.
Wat is het precies?
Een neuraal netwerk, het type wiskundig model achter chatbots en taalmodellen, bestaat in de kern uit lagen die vooral één ding doen: getallen vermenigvuldigen met een zogeheten gewichtsmatrix. Zo'n matrix is een groot rooster van getallen (de 'geleerde kennis' van het model) waarmee de invoer wordt vermenigvuldigd om tot een uitvoer te komen. Bij de grootste modellen bevatten die matrices honderden miljarden getallen, veel te veel om in het geheugen van één GPU te passen.
Tensor parallellisme lost dit op door zo'n gewichtsmatrix in stukken te knippen, bijvoorbeeld kolomsgewijs, en elk stuk naar een andere GPU te sturen. Elke GPU vermenigvuldigt de invoer met zijn eigen stukje matrix en berekent zo een deel van het antwoord. Vervolgens moeten die deelantwoorden weer bij elkaar worden opgeteld tot het volledige resultaat, een stap die in de techniek 'all-reduce' heet: alle GPU's sturen hun deeluitkomst naar elkaar en tellen die op, zodat iedereen weer verder kan met hetzelfde complete resultaat.
Die uitwisseling moet razendsnel gebeuren, want dit gebeurt niet één keer maar bij vrijwel elke laag van het netwerk, dus miljoenen keren tijdens een training. Daarom zijn GPU's die tensor parallellisme gebruiken meestal verbonden via extreem snelle verbindingen zoals NVIDIA's NVLink, een soort supersnelle datakabel tussen chips die veel meer bandbreedte biedt dan een gewone netwerkkabel.
Het is nuttig om tensor parallellisme te onderscheiden van twee andere technieken waarmee het vaak wordt gecombineerd. Bij data-parallellisme krijgt elke GPU een kopie van het hele model, maar werkt elke GPU aan een ander deel van de trainingsvoorbeelden. Bij pipeline-parallellisme wordt het model in lagen opgeknipt, waarbij de ene GPU bijvoorbeeld de eerste tien lagen doet en de volgende GPU de daaropvolgende tien. Tensor parallellisme is fijnmaziger: het knipt dwars door één enkele laag heen, zodat zelfs de grootste individuele rekenstap over meerdere chips verdeeld kan worden.
Wat wil men ermee bereiken?
Het belangrijkste doel is simpel: sommige taalmodellen zijn zo groot geworden dat ze domweg niet meer in het geheugen van één GPU passen, hoe krachtig die ene chip ook is. Zonder een manier om het model op te splitsen, zou het trainen of draaien van zulke modellen onmogelijk zijn. Tensor parallellisme maakt het mogelijk om modellen met honderden miljarden parameters (de instelbare 'knoppen' die tijdens training worden bijgesteld) toch te trainen, door de rekenlast te verspreiden over tientallen of honderden GPU's tegelijk.
Daarnaast wil men de trainingstijd terugbrengen. Een model met honderden miljarden parameters trainen op één chip zou letterlijk jaren duren. Door de rekenkracht van veel chips tegelijk te benutten, kan die tijd worden teruggebracht tot enkele weken of maanden. Ook bij het daadwerkelijk gebruiken van zo'n model (inference, het moment dat het model een antwoord genereert) zorgt tensor parallellisme voor kortere wachttijden, omdat meerdere chips tegelijk aan hetzelfde antwoord rekenen in plaats van na elkaar.
Tensor parallellisme wordt zelden alleen ingezet. In de praktijk combineren onderzoekers het met data- en pipeline-parallellisme tot wat vaak '3D-parallellisme' wordt genoemd: drie manieren om werk te verdelen, elk voor een ander deel van het probleem, tegelijk toegepast op honderden of duizenden GPU's in een datacenter.
Voorbeelden uit de praktijk
Megatron-LM, ontwikkeld door chipfabrikant NVIDIA, geldt als een van de projecten die tensor parallellisme voor taalmodellen populair maakten. In een paper uit 2019 van onderzoeker Mohammad Shoeybi en collega's liet NVIDIA zien hoe de techniek gebruikt kon worden om GPT-achtige modellen tot ruim acht miljard parameters te trainen door gewichtsmatrices over meerdere GPU's te verdelen.
In 2021 combineerden Microsoft en NVIDIA hun krachten voor Megatron-Turing NLG, een model met 530 miljard parameters. Dit project gebruikte tegelijk tensor-, pipeline- en data-parallellisme, verspreid over honderden GPU's in Microsofts Selene-supercomputer, om een model te trainen dat destijds tot de grootste taalmodellen ter wereld behoorde.
Een opvallend open voorbeeld is BLOOM, een taalmodel met 176 miljard parameters uit 2022, gebouwd door het BigScience-initiatief. Dit was een internationale samenwerking van honderden onderzoekers, georganiseerd via AI-bedrijf Hugging Face. BLOOM werd getraind op de Franse Jean Zay-supercomputer, met steun van de Franse overheid via het rekencentrum GENCI en onderzoeksinstituut CNRS, en maakte gebruik van het Megatron-DeepSpeed framework, dat tensor parallellisme combineert met geheugenbesparende technieken.
Ook Meta gebruikte volgens de eigen publicaties over de LLaMA-modellen (2023) vormen van model- en tensor-parallellisme in de trainingsinfrastructuur om deze modellen efficiënt over grote GPU-clusters te verdelen. En Google past voor eigen modellen zoals PaLM (2022, 540 miljard parameters) een vergelijkbare aanpak toe met een eigen systeem genaamd GSPMD, in combinatie met het Pathways-platform, toegespitst op Googles eigen TPU-chips (Tensor Processing Units, chips die Google speciaal voor AI-berekeningen ontwierp). Over de precieze interne trainingsdetails van sommige recentere, gesloten modellen is minder publiek bekend, dus daarover kunnen alleen algemene uitspraken worden gedaan.
Hoe ver is de techniek?
Tensor parallellisme is inmiddels geen experimentele curiositeit meer, maar een standaardonderdeel van het trainen van vrijwel elk groot taalmodel sinds ongeveer 2020. Raamwerken zoals Megatron-LM en DeepSpeed (ontwikkeld door Microsoft) zijn breed gebruikte, volwassen softwarebibliotheken geworden. Ook het populaire machine-learning raamwerk PyTorch heeft sinds versie 2.x eigen ondersteuning toegevoegd, via een functionaliteit genaamd DTensor en een tensor-parallelle API binnen PyTorch Distributed, waardoor de techniek toegankelijker is geworden voor een breder publiek onderzoekers.
Het grootste technische obstakel blijft communicatie-overhead: omdat GPU's bij elke laag opnieuw resultaten moeten uitwisselen, is tensor parallellisme sterk afhankelijk van zeer snelle verbindingen zoals NVLink. Zulke verbindingen zijn doorgaans alleen beschikbaar tussen GPU's binnen één server, een zogeheten 'node'. Daarom wordt tensor parallellisme in de praktijk meestal beperkt tot GPU's binnen dezelfde node, terwijl de tragere verbindingen tussen verschillende servers worden opgevuld met pipeline- en data-parallellisme, die minder vaak hoeven te communiceren.
Er wordt actief onderzoek gedaan naar het verder verminderen van deze communicatie-overhead, bijvoorbeeld door berekeningen en datauitwisseling slimmer te laten overlappen in tijd. Ook lopen er projecten, zoals het academische Alpa-project van onder meer UC Berkeley, die proberen automatisch te bepalen hoe een model het beste over beschikbare chips verdeeld kan worden, zodat onderzoekers dit niet meer handmatig hoeven te ontwerpen. Dat automatiseren is nog niet volledig opgelost en blijft een actief onderzoeksgebied.
Wie werken eraan?
NVIDIA speelt een centrale rol, zowel met de Megatron-LM software als met de NVLink-hardware die de techniek praktisch mogelijk maakt. Microsoft draagt bij via het DeepSpeed-raamwerk en de samenwerking met NVIDIA aan Megatron-Turing NLG. Google ontwikkelt met GSPMD en Pathways een eigen aanpak, toegesneden op de eigen TPU-chips. Meta AI gebruikt en publiceert over vergelijkbare technieken in de trainingsinfrastructuur van de LLaMA-modellen.
Daarnaast is er een belangrijke rol voor open, niet-commerciële initiatieven. BigScience, georganiseerd via Hugging Face en met financiële en rekenkundige steun van de Franse overheid (via GENCI en het CNRS-rekencentrum Jean Zay), liet met BLOOM zien dat ook academische en non-profitcoalities dit soort grootschalige, verdeelde training kunnen uitvoeren. Onderzoeksgroepen aan universiteiten zoals UC Berkeley en Stanford dragen bij met projecten als Alpa, en het open-source onderzoekscollectief EleutherAI heeft eigen grote taalmodellen getraind met vergelijkbare parallellisatietechnieken.