1,7 kata kadar hız: NVIDIA($NVDA), B300’te moleküler dinamik hesabını hızlandırdı

NVIDIA($NVDA), 17.280 atomdan oluşan bir su sisteminde B300 GPU’nun BF16x9 yöntemini kullanarak PET moleküler dinamik hesaplamalarını FP32’ye kıyasla 1,7 kata kadar hızlı gerçekleştirdi. Bu ölçüm yalnızca matris çarpımlarını değil, moleküler dinamikteki bir tam hesaplama adımını kapsıyor.

NVIDIA, 29’unda yayımladığı teknik blogda makine öğrenimi tabanlı atomlar arası potansiyel PET’te (Point Edge Transformer) tensör çekirdeklerine dayalı iki matris çarpımı yönteminin sonuçlarını açıkladı. BF16x9, B300 (Blackwell) GPU’da; TF32x3 ise A100 GPU’da denendi. Performans değerleri, belirli modeller ve donanımlarla yapılan NVIDIA ölçümlerine dayanıyor.

Makine öğrenimi tabanlı atomlar arası potansiyeller, atom yapılarından enerji ve atomlar arasındaki kuvvetleri tahmin ederek moleküllerin ve malzemelerin hareketini hesaplamak için kullanılıyor. Deneyde, yoğun hesaplama gerektiren FP32 matris çarpımlarının tensör çekirdeklerinde çalıştırılmasının performans ve sonuçlar üzerindeki etkisi incelendi.

PET, İsviçre’deki Lozan Federal Teknoloji Enstitüsü’nün (École polytechnique fédérale de Lausanne, EPFL) COSMO araştırma grubu tarafından geliştirildi. NVIDIA, 26 milyon parametreli küçük PET-MAD modelini 192 ve 17.280 atomlu su sistemlerinde değerlendirdi. 255 milyon parametreli büyük model ise 192 atomlu sistemde test edildi.

BF16x9 yöntemi, FP32 değerlerini üç BF16 bileşenine ayırıyor; dokuz kısmi çarpımın ardından sonuçları FP32 olarak topluyor. Yöntem, büyük matris çarpımlarını tensör çekirdeklerinde çalıştırmak üzere tasarlandı.

17.280 atomlu sistemde BF16x9, temel matris çarpımlarını FP32’ye göre 2,2 ila 2,5 kat hızlandırdı. Matris çarpımları dışındaki işlemleri de içeren bir moleküler dinamik adımı, çalışma biçimine bağlı olarak 1,4 ila 1,7 kat hızlandı.

192 atomlu sistemde BF16x9 ile temel matris çarpımları 1,2 ila 1,4 kat, moleküler dinamik adımları ise 1,1 ila 1,2 kat daha hızlı tamamlandı. NVIDIA, bu yöntemle tek bir zaman noktasında hesaplanan enerji ve kuvvetlerdeki hatanın PET modelinin kendi deterministik olmayan davranışından kaynaklanan değişkenlikten daha düşük olduğunu belirtti.

A100’de denenen TF32x3, FP32 değerlerini ana bileşen ve artık bileşenlere ayırıyor; üç TF32 çarpımının sonuçlarını FP32 olarak topluyor. Bu yöntem, BF16x9’dan farklı bir hassasiyet biçimi kullanıyor.

17.280 atomlu sistemde TF32x3, temel matris çarpımlarını 1,7 ila 1,9 kat hızlandırırken moleküler dinamik adımında en fazla 1,5 kat hız artışı sağladı. 192 atomlu sistemde moleküler dinamik adımları küçük modelde 1,2 kat, büyük modelde 1,4 kat daha hızlı tamamlandı.

Her iki yöntem de eğitilmiş modelin ağırlıklarını veya yapısını değiştirmeden matris çarpımlarının yürütülme biçimini değiştirdi. Modeli yeniden eğitmek gerekmiyor, ancak uygulama yöntemleri farklı. BF16x9, cuBLAS’ın desteklediği bir işlevden yararlanırken TF32x3, PET’in doğrusal katmanlarına özel bir çekirdek uyguladı.

Uzun süreli enerji korunumu ve yapı optimizasyonu testleri yalnızca TF32x3 için paylaşıldı. 192 atomlu su sistemindeki enerji değişimi FP32 ile benzer aralıkta kaldı. Platin yüzeyinde hareket ettirilen karbonmonoksit molekülü için ölçülen enerji eğrisi de FP32 sonuçları kadar düzgün seyretti.

539 atomlu Cantor alaşımı yüzeyinin yapı optimizasyonunda FP32, kuvvet eşiği olan 1,0×10⁻³ eV/Å değerine 218 iterasyonda; TF32x3 ise 218 ila 222 iterasyonda ulaştı. Tek geçişli TF32, 500 iterasyon içinde yakınsamadı ve hesaplama 1.000 iterasyona çıkarıldığında da aynı eşiği karşılamadı.

BF16x9 için aynı türde uzun süreli enerji korunumu ve yapı optimizasyonu test sonuçları paylaşılmadı. Bu nedenle iki yöntemin fiziksel tutarlılığının aynı kapsamda doğrulandığı söylenemez.

Performans farkı sistemin büyüklüğüne ve toplam hesaplama içinde matris çarpımlarının payına göre değişiyor. NVIDIA, paylaşılan değerlerin belirtilen PET modelleri, GPU’lar ve çekirdek ayarlarıyla elde edildiğini bildirdi.

NVIDIA’nın CUDA cuBLAS belgelerinde BF16x9 FP32 emülasyonunun desteklendiği GPU’lar için 10.0 ve 10.3 hesaplama yetenekleri ile CUDA 12.9 veya üzeri sürüm koşulu belirtiliyor.