NVIDIA testet Tensor-Core-FP32 mit PET: Kernel- und Gesamtschritt-Gewinne unterscheiden sich

NVIDIA veröffentlichte am 29. Oktober Versuche zur Genauigkeit und Geschwindigkeit maschinell gelernter interatomarer Potenziale. Unsere technische Ausgabe vom 1. November behält dieses Originaldatum bei. Untersucht wurden Tensor-Core-Verfahren für FP32-ähnliche Matrixmultiplikation in bestimmten Point-Edge-Transformer- oder PET-Lasten. Das ist eine Messung des Anbieters auf bestimmten Modellen und GPUs, kein allgemeiner Nachweis gleicher Beschleunigung für jede wissenschaftliche Anwendung. Werte für Matrixkerne und ganze Schritte der Molekulardynamik müssen getrennt bleiben.

Ein interatomares Potenzial sagt die Energie einer Atomanordnung voraus; daraus abgeleitete Kräfte bestimmen über viele Schritte hinweg Position und Geschwindigkeit. Kleine Rechenfehler können sich in einer Trajektorie verstärken, weshalb viele Modelle FP32 verwenden. NVIDIA trainiert PET für diesen Ansatz nicht neu und ändert seine Architektur nicht. FP32-Operanden werden in niedriger präzise Komponenten zerlegt, ausgewählte Produkte auf Tensor Cores berechnet und in FP32 aufsummiert. BF16x9 nutzt auf Blackwell je drei BF16-Teile und neun Produkte. TF32x3 auf Ampere nutzt drei Produkte und lässt das kleinste Produkt der Restanteile weg. Beide Verfahren besitzen unterschiedliche Genauigkeits- und Leistungsprofile.

Die Versuchsgrößen sind ausdrücklich angegeben. Ein kleiner PET-MAD-Checkpoint mit 26 Millionen Parametern wird auf 192 Wassermolekül-Atomen und einer Superzelle mit 17.280 Atomen getestet. Ein großer Checkpoint mit 255 Millionen Parametern wird für 192 Atome geprüft. Vergleich ist jeweils der zugehörige native FP32-Pfad. Mit BF16x9 auf Blackwell B300 liefen große GEMM-Kerne der 17.280-Atom-Last 2,2- bis 2,5-mal schneller. Der vollständige Molekulardynamik-Schritt wurde dort mit der Strategie performant 1,7-mal und mit eager 1,4-mal schneller. Für 192 Atome betrugen die Gewinne des ganzen Schrittes lediglich 1,2 und 1,1. Ein 2,5-facher Kernelgewinn darf also nicht als Gewinn der gesamten Simulation wiedergegeben werden.

Das Ampere-A100-Experiment setzt TF32x3 über einen CUTLASS-Kern `OpMultiplyAddFastF32` und einen eigens angepassten `TF32Linear`-Ersatz für passende lineare Schichten um. Ein- und Ausgaben bleiben FP32. Große GEMM-Operationen im 17.280-Atom-Fall wurden gegenüber FP32 ungefähr 1,7- bis 1,9-mal schneller; der ganze Schritt gewann etwa 1,5-mal. Beim 255-Millionen-Parameter-Modell auf 192 Atomen lag der Schrittgewinn bei 1,4 gegenüber 1,2 für das kleine Modell gleicher Systemgröße. Eine einzelne TF32-Passage kann zwar schneller sein, hat laut NVIDIA aber größere Nachteile bei Genauigkeit und Trajektorienstabilität. Die schnellste Rechenart liefert nicht zwangsläufig die brauchbarste physikalische Simulation.

Darum wurden weitere physikalische Eigenschaften untersucht. Ein CO-Molekül wurde von einer Pt(111)-Oberfläche entfernt und die Energiekurve betrachtet; außerdem lief eine 30-Pikosekunden-NVE-Simulation eines 192-Atom-Wassersystems. TF32x3 folgte einer glatten, FP32-ähnlichen Kurve und vergleichbar begrenzten Energieschwankungen. Einmaliges TF32 zeigte stärkere Schwankungen oder systematische Drift. Solche Ergebnisse gelten für die geprüften Stoffe und Zeitspannen, nicht für jedes Material und beliebig lange Prozesse.

Die Strukturrelaxation einer Legierungsplatte mit 539 Atomen prüft zusätzlich die Konvergenz. In drei Wiederholungen erreichte FP32 das Ziel nach 218 Schritten, TF32x3 nach 218 bis 222 Schritten bei einer Kraftschwelle von 1,0 × 10⁻³ eV/Å. Alle drei Läufe mit einfachem TF32 verfehlten das Ziel innerhalb von 500 Schritten und laut Quelle auch bei einer Grenze von 1.000. TF32x3 beschleunigte jeden Optimierungsschritt um den Faktor 1,5 und behielt die beobachtete Konvergenz; einfaches TF32 konnte einen Schritt um den Faktor 2 beschleunigen, ohne das Ziel zu erreichen. Nur Schrittzeiten zu vergleichen, wäre hier irreführend.

Für eigene Anwendungen sollten Teams den Anteil großer GEMM-Operationen messen und Energien, Kräfte, längere Trajektorien und Konvergenz prüfen. Kleine Systeme gewinnen oft weniger, weil Nachbarschaftsberechnung, Grapharbeit und Startaufwand mehr Zeit beanspruchen. Hardware und Kernelabstimmung beeinflussen das Resultat. NVIDIA hat den kombinierten Effekt mit anderen möglichen Optimierungen ausdrücklich nicht gemessen. Unsere Redaktion hat die PET-Versuche nicht selbst wiederholt; jeder Quotient bleibt an Anbieter, Hardware, Modellgröße, Atomzahl, Metrik und FP32-Ausgangswert gebunden.