Inflection AI hat in enger Zusammenarbeit mit Intel die Inferenz-Stack des eigenen Large Language Models auf Intels Gaudi-Accelerator (HPU) portiert. Das technische Team teilt Einblicke in die Herausforderungen des Wechsels von NVIDIA zu einer alternativen Beschleunigerarchitektur und wie sie hohe Performance-Standards erreicht haben.

Die Herausforderung: Vom NVIDIA-Ökosystem zum Gaudi-Accelerator

Das Mapping von KI-Architekturen auf die zugrundeliegende Hardware ist entscheidend für effiziente und skalierbare KI-Workloads. Die Anzahl der PyTorch-Operations ist in den letzten zehn Jahren auf über 2000 angewachsen. SynapseAI, die Software-Backend für Gaudi, unterstützt nur eine Teilmenge, die über die anderen Operations generalisiert.

Zwar war ein Großteil von Inflections Modellarchitektur nativ von Habonas Software-Stack unterstützt, Operations wie pythonisches (numpy-basiertes) Tensor-Slicing und weniger gängige Funktionen wie torch.triu_indices fehlten jedoch. Bei großen Tensoren führten einige dieser Operations zu schwer diagnosebaren Segfaults, während nicht unterstützte Operations zu implizitem Fallback auf die CPU führten.

Performance-Optimierung: Bis zu 15x schneller

Der Fallback auf die CPU führte zu enormen Overhead: Tensor-Übertragung zur CPU, Ausführung und Rückübertragung zur HPU verursachte Latenz, die die native HPU-Berechnungszeiten bei Weitem übertraf. Inflection löste beide Probleme durch Neuschreiben der nicht unterstützten Operationen in performantere Äquivalente – mit einem 15-fachen Speedup und ohne Segfaults.

Execution Frameworks: Eager vs. Lazy

Anfangs nutzte das Team den Eager-Execution-Modus (PyTorch-Standard), der Out-of-the-Box funktionierte, aber höhere Latenz als NVIDIA aufwies. Gaudi-Hardware ist jedoch für Graph-Execution optimiert – das sequenzielle Ausführen von Operations einzeln verursacht Python-Overhead und reduzierte Fusionsmöglichkeiten.

Der Wechsel zum Lazy-Modus (Gaudis Standard) brachte eigene Probleme mit sich: Dynamische Operationen wie datenabhängige Verzweigungen oder variable Tensor-Shape-Breiten brechen den Lazy-Graphen. Die naive Lazy-Implementierung war doppelt so langsam wie Eager. Nach iterativem Entfernen aller dynamischen Operationen und Integration von HPU-Graphen (vergleichbar mit CUDA-Graphen) erreichten sie einen 4-fachen Speedup, der die NVIDIA-Performance einholte.

Profilierung und Hardware-Auslastung

Ähnlich wie die H100, die TMA-basierte GEMMs auf den Tensor Cores mit exponentiellen Berechnungen der speziellen Multi-Function-Unit überlappt, nutzt Gaudi gleichzeitig die MME für GEMM-Operationen und separate Tensor Processing Cores (TPC) für elementweise Operationen. Inflection arbeitete mit dem Intel Habana-Team zusammen, um den Attention-Mechanismus in zwei Teile zu splitten – während die MME die schweren Matrixmultiplikationen bearbeitet, beginnt die TPC parallel mit der nächsten Layer.

Ausblick: Hardware-agnostische Optimierung

Der Erfolg unterstreicht ein breiteres Prinzip: Hardware-Design-Patterns wiederholen sich über verschiedene Beschleuniger hinweg. Die gewonnenen Erkenntnisse aus der Optimierung einer Plattform leiten und beschleunigen Bemühungen auf jeder neuen Architektur. Die entwickelten Methoden skalieren mit Modellgrößen von 1B bis über 100B+ Parameter.

Quelle: Inflection AI Blog