NVIDIA: Tensor Core módszerekkel gyorsítható az atomisztikus szimuláció

A pontosság és a sebesség közötti kompromisszum

Az atomisztikus szimulációk segítségével a kutatók molekulák és anyagok viselkedését vizsgálhatják, például kémiai reakciókat, fázisátalakulásokat vagy mechanikai meghibásodásokat. A sűrűségfunkcionál-elmélet pontos energiákat és erőket képes adni, számításigénye azonban korlátozza a nagy rendszerek és hosszú szimulációk használatát.

A klasszikus interatomikus potenciálok gyorsabbak, de rögzített funkcionális formájuk korlátozhatja a pontosságot és az általánosíthatóságot. A gépi tanulási interatomikus potenciálok, röviden MLIP-k, neurális hálózatokkal becsülik a potenciális energiát és az atomok közötti erőket, így a forrás szerint a kvantummechanikai módszereknél jóval kisebb költséggel teszik lehetővé nagyobb rendszerek vizsgálatát.

A pontosság különösen fontos a molekuladinamikai szimulációkban. Az erőket a becsült energiából származtatják, ezek pedig az egymást követő időlépésekben meghatározzák az atomok sebességét és helyzetét. A kis numerikus hibák ezért idővel felhalmozódhatnak, és a szimuláció pályájának stabilitására is hatással lehetnek.

A Tensor Core-ok használata FP32 mellett

Az NVIDIA fejlesztői olyan Tensor Core-módszereket vizsgáltak, amelyek az FP32-es mátrixszorzás sebességkülönbségét igyekeznek csökkenteni. Az FP32 operandusokat alacsonyabb pontosságú komponensekre bontják, a részeredményeket Tensor Core-okon számítják ki, majd FP32-ben összegezik.

A megközelítés célja, hogy a natív FP32-höz közeli pontosság mellett nagyobb mátrixszorzási teljesítményt biztosítson. A forrás szerint ehhez nincs szükség a modell újratanítására, súlyainak módosítására vagy architektúrájának megváltoztatására.

A Blackwell GPU-khoz az NVIDIA cuBLAS BF16x9 Tensor Core-útvonalat kínál. Ennél minden FP32 operandus három BF16 komponensre bomlik, amelyekből kilenc páros szorzatot állítanak elő, majd a megfelelő skálázással és FP32-akkumulációval kombinálják őket. A funkció környezeti változóval is engedélyezhető, így a kompatibilis PyTorch GEMM-műveletekhez nem kell a modell kódját módosítani.

Az Ampere GPU-kon a TF32x3 hasonló elven működik. Az FP32 operandus egy vezető és egy maradék TF32-komponensből áll, a módszer pedig három keresztterméket számol ki, a maradék és a maradék közötti szorzatot pedig elhagyja. Az NVIDIA CUTLASS ezt az OpMultiplyAddFastF32 operátorral biztosítja.

A PET modellen mért eredmények

A vizsgálatban a Lab COSMO, az EPFL kutatócsoportja által fejlesztett, transformer-alapú Point Edge Transformer, vagyis PET gépi tanulási interatomikus potenciál szerepelt. A kutatók egy 26 millió paraméteres kis modellt teszteltek 192 atomos vízdobozon és egy 17 280 atomos szupercellán, továbbá egy 255 millió paraméteres nagy modellt is futtattak a 192 atomos rendszeren.

A PET futásidejének jelentős részét optimalizált általános mátrixszorzások, GEMM-ek adják. A Blackwell GPU-kon a BF16x9 a nagy GEMM-műveleteket a 17 280 atomos munkaterhelésnél 2,2 és 2,5-szeresére, a 192 atomos rendszernél 1,2 és 1,4-szeresére gyorsította a forrás szerint.

A teljes molekuladinamikai lépésnél a 192 atomos rendszeren a performant stratégia 1,2-szeres, az eager stratégia 1,1-szeres gyorsulást ért el. A 17 280 atomos rendszeren ugyanezek az értékek 1,7-szeres, illetve 1,4-szeres gyorsulást jelentettek. A performant mód akkor használja a BF16x9-et, amikor a cuBLAS várakozása szerint az gyorsabb, az eager mód pedig minden támogatott műveletnél alkalmazza.

Az NVIDIA szerint a BF16x9 által okozott hiba kisebb volt a PET saját nemdeterminisztikusságánál. A vállalat összefoglalója alapján az Ampere GPU-kon vizsgált TF32x3 a nagy GEMM-eket 1,7 és 1,9-szeresére, a teljes molekuladinamikai lépést pedig legfeljebb 1,5-szeresére gyorsította, miközben az erőhibát az FP32-alapérték közelében tartotta.

Mit jelent ez a felhasználóknak?

Az eredmények azt mutatják, hogy az MLIP-modellek futtatói a GPU-k alacsonyabb pontosságú Tensor Core-egységeit úgy is kihasználhatják, hogy közben megtartják az FP32-höz közeli számítási viselkedést. A BF16x9 használatához a forrás szerint a cuBLAS dokumentációja, a TF32x3 integrációjához pedig az NVIDIA ALCHEMI Toolkit dokumentációja ad példákat.

Az NVIDIA ugyanakkor jelzi, hogy az optimalizált gráfkonstrukció, az összevont műveletek, az egyedi kernelek és a CUDA Graphs együttes hatását ebben az értékelésben nem mérték. Az ilyen kombinált gyorsulást ezért minden munkaterhelésnél külön, teljes folyamaton kell meghatározni.