Precision autotuning for linear solvers via contextual bandit-based RL

Titel: Präzises Autotuning für Lineare Löser mittels kontextbasiertem Bandit-basiertem Reinforcement Learning

1. Problemstellung

In der wissenschaftlichen Rechenleistung und beim maschinellen Lernen führt der Einsatz von Floating-Point-Arithmetik zu einem erheblichen Energieverbrauch und Kommunikationsaufwand. Moderne Hardware (z. B. NVIDIA A100, H100) bietet massive Beschleunigungen bei reduzierter Genauigkeit (z. B. FP16, BF16, FP8), was bis zu einem 64-fachen Speedup gegenüber FP64 ermöglicht.
Das zentrale Problem besteht jedoch darin, dass die Verwendung von niedriger Genauigkeit zu Rundungsfehlern und numerischer Instabilität führen kann, insbesondere bei schlecht konditionierten linearen Systemen ($Ax=b$).

• Herausforderung: Bestehende Methoden zur gemischten Genauigkeit (Mixed-Precision) basieren oft auf theoretischen Analysen oder manuellen Heuristiken, die für spezifische Problemklassen gelten, aber nicht generalisierbar sind.
• Limitierung bestehender Tools: Automatisierte Tuning-Tools (wie Precimonious oder HiFPTuner) erfordern oft viele Trial-and-Error-Läufe, skalieren schlecht bei großen Programmen und können ihre Wirksamkeit auf ungesehene Datensätze nicht garantieren.
• Ziel: Entwicklung eines Frameworks, das automatisch und adaptiv die optimale Genauigkeit für einzelne Rechenschritte eines Algorithmus wählt, um einen Kompromiss zwischen Recheneffizienz und numerischer Genauigkeit zu finden, ohne auf umfangreiche manuelle Konfiguration oder wiederholtes Training für jeden neuen Datensatz angewiesen zu sein.

2. Methodik

Die Autoren schlagen ein Reinforcement Learning (RL)-Framework vor, das das Problem der Genauigkeitsauswahl als kontextbasiertes Bandit-Problem (Contextual Bandit) formuliert.

• Formulierung als Kontext-Bandit:

  • Kontext (State $s$): Ein Merkmalsvektor, der die Eigenschaften des linearen Systems beschreibt. Im Fall von Iterativen Verfeinerung (Iterative Refinement) werden der logarithmierte Konditionszahl ($\log_{10}(\kappa(A))$) und die Matrixnorm ($\log_{10}(\|A\|_\infty$)) verwendet.
  • Aktion (Action $a$): Eine Konfiguration der Genauigkeiten für $k$ Rechenschritte des Algorithmus (z. B. LU-Faktorisierung, Residuenberechnung, GMRES-Lösung, Update). Der Aktionsraum ist das kartesische Produkt der verfügbaren Genauigkeitsstufen.
  • Belohnung (Reward $r$): Eine skalare Funktion, die Genauigkeit (Vorwärts- und Rückwärtsfehler), Konvergenzverhalten (Anzahl der Iterationen) und Rechenaufwand (Gewichtung der Genauigkeitsstufen) balanciert. Ein Strafterm für zu viele Iterationen wird eingeführt, um unnötige Rechenzeit zu vermeiden.

• Lernalgorithmus:

  • Es wird eine tabellarische Q-Tabelle verwendet, um den erwarteten Reward für diskretisierte Zustände und Aktionen zu schätzen.
  • Der kontinuierliche Zustandsraum wird durch Binning (Diskretisierung) in einen endlichen Raum überführt, um die Komplexität zu beherrschen.
  • Die Aktionsauswahl erfolgt mittels einer $\epsilon$-greedy-Strategie mit abklingendem $\epsilon$, um Exploration (Testen neuer Konfigurationen) und Exploitation (Nutzung der besten bekannten Konfiguration) zu balancieren.
  • Reduktion des Aktionsraums: Um die exponentielle Explosion des Aktionsraums zu vermeiden, wird eine Ordnungsbeschränkung eingeführt ($u_1 \le u_2 \le \dots \le u_k$), was die Anzahl der gültigen Kombinationen drastisch reduziert (z. B. von 256 auf 35).

• Anwendungsfall (Case Study):
  Das Framework wird auf GMRES-basierte Iterative Verfeinerung (GMRES-IR) angewendet. Der Agent wählt für vier Schritte (Faktorisierung, Residuum, GMRES-Inner-Solver, Update) die Genauigkeit aus dem Set $\{BF16, TF32, FP32, FP64\}$.

3. Wichtige Beiträge

1. Erste RL-basierte Lösung mit Generalisierung: Dies ist laut Autoren die erste Arbeit, die RL zur automatischen Genauigkeitsauswahl (Precision Autotuning) in numerischen Algorithmen einsetzt und die Generalisierungsfähigkeit auf ungesehene Datensätze (Out-of-Sample) verifiziert.
2. Kontext-Bandit-Formulierung: Die formale Definition des Problems als kontextbasiertes Bandit, das die Unabhängigkeit von Probleminstanzen nutzt und keine langfristigen Folgen (wie bei vollem MDP) betrachtet, sondern den sofortigen Reward optimiert.
3. Theoretische Absicherung: Es wird eine theoretische Fehlergrenze für die Diskretisierung des Zustandsraums hergeleitet, die zeigt, dass der Performance-Verlust durch die Diskretisierung durch die Granularität der Bins kontrolliert wird (unter der Annahme einer Lipschitz-Stetigkeit des erwarteten Rewards).
4. Adaptive Strategie: Das Framework lernt, dass bei gut konditionierten Systemen niedrigere Genauigkeiten (BF16/TF32) sicher sind, während es bei schlecht konditionierten Systemen automatisch auf FP64 zurückfällt, um Stabilität zu gewährleisten.

4. Ergebnisse

Die Experimente wurden mit synthetischen dichten und dünnbesetzten (sparse) linearen Systemen durchgeführt.

• Dichte Systeme (Dense Systems):

  • Der RL-Agent lernte erfolgreich, die Genauigkeit basierend auf der Konditionszahl anzupassen.
  • Bei gut konditionierten Systemen ($\kappa \approx 10^2 - 10^3$) nutzte das aggressive Modell (W2) erfolgreich BF16 und TF32, was zu einer Reduktion der Rechenkosten führte, während die Genauigkeit erhalten blieb.
  • Bei schlecht konditionierten Systemen ($\kappa > 10^6$) wechselte der Agent automatisch zu FP64, um Konvergenz und Stabilität zu sichern.
  • Die Erfolgsrate (Success Rate) lag bei konservativen Einstellungen bei 100 % und war mit der FP64-Baseline vergleichbar.

• Dünnbesetzte Systeme (Sparse Systems):

  • Diese Datensätze waren extrem schlecht konditioniert ($\kappa \approx 10^8 - 10^9$).
  • Der RL-Agent zeigte hier ein hochkonservatives Verhalten: Er wählte in fast allen Fällen FP64 für alle Schritte, da jede Reduktion der Genauigkeit zu Divergenz oder Stagnation geführt hätte.
  • Dies beweist die Robustheit des Frameworks: Es priorisiert numerische Stabilität über Recheneffizienz, wenn die Problemstellung dies erfordert.

• Ablationsstudie (Strafterm):

  • Eine Studie ohne den Strafterm für Iterationen zeigte, dass der Agent dann mehr Iterationen in Kauf nahm, um niedrigere Genauigkeiten zu nutzen. Dies unterstreicht die Notwendigkeit einer sorgfältigen Modellierung der Reward-Funktion, um einen echten Effizienzgewinn zu erzielen.

5. Bedeutung und Ausblick

• Praktische Relevanz: Das Framework ermöglicht es, numerische Algorithmen ohne tiefes Expertenwissen über die spezifischen Fehlergrenzen für neue, ungesehene Probleme effizient zu konfigurieren. Es reduziert die Abhängigkeit von manuellen Heuristiken.
• Black-Box-Ansatz: Es kann saliente Merkmale identifizieren, die für die Genauigkeitswahl entscheidend sind, auch wenn diese nicht a priori bekannt sind.
• Zukunftsaussichten: Die Autoren planen, Deep RL für kontinuierliche Zustandsräume zu erforschen, weitere Matrixmerkmale (z. B. Fill-in-Ratio, Pivot-Wachstum) einzubeziehen und das Framework auf andere numerische Algorithmen jenseits von linearen Lösern zu erweitern.

Fazit: Die Arbeit demonstriert erfolgreich, dass kontextbasiertes Reinforcement Learning eine leistungsfähige Methode ist, um gemischte Genauigkeitsstrategien für lineare Löser automatisch zu optimieren, wobei ein optimaler Kompromiss zwischen Geschwindigkeit und Genauigkeit erreicht wird, der sich dynamisch an die Eigenschaften des Eingabedatensatzes anpasst.