Gradient checkpointing: hoe bespaart het GPU-geheugen?
Automatische vertaling
Dit artikel is automatisch vertaald vanuit de oorspronkelijke Engelse versie.
Gradient checkpointing bewaart minder tussenliggende activations tijdens de forward pass en berekent ontbrekende waarden opnieuw tijdens de backward pass. Gebruik het wanneer activations ervoor zorgen dat de trainingsbatch of sequentielengte niet in het GPU-geheugen past en de extra berekeningen acceptabel zijn.
Het vermindert het geheugen voor opgeslagen activations. Het verwijdert de model weights en optimizerstatus niet, en ook niet alle tijdelijke buffers.
Wat wordt opnieuw berekend?
De backward pass heeft tussenresultaten van de forward pass nodig om gradiënten te berekenen. Zonder checkpointing bewaart het framework de benodigde resultaten. Met checkpointing slaat het geselecteerde invoer of tussenresultaten op en voert het de betreffende bewerkingen opnieuw uit wanneer hun waarden nodig zijn.
Een framework kan bijvoorbeeld de invoer van een groep transformerlagen opslaan en hun interne activations verwijderen. Tijdens de backward pass voert het die groep opnieuw uit vanaf de opgeslagen invoer. De opgeslagen invoer is een activation checkpoint. Dat verschilt van een model checkpoint dat voor herstel naar schijf wordt geschreven.
De oorspronkelijke analyse van Chen et al. toont een strategie met gelijkmatige segmenten die O(sqrt(n)) geheugen voor activations gebruikt voor een netwerk met n lagen en ongeveer één extra forward pass per trainingsstap vereist. Echte transformers hebben activations van verschillende groottes, verschillende attention-implementaties en frameworkbeleid. Meet het resulterende geheugenverbruik en de tijd in plaats van de asymptotische schatting over te nemen als voorspelling voor een deployment.
Schakel de integratie in die je stack verwacht
Voor models die Hugging Face Trainer ondersteunt, schakelt gradient_checkpointing=True het algemene mechanisme in. Controleer bij FSDP in plaats daarvan de integratie van het framework voor activation checkpointing. De huidige documentatie van Transformers Trainer raadt activation_checkpointing aan voor FSDP, omdat generieke gradient checkpointing een extra all-gather kan veroorzaken tijdens de backward pass.
Vergelijk dezelfde trainingsconfiguratie met en zonder checkpointing. Noteer de hoogste GPU-geheugentoewijzing, het aantal voorbeelden of tokens per seconde en of de beoogde sequentielengte past. Controleer ook of de gradiënten geldig blijven wanneer je model aangepaste bewerkingen of adaptertraining gebruikt.
Als de optimizerstatus het grootste deel van het geheugen gebruikt, overweeg dan het verdelen van die status over apparaten of parameter-efficiënte afstemming; het opnieuw berekenen van activations verwijdert die geheugentoewijzing niet.
Het gedeelte over gradient checkpointing in de LLM Engineering Guide legt de berekening per segment uit.