Trainingscheckpoint: het geheugen van een lerend AI-model
Wie weleens een computerspel heeft gespeeld kent het principe: je slaat je voortgang op, zodat je na een crash of een verloren gevecht niet helemaal opnieuw hoeft te beginnen. Een trainingscheckpoint is in de wereld van kunstmatige intelligentie precies dat, maar dan voor een AI-model dat wordt getraind. Het is een momentopname van alles wat een neuraal netwerk op een bepaald punt in het trainingsproces heeft geleerd, opgeslagen als bestand op een harde schijf.
Stel dat een onderzoeksteam een taalmodel drie weken lang traint op honderden grafische rekenchips (GPU's). Halverwege dag negen valt de stroom uit. Zonder checkpoints zou al het werk van die negen dagen verloren gaan. Met checkpoints kan het team het trainingsproces hervatten vanaf het laatst opgeslagen punt, bijvoorbeeld een paar uur eerder, in plaats van weer bij nul te beginnen. Checkpoints zijn dus geen spectaculaire doorbraaktechnologie, maar een onmisbaar stuk gereedschap dat het trainen van moderne AI-modellen praktisch haalbaar maakt.
Wat is het precies?
Een neuraal netwerk, het type wiskundig model achter de meeste moderne AI-systemen, bestaat uit miljoenen tot honderden miljarden instelbare getallen: de zogeheten parameters of gewichten. Tijdens het trainen worden deze getallen stapsgewijs aangepast, telkens een klein beetje, op basis van voorbeelden die het model te zien krijgt. Dit proces heet gradient descent: het model probeert zijn fouten steeds iets kleiner te maken.
Een checkpoint legt de waarden van al die parameters op een gegeven moment vast in een bestand. Daarbij wordt vaak meer opgeslagen dan alleen de gewichten zelf. Ook de status van de optimizer gaat mee: dat is het onderdeel van het trainingsproces dat bepaalt hoe groot en in welke richting elke aanpassing moet zijn. Sommige optimizers, zoals de veelgebruikte methode Adam, houden daarvoor een soort geheugen bij van eerdere aanpassingen. Zonder dat geheugen mee op te slaan zou het hervatten van training minder soepel verlopen. Verder bevat een checkpoint doorgaans het nummer van de laatste trainingsstap of -ronde (epoch), zodat duidelijk is waar het proces gebleven was.
Checkpoints worden opgeslagen in verschillende bestandsformaten, afhankelijk van het softwareraamwerk dat wordt gebruikt. Bij PyTorch, een populaire ontwikkelomgeving van Meta, zijn dat vaak bestanden met extensie .pt of .pth. Bij TensorFlow en Keras, ontwikkeld door Google, kom je .h5- of .ckpt-bestanden tegen. De laatste jaren wint het formaat .safetensors terrein, omdat het veiliger is: oudere formaten konden bij het inladen ongewenst programmacode uitvoeren, iets wat .safetensors bewust voorkomt.
De bestandsgrootte van een checkpoint hangt direct samen met het aantal parameters van een model. Een klein model kan met een checkpoint van enkele megabytes toe, maar een groot taalmodel met honderden miljarden parameters produceert checkpoints van honderden gigabytes tot enkele terabytes. Dat maakt opslag, verplaatsing en back-up van checkpoints bij de grootste modellen een serieuze technische uitdaging op zich.
Wat wil men ermee bereiken?
Het belangrijkste doel van checkpointing is robuustheid. Trainingsruns voor grote modellen kunnen dagen tot maanden duren en draaien op honderden of duizenden chips tegelijk. De kans dat ergens een chip uitvalt, een server crasht of een stroomstoring optreedt is bij zulke lange, grootschalige runs niet verwaarloosbaar. Regelmatig checkpoints opslaan beperkt de schade: bij een storing gaat hooguit het werk sinds het laatste checkpoint verloren, niet de hele training.
Checkpoints maken het ook mogelijk om een model tussentijds te evalueren. Onderzoekers kunnen op verschillende momenten tijdens de training kijken hoe goed het model al presteert, zonder de training te moeten stoppen of te herhalen.
Een ander belangrijk doel is fine-tuning, ook wel transfer learning genoemd: het verder trainen van een bestaand, al getraind model op een specifieke taak of dataset. In plaats van een compleet nieuw model vanaf nul te trainen, wat enorm veel rekenkracht en dus geld kost, laadt men een bestaand checkpoint in en traint men dat verder met relatief weinig extra data. Dit is de manier waarop veel gespecialiseerde AI-toepassingen tegenwoordig worden gemaakt.
Ten slotte spelen checkpoints een rol bij reproduceerbaarheid: wanneer onderzoekers een checkpoint publiceren, kunnen anderen precies hetzelfde model gebruiken, de resultaten controleren, of erop voortbouwen zonder de oorspronkelijke, kostbare training te moeten herhalen.
Voorbeelden uit de praktijk
Meta bracht in 2023 zijn LLaMA-taalmodellen uit, aanvankelijk beperkt beschikbaar voor onderzoekers en later breder. Wat onderzoekers en ontwikkelaars in feite ontvingen, waren de getrainde checkpoints: de opgeslagen parameters waarmee zij het model direct konden gebruiken of verder konden fine-tunen, zonder de oorspronkelijke, zeer dure trainingsrun zelf te hoeven uitvoeren.
Stability AI publiceerde in 2022 Stable Diffusion, een model dat afbeeldingen genereert op basis van tekst. Omdat het onderliggende checkpoint openbaar werd gemaakt, ontstond een levendige gemeenschap die eigen, gespecialiseerde versies (fine-tunes) maakte, bijvoorbeeld gericht op een bepaalde tekenstijl. Al die varianten zijn in essentie nieuwe checkpoints, voortgebouwd op het originele.
Het onderzoeksinitiatief BigScience, een internationale samenwerking van honderden onderzoekers, publiceerde in 2022 het taalmodel BLOOM. Nadrukkelijk als open-sciencedoel werd niet alleen het model zelf, maar ook informatie over het trainingsproces en verschillende checkpoints tijdens de training gedeeld, zodat andere onderzoekers konden bestuderen hoe het model zich gedurende de training ontwikkelde.
Het platform Hugging Face fungeert inmiddels als een centrale bibliotheek waar bedrijven en onderzoekers checkpoints van duizenden modellen delen, van kleine experimentele netwerken tot grote taalmodellen. Dit heeft het hergebruiken en fine-tunen van bestaande modellen sterk vereenvoudigd.
Hoe ver is de techniek?
Checkpointing zelf is geen nieuwe of onbewezen techniek. Het wordt al sinds het begin van de opkomst van deep learning, ergens in de jaren 2010, standaard toegepast en is een basisonderdeel van vrijwel elk trainingsraamwerk. In die zin is de techniek volwassen en betrouwbaar.
Wat wel is veranderd, is de schaal. Naarmate modellen groeiden van miljoenen naar honderden miljarden parameters, groeide ook de omvang van checkpoints navenant. Bij de grootste taalmodellen kan een enkel checkpoint honderden gigabytes tot enkele terabytes beslaan. Het wegschrijven van zo'n bestand kost tijd, en tijdens het schrijven kan de training vaak niet gewoon doorgaan, wat rekenkracht verspilt.
Dit heeft geleid tot praktische innovaties, eerder infrastructureel dan fundamenteel van aard. Asynchroon checkpointen laat de training doorgaan terwijl het checkpoint op de achtergrond wordt weggeschreven. Sharding verdeelt een groot checkpoint over meerdere bestanden of machines, zodat het sneller parallel kan worden weggeschreven en ingeladen. Quantisatie, het opslaan van getallen met minder precisie (bijvoorbeeld 8 bits in plaats van 32 bits per getal), verkleint checkpoints aanzienlijk, met een kleine impact op nauwkeurigheid.
Een reëel obstakel blijft de opslag- en netwerkinfrastructuur: bij trainingsruns met duizenden chips die frequent checkpoints van terabytes moeten wegschrijven en soms ook delen over datacenters, is dit een serieuze technische en financiële last. Onderzoek naar efficiëntere checkpointmethoden loopt door, maar het is vooral optimalisatie van bestaande aanpak, geen doorbraak die op korte termijn te verwachten valt.
Wie werken eraan?
De belangrijkste softwareraamwerken voor het trainen van AI-modellen, en daarmee ook voor checkpointing, komen van een beperkt aantal grote spelers. PyTorch, ontwikkeld door Meta AI, en TensorFlow, ontwikkeld door Google, zijn de twee dominante raamwerken. Google gebruikt daarnaast intern ook JAX, een raamwerk dat specifiek is toegesneden op het trainen van zeer grote modellen op gespecialiseerde chips.
Grote AI-bedrijven zoals OpenAI, Google DeepMind, Microsoft, Meta AI, Stability AI en het Franse Mistral AI trainen dagelijks modellen waarbij checkpointing een centrale rol speelt in hun infrastructuur, al is niet elk bedrijf even open over de gebruikte checkpoints zelf.
Het platform Hugging Face speelt een aparte rol als tussenpartij: geen ontwikkelaar van modellen zelf, maar een plek waar checkpoints van talloze andere partijen worden gehost, gedeeld en doorzoekbaar gemaakt. Daarnaast dragen academische samenwerkingsverbanden zoals BigScience bij aan meer openheid rond het delen van checkpoints en het bijbehorende onderzoeksproces.