Het artikel bespreekt een optimalisatie voor PyTorch's activation memory budget solver. Bij het bouwen van een berekeningsgrafiek moet PyTorch bepalen welke operaties bewaard worden in het geheugen en welke opnieuw berekend moeten worden om de runtime te optimaliseren binnen een beperkt RAM-budget. Dit is in essentie een 0/1 Knapsack-probleem.
De huidige standaard, dpknapsack, maakt gebruik van dynamisch programmeren (DP) met een 2D-tabel, wat bij grote modellen tot crashes kan leiden door een te hoog piekgeheugengebruik. De nieuwe solver, dpknapsackslidinghirschberg, lost dit op door twee technieken te combineren:
- Sliding window: Alleen de vorige en huidige rij van de DP-tabel worden bijgehouden.
- Hirschberg-truc: Een divide and conquer-benadering die backtracking overbodig maakt.
Resultaten:
- Het piekgeheugengebruik daalt met een factor 20x.
- De solver kan veel grotere problemen aan (tot 2000 items tegenover 100 bij de oude versie).
- Er is een runtime-snelheidswinst van ongeveer 37%.
Het artikel concludeert dat gebruikers die een exacte oplossing willen zonder externe afhankelijkheden zoals SciPy (nodig voor ilpknapsack), kunnen overstappen op deze nieuwe solver via de torch.functorch.config.
Gebruik 20x minder piek-RAM met dpknapsacksliding_hirschberg, een nieuwe activation memory budget solver voor PyTorch
Wat is dit?
Wanneer PyTorch een berekeningsgrafiek (computation graph) van je ML-model bouwt, probeert het verschillende methoden om deze snel en geheugenefficiënt te maken. Eén techniek is om een gezamenlijke forward- en backward-berekeningsgrafiek te nemen en te bepalen welke operaties bewaard moeten worden (opgeslagen in het geheugen en hergebruikt) en welke operaties opnieuw berekend moeten worden.
Elke operatie kost tijd om uit te voeren en neemt geheugen in beslag als deze wordt opgeslagen, terwijl we slechts een beperkte hoeveelheid geheugen tot onze beschikking hebben.
Het 0/1 Knapsack-probleem
Stel je voor dat je een dief bent die zojuist is ingebroken in een pand met veel waardevolle spullen. Je hebt een rugzak (knapsack) die maximaal 10 kilogram kan dragen. Je wilt de waarde van de spullen die je steelt maximaliseren. Ideaal gezien zou je alles meenemen, maar als je meer dan 10 kg in je rugzak stopt, scheurt deze en blijf je met niets achter.
Alternatief kun je heuristieken gebruiken om items te kiezen (zoals alleen kleine items of alleen groene items), maar dan is je oplossing mogelijk niet optimaal; je steelt wel waarde, maar wellicht minder dan wat je met een optimale keuze had kunnen behalen.
Hier komen 0/1 knapsack-solvers kijken. Dit zijn algoritmen die de optimale keuze van items berekenen binnen een gegeven capaciteit (maximaal gewicht van de spullen in de rugzak) en een lijst met items, waarbij elk item een waarde en een kostprijs (gewicht) heeft.
In PyTorch doen we in feite hetzelfde: binnen een RAM-budget (de capaciteit van de rugzak) kiezen we welke operaties (de items) bewaard moeten worden in het geheugen om de tijdwinst tijdens de runtime (de waarde) te maximaliseren.
dp_knapsack
Momenteel is de standaardimplementatie van de knapsack-solver in PyTorch dp_knapsack, een benadering op basis van dynamisch programmeren (DP) om het 0/1 knapsack-probleem op te lossen. Het algoritme is eenvoudig zodra je het begrijpt, zeker als je al ervaring hebt met dynamisch programmeren.
Stel dat er een lijst is met 4 items, waarbij i de index is, v de waarde en w het gewicht. De capaciteit is 6. In een DP-tabel worden daarom 6 kolommen gebruikt (elk een deel van de totale capaciteit) en 4 rijen (één per item). In de praktijk zijn dit er vaak 7 en 5 vanwege 'padding' met nullen voor het gemak van de berekeningen (zodat elke rij en kolom een ouder heeft en de index-aritmetiek eenvoudiger wordt).
Het bouwen van een DP-tabel
Het DP 0/1 Knapsack-algoritme werkt als volgt:
- Neem een item.
- Ga in de DP-tabel door de bijbehorende rij, kolom voor kolom.
- Vergelijk de waarde van het item met de waarde in de huidige cel van de DP-tabel (rij, kolom; waarbij de huidige kolomindex het huidige maximale gewicht/capaciteit van de rugzak is).
- Als het gewicht van het item groter is dan de huidige kolomindex, ga dan door naar de volgende kolom.
- Als het gewicht van het item gelijk is aan een kolomindex en de waarde van het item hoger is dan de huidige cel, vervang dan de cel door de waarde van het item.
- Als het gewicht van het item kleiner is dan de kolomindex en het past in de resterende capaciteit, voeg het dan toe aan de waarde van de cel (vrije capaciteit is de huidige capaciteit minus de capaciteit van de reeds gekozen items).
- Herhaal dit totdat alle items zijn doorlopen.
Op dit punt weet je wat de maximale waarde is die je kunt behalen binnen het gewichtsbudget. Als het doel enkel het meten van de maximale waarde was, zouden we hier stoppen. Echter, we willen meer weten.
Backtracking
We moeten weten welke items we moeten nemen om deze maximale waarde te bereiken. In het geval van PyTorch betekent dit: welke operaties we in het geheugen moeten opslaan en welke operaties we opnieuw moeten berekenen.
Om te achterhalen welke items bijdragen aan de maximale waarde (welke items zijn gekozen tijdens het bouwen van de DP-tabel), moeten we terugwaarts door de DP-tabel navigeren:
- We beginnen bij de laatste rij en laatste kolom en bewegen omhoog.
- We vergelijken een cel met de cel erboven (dezelfde kolom, vorige rij).
- Als ze verschillen, draagt het huidige item bij aan de maximale waarde (omdat elke rij correspondeert met een enkel item).
- Als ze gelijk zijn, wordt dit item niet in het geheugen opgeslagen en is het opnieuw berekenbaar.
- Als het item is gekozen, bewegen we één rij omhoog en verschuiven we de positie van de huidige cel naar links met het gewicht van het gekozen item.
- Als het item niet is gekozen, bewegen we één rij omhoog en blijven we op dezelfde positie.
- Dit proces wordt herhaald tot we de bovenkant van de DP-tabel bereiken.
Optimalisatie
Deze solver werkt correct en biedt een exacte oplossing. Voor de meeste gevallen is dit voldoende, maar er zijn verbeteringen mogelijk op het gebied van snelheid en geheugengebruik.
Het grootste probleem is dat er voor het resultaat een volledige 2D-tabel moet worden gealloceerd met de vorm (aantalitems, maxgewicht), wat niet efficiënt is. Horace He (Chillee), die een groot deel van de memory planning in PyTorch heeft gebouwd, liet hiervoor de volgende notitie achter in de code:
# TODO(chilli): I think if needed, this memory can be optimized with sliding window trick + Hirschberg trick
Hij stelt twee optimalisaties voor: de sliding window-truc en de Hirschberg-truc.
Sliding window
De window-truc houdt in dat we in plaats van een volledige DP-tabel te bouwen, over een tabel "glijden" en slechts de vorige rij en de huidige rij gebruiken. Wanneer we naar de volgende rij gaan, wordt de huidige rij de vorige rij.
Deze verbetering alleen zou voldoende zijn als we alleen de maximale waarde wilden berekenen. Maar omdat we moeten weten welke items bijdragen aan die waarde, is de sliding window-truc alleen niet genoeg. Daarom wordt deze gecombineerd met de Hirschberg-truc.
De Hirschberg-truc
Dit is een divide and conquer-benadering, vergelijkbaar met quicksort, waarbij het probleem wordt opgesplitst in kleinere deelproblemen die vervolgens worden opgelost. Het belangrijkste voordeel van de Hirschberg-truc is dat het backtracking overbodig maakt.
Vanuit implementatieperspectief wordt dit gedaan via een recursie, maar zonder nieuwe frames op de stack te plaatsen. Er wordt een Python-lijst gebruikt als expliciete stack (last in, first out), die een bereik van item-indexen en de huidige beschikbare capaciteit opslaat.
Het proces verloopt als volgt:
- Haal een item van de stack.
- Als het een enkel item is (bereik index einde - start = 1), controleer dan of het gewicht van het item in de beschikbare capaciteit past. Zo ja, voeg het toe aan de lijst met bewaarde items; zo nee, voeg het toe aan de lijst met opnieuw berekenbare items. Dit is de grootste kracht: de lijsten worden tijdens het algoritme gebouwd, waardoor backtracking achteraf niet nodig is.
- Als het bereik van items groter is dan 1, splitst het algoritme de items in twee delen en berekent voor elk deel een DP-tabel. Hierbij wordt de sliding window-truc toegepast om alleen de laatste DP-rij te berekenen met behulp van een huidige en vorige rij.
- Zodra beide DP-rijen (links en rechts) beschikbaar zijn, worden ze bij elkaar opgeteld. De positie van de hoogste waarde wordt het nieuwe splitsingspunt voor het geheugen.
- Het geheugen wordt gesplitst op basis van deze positie en het algoritme wordt opnieuw uitgevoerd voor beide delen van het geheugen en de items.
- Dit wordt herhaald tot het proces is voltooid.
(Opmerking: het kan gebeuren dat een van de gesplitste delen leeg is als het beste splitsingspunt het eerste of laatste item is. Dit is toegestaan; dit stack-frame wordt dan simpelweg overgeslagen.)
Is dpknapsacksliding_hirschberg nu dé solver?
Nee. Met deze twee optimalisaties kunnen we de grootte van de DP-tabel verlagen van (aantalitems, maxgewicht) naar (2, max_gewicht). Dit vermindert het piekgeheugengebruik met een factor 20x.
Op een machine met 64 GB RAM crasht dpknapsack bij problemen met 100 items, terwijl dpknapsackslidinghirschberg problemen met 2000 items kan afhandelen binnen hetzelfde budget. Er is ook een winst in runtime-snelheid, ongeveer 37% (hoewel dit niet strikt rigoureus is gemeten).
Dat gezegd hebbende: als je bereid bent SciPy te gebruiken, is ilp_knapsack waarschijnlijk de betere keuze. Deze is VEEL sneller dan elke bestaande DP-knapsack-implementatie in PyTorch. Het is iets langzamer dan een greedy-algoritme, maar biedt, net als de DP-solvers, een exacte oplossing.
Hoe kies je de juiste solver?
- Snelheid is belangrijker dan correctheid: gebruik
greedy_knapsack.
- Je wilt een exacte oplossing en SciPy toevoegen aan je project is geen probleem: gebruik
ilp_knapsack.
- Je installeert PyTorch via een package manager en wilt geen extra configuratie: doe niets; PyTorch gebruikt standaard
dp_knapsack.
- Je wilt een exacte oplossing zonder SciPy, maar met een veel lager geheugengebruik dan de standaard: gebruik
dpknapsacksliding_hirschberg.
Hoe gebruik je dpknapsacksliding_hirschberg?
import torch
import torch._functorch.config as fconfig
fconfig.activation_memory_budget_solver = fconfig.dp_knapsack_sliding_hirschberg
Gebruik 20x minder piek-RAM met dpknapsacksliding_hirschberg, een nieuwe activation memory budget solver voor PyTorch
Wat is dit?
Wanneer PyTorch een berekeningsgrafiek (computation graph) van je ML-model bouwt, probeert het verschillende methoden om deze snel en geheugenefficiënt te maken. Eén techniek is om een gezamenlijke forward- en backward-berekeningsgrafiek te nemen en te bepalen welke operaties bewaard moeten worden (opgeslagen in het geheugen en hergebruikt) en welke operaties opnieuw berekend moeten worden.
Elke operatie kost tijd om uit te voeren en neemt geheugen in beslag als deze wordt opgeslagen, terwijl we slechts een beperkte hoeveelheid geheugen tot onze beschikking hebben.
Het 0/1 Knapsack-probleem
Stel je voor dat je een dief bent die zojuist is ingebroken in een pand met veel waardevolle spullen. Je hebt een rugzak (knapsack) die maximaal 10 kilogram kan dragen. Je wilt de waarde van de spullen die je steelt maximaliseren. Ideaal gezien zou je alles meenemen, maar als je meer dan 10 kg in je rugzak stopt, scheurt deze en blijf je met niets achter.
Alternatief kun je heuristieken gebruiken om items te kiezen (zoals alleen kleine items of alleen groene items), maar dan is je oplossing mogelijk niet optimaal; je steelt wel waarde, maar wellicht minder dan wat je met een optimale keuze had kunnen behalen.
Hier komen 0/1 knapsack-solvers kijken. Dit zijn algoritmen die de optimale keuze van items berekenen binnen een gegeven capaciteit (maximaal gewicht van de spullen in de rugzak) en een lijst met items, waarbij elk item een waarde en een kostprijs (gewicht) heeft.
In PyTorch doen we in feite hetzelfde: binnen een RAM-budget (de capaciteit van de rugzak) kiezen we welke operaties (de items) bewaard moeten worden in het geheugen om de tijdwinst tijdens de runtime (de waarde) te maximaliseren.
dp_knapsack
Momenteel is de standaardimplementatie van de knapsack-solver in PyTorch dp_knapsack, een benadering op basis van dynamisch programmeren (DP) om het 0/1 knapsack-probleem op te lossen. Het algoritme is eenvoudig zodra je het begrijpt, zeker als je al ervaring hebt met dynamisch programmeren.
Stel dat er een lijst is met 4 items, waarbij i de index is, v de waarde en w het gewicht. De capaciteit is 6. In een DP-tabel worden daarom 6 kolommen gebruikt (elk een deel van de totale capaciteit) en 4 rijen (één per item). In de praktijk zijn dit er vaak 7 en 5 vanwege 'padding' met nullen voor het gemak van de berekeningen (zodat elke rij en kolom een ouder heeft en de index-aritmetiek eenvoudiger wordt).
Het bouwen van een DP-tabel
Het DP 0/1 Knapsack-algoritme werkt als volgt:
- Neem een item.
- Ga in de DP-tabel door de bijbehorende rij, kolom voor kolom.
- Vergelijk de waarde van het item met de waarde in de huidige cel van de DP-tabel (rij, kolom; waarbij de huidige kolomindex het huidige maximale gewicht/capaciteit van de rugzak is).
- Als het gewicht van het item groter is dan de huidige kolomindex, ga dan door naar de volgende kolom.
- Als het gewicht van het item gelijk is aan een kolomindex en de waarde van het item hoger is dan de huidige cel, vervang dan de cel door de waarde van het item.
- Als het gewicht van het item kleiner is dan de kolomindex en het past in de resterende capaciteit, voeg het dan toe aan de waarde van de cel (vrije capaciteit is de huidige capaciteit minus de capaciteit van de reeds gekozen items).
- Herhaal dit totdat alle items zijn doorlopen.
Op dit punt weet je wat de maximale waarde is die je kunt behalen binnen het gewichtsbudget. Als het doel enkel het meten van de maximale waarde was, zouden we hier stoppen. Echter, we willen meer weten.
Backtracking
We moeten weten welke items we moeten nemen om deze maximale waarde te bereiken. In het geval van PyTorch betekent dit: welke operaties we in het geheugen moeten opslaan en welke operaties we opnieuw moeten berekenen.
Om te achterhalen welke items bijdragen aan de maximale waarde (welke items zijn gekozen tijdens het bouwen van de DP-tabel), moeten we terugwaarts door de DP-tabel navigeren:
- We beginnen bij de laatste rij en laatste kolom en bewegen omhoog.
- We vergelijken een cel met de cel erboven (dezelfde kolom, vorige rij).
- Als ze verschillen, draagt het huidige item bij aan de maximale waarde (omdat elke rij correspondeert met een enkel item).
- Als ze gelijk zijn, wordt dit item niet in het geheugen opgeslagen en is het opnieuw berekenbaar.
- Als het item is gekozen, bewegen we één rij omhoog en verschuiven we de positie van de huidige cel naar links met het gewicht van het gekozen item.
- Als het item niet is gekozen, bewegen we één rij omhoog en blijven we op dezelfde positie.
- Dit proces wordt herhaald tot we de bovenkant van de DP-tabel bereiken.
Optimalisatie
Deze solver werkt correct en biedt een exacte oplossing. Voor de meeste gevallen is dit voldoende, maar er zijn verbeteringen mogelijk op het gebied van snelheid en geheugengebruik.
Het grootste probleem is dat er voor het resultaat een volledige 2D-tabel moet worden gealloceerd met de vorm (aantalitems, maxgewicht), wat niet efficiënt is. Horace He (Chillee), die een groot deel van de memory planning in PyTorch heeft gebouwd, liet hiervoor de volgende notitie achter in de code:
# TODO(chilli): I think if needed, this memory can be optimized with sliding window trick + Hirschberg trick
Hij stelt twee optimalisaties voor: de sliding window-truc en de Hirschberg-truc.
Sliding window
De window-truc houdt in dat we in plaats van een volledige DP-tabel te bouwen, over een tabel "glijden" en slechts de vorige rij en de huidige rij gebruiken. Wanneer we naar de volgende rij gaan, wordt de huidige rij de vorige rij.
Deze verbetering alleen zou voldoende zijn als we alleen de maximale waarde wilden berekenen. Maar omdat we moeten weten welke items bijdragen aan die waarde, is de sliding window-truc alleen niet genoeg. Daarom wordt deze gecombineerd met de Hirschberg-truc.
De Hirschberg-truc
Dit is een divide and conquer-benadering, vergelijkbaar met quicksort, waarbij het probleem wordt opgesplitst in kleinere deelproblemen die vervolgens worden opgelost. Het belangrijkste voordeel van de Hirschberg-truc is dat het backtracking overbodig maakt.
Vanuit implementatieperspectief wordt dit gedaan via een recursie, maar zonder nieuwe frames op de stack te plaatsen. Er wordt een Python-lijst gebruikt als expliciete stack (last in, first out), die een bereik van item-indexen en de huidige beschikbare capaciteit opslaat.
Het proces verloopt als volgt:
- Haal een item van de stack.
- Als het een enkel item is (bereik index einde - start = 1), controleer dan of het gewicht van het item in de beschikbare capaciteit past. Zo ja, voeg het toe aan de lijst met bewaarde items; zo nee, voeg het toe aan de lijst met opnieuw berekenbare items. Dit is de grootste kracht: de lijsten worden tijdens het algoritme gebouwd, waardoor backtracking achteraf niet nodig is.
- Als het bereik van items groter is dan 1, splitst het algoritme de items in twee delen en berekent voor elk deel een DP-tabel. Hierbij wordt de sliding window-truc toegepast om alleen de laatste DP-rij te berekenen met behulp van een huidige en vorige rij.
- Zodra beide DP-rijen (links en rechts) beschikbaar zijn, worden ze bij elkaar opgeteld. De positie van de hoogste waarde wordt het nieuwe splitsingspunt voor het geheugen.
- Het geheugen wordt gesplitst op basis van deze positie en het algoritme wordt opnieuw uitgevoerd voor beide delen van het geheugen en de items.
- Dit wordt herhaald tot het proces is voltooid.
(Opmerking: het kan gebeuren dat een van de gesplitste delen leeg is als het beste splitsingspunt het eerste of laatste item is. Dit is toegestaan; dit stack-frame wordt dan simpelweg overgeslagen.)
Is dpknapsacksliding_hirschberg nu dé solver?
Nee. Met deze twee optimalisaties kunnen we de grootte van de DP-tabel verlagen van (aantalitems, maxgewicht) naar (2, max_gewicht). Dit vermindert het piekgeheugengebruik met een factor 20x.
Op een machine met 64 GB RAM crasht dpknapsack bij problemen met 100 items, terwijl dpknapsackslidinghirschberg problemen met 2000 items kan afhandelen binnen hetzelfde budget. Er is ook een winst in runtime-snelheid, ongeveer 37% (hoewel dit niet strikt rigoureus is gemeten).
Dat gezegd hebbende: als je bereid bent SciPy te gebruiken, is ilp_knapsack waarschijnlijk de betere keuze. Deze is VEEL sneller dan elke bestaande DP-knapsack-implementatie in PyTorch. Het is iets langzamer dan een greedy-algoritme, maar biedt, net als de DP-solvers, een exacte oplossing.
Hoe kies je de juiste solver?
- Snelheid is belangrijker dan correctheid: gebruik
greedy_knapsack.
- Je wilt een exacte oplossing en SciPy toevoegen aan je project is geen probleem: gebruik
ilp_knapsack.
- Je installeert PyTorch via een package manager en wilt geen extra configuratie: doe niets; PyTorch gebruikt standaard
dp_knapsack.
- Je wilt een exacte oplossing zonder SciPy, maar met een veel lager geheugengebruik dan de standaard: gebruik
dpknapsacksliding_hirschberg.
Hoe gebruik je dpknapsacksliding_hirschberg?
import torch
import torch._functorch.config as fconfig
fconfig.activation_memory_budget_solver = fconfig.dp_knapsack_sliding_hirschberg
Groetjes,