Precision autotuning for linear solvers via contextual bandit-based RL

Probleemstelling

In wetenschappelijk rekenen en machine learning wordt floating-point rekenen (zoals FP64) intensief gebruikt, wat leidt tot aanzienlijke energieconsumptie en communicatiekosten. Moderne hardware (zoals NVIDIA A100 en H100) biedt aanzienlijke snelheidswinsten (tot 64x) bij het gebruik van verminderde precisie-formaten (zoals FP16, BF16, FP8). Echter, het gebruik van verminderde precisie introduceert afrondingsfouten en numerieke instabiliteit, wat de convergentie en nauwkeurigheid van lineaire oplosmethoden (zoals het oplossen van $Ax = b$) kan compromitteren.

Bestaande methoden voor mixed-precision computing (combinatie van hoge en lage precisie) hebben twee grote tekortkomingen:

1. Theoretische analyses zijn vaak te restrictief en vereisen vooraf kennis van systeemkenmerken (zoals de conditienummer), die in de praktijk niet altijd beschikbaar of snel te berekenen zijn.
2. Autotuning-tools (zoals Precimonious of HiFPTuner) werken vaak via trial-and-error, wat rekenkundig duur is en niet goed schaalbaar is naar grote programma's of nieuwe, ongezette datasets.

Er is dus behoefte aan een methode die dynamisch en adaptief de juiste precisie kiest voor verschillende stappen in een algoritme, gebaseerd op de eigenschappen van het specifieke probleem, zonder dat er uitgebreide handmatige tuning nodig is.

Methodologie

De auteurs stellen een Reinforcement Learning (RL) framework voor dat het probleem van precisiekeuze formuleert als een Contextual Bandit-probleem. Dit is een vereenvoudigde vorm van RL waarbij de agent beslissingen neemt op basis van de huidige context (toestand) om een directe beloning te maximaliseren, zonder rekening te houden met toekomstige contexten.

Kerncomponenten van het framework:

1. Contextuele Bandit Formulering:

   • Context (Toestand $s$): Een feature-vector die het probleem beschrijft. Voor lineaire systemen zijn dit onder andere het geschatte conditienummer ($\kappa(A)$) en de matrixnorm ($\|A\|_\infty$). Deze continue waarden worden gediskretiseerd in "bins" om de zoekruimte beheersbaar te maken.
   • Actie ($a$): Een configuratie van precisies voor de verschillende rekenstappen van de solver. In het geval van GMRES-IR (Iterative Refinement) worden vier stappen gedefinieerd: LU-factorisatie, residu-berekening, GMRES-oplossing en update van de oplossing. De actie is een tuple $(u_f, u, u_g, u_r)$.
   • Beloning (Reward $r$): Een gewogen som van meerdere doelen:
     • Nauwkeurigheid: Straft fouten in de oplossing (forward/backward error).
     • Kosten: Belonigt het gebruik van lagere precisie (minder bits).
     • Strafterm: Straft het aantal iteraties of falen van convergentie om te voorkomen dat de agent te veel iteraties gebruikt om lage precisie te compenseren.

2. Leeralgoritme:

   • De agent leert een Q-tabel (waarde-functie) die de verwachte beloning schat voor een gegeven gediskretiseerde toestand en een gekozen precisie-configuratie.
   • Er wordt gebruik gemaakt van een $\epsilon$-greedy strategie om exploratie (proberen van nieuwe combinaties) en exploitatie (gebruik van de beste bekende keuze) te balanceren. De $\epsilon$-waarde neemt af tijdens het trainen.
   • Actieruimte-reductie: Om de exponentiële groei van mogelijke precisiecombinaties te beperken, wordt een orde-beperking opgelegd ($u_f \leq u \leq u_g \leq u_r$). Dit reduceert het aantal mogelijke acties aanzienlijk (in de experimenten van 256 naar 35).

3. Toepassing op GMRES-IR:
   Het framework wordt getest op GMRES-based Iterative Refinement. De agent leert welke precisie (BF16, TF32, FP32, FP64) het beste is voor elke stap in het iteratieve proces, afhankelijk van de eigenschappen van de matrix.

Belangrijkste Resultaten

De auteurs hebben experimenten uitgevoerd op zowel dichte als verspreide (sparse) lineaire systemen met variërende conditienummers.

1. Dichte Systemen (Dense Systems):

   • Het RL-agent leert een conditie-afhankelijke strategie. Voor goed geconditioneerde systemen (lage $\kappa$) kiest de agent agressief voor lagere precisie (BF16, TF32), wat leidt tot aanzienlijke besparingen in rekentijd.
   • Voor slecht geconditioneerde systemen (hoge $\kappa$) schakelt de agent automatisch over naar hoge precisie (FP64) om convergentie en nauwkeurigheid te garanderen.
   • Twee beleidsinstellingen werden getest:
     • Conservatief (W1): Behoudt nauwkeurigheid vergelijkbaar met FP64-baselines.
     • Agressief (W2): Kiest vaker voor lage precisie, wat leidt tot iets hogere fouten maar wel binnen acceptabele grenzen, met een duidelijke winst in efficiëntie.

2. Verspreide Systemen (Sparse Systems):

   • Voor zeer slecht geconditioneerde verspreide systemen (conditienummers tot $10^9$) gedraagt de agent zich uiterst conservatief.
   • De agent leert dat elke reductie in precisie hier leidt tot divergentie of stagnatie. Daarom kiest hij bijna uitsluitend voor FP64, ongeacht de beloningsinstelling. Dit toont aan dat het framework de "overlevingsgrens" van het algoritme correct identificeert en niet blindelings optimaliseert voor snelheid ten koste van stabiliteit.

3. Generalisatie:
   Het model, getraind op een beperkte dataset, generaliseert succesvol naar ongezette testdata. Het vermogen om de juiste precisie te infereren op basis van features zonder dat de exacte oplossing bekend is, is een sterk punt.

4. Impact van de Strafterm:
   Een ablatiestudie toonde aan dat het verwijderen van de strafterm voor het aantal iteraties leidt tot een ongewenst gedrag: de agent kiest dan wel lage precisie, maar moet hierdoor veel meer iteraties uitvoeren om dezelfde nauwkeurigheid te bereiken, wat de winst tenietdoet.

Bijdragen

1. Eerste werk met RL voor precisie-autotuning: Dit is, naar de kennis van de auteurs, het eerste werk dat Reinforcement Learning toepast voor precisiekeuze in numerieke algoritmen met verificatie op ongezette datasets.
2. Formulering als Contextual Bandit: Het biedt een generiek raamwerk dat het probleem van precisiekeuze formaliseert als een bandit-probleem, wat efficiënter is dan volledige MDP's (Markov Decision Processes) voor deze specifieke toepassing.
3. Adaptieve Strategie: Het systeem leert dynamisch om te schakelen tussen lage en hoge precisie op basis van de numerieke stabiliteit van het specifieke probleem, zonder dat menselijke experts handmatige regels hoeven op te stellen.
4. Actieruimte-Optimalisatie: De introductie van een orde-beperking voor precisiecombinaties maakt het trainen van het model haalbaar voor complexe algoritmen.

Betekenis en Toekomstperspectief

Deze studie toont aan dat AI-gedreven methoden (specifiek RL) effectief kunnen worden ingezet om de balans te vinden tussen rekenkracht en nauwkeurigheid in wetenschappelijk rekenen. Het framework biedt een oplossing voor het "black-box" probleem van het kiezen van de juiste precisie in mixed-precision algoritmen.

Toekomstig werk omvat:

• Het toepassen van Deep RL om te werken met continue toestandsruimtes in plaats van gediskretiseerde bins.
• Het integreren van meer complexe matrixkenmerken (zoals fill-in ratio, pivot-groei en spectrale eigenschappen) om de besluitvorming te verfijnen.
• Het uitbreiden van het framework naar andere numerieke algoritmen buiten lineaire oplossers.

Samenvattend biedt dit paper een robuust en generaliseerbaar raamwerk dat de weg vrijmaakt voor efficiëntere, energiebesparende numerieke methoden in de toekomstige high-performance computing (HPC) en AI-landschappen.