“Ölçeklenebilir AI Tabanlı Moleküler Dinamik Simülasyonlarının Aktif Hale Getirilmesi”

Moleküler Dinamik Simülasyonları, hesaplamalı kimya ve malzeme bilimi alanında güçlü bir araçtır. Kimyasal reaksiyonlar, malzeme özellikleri ve biyolojik etkileşimleri mikroskobik düzeyde incelemek için gereklidir. Ancak, bu simülasyonların karmaşıklığı ve hesaplama gereksinimleri, genellikle makine öğrenimi etkileşimsel potansiyelleri (MLIP’ler) gibi ileri düzey tekniklerin kullanılmasını gerektirir; bu teknikler, ölçeklenebilirlik, verimlilik ve doğruluk sağlar.

PyTorch tabanlı MLIP’leri LAMMPS moleküler dinamik paketine entegre etmek, ML-IAP-Kokkos arayüzü aracılığıyla kimya ve malzeme bilimi araştırmaları için hızlı ve ölçeklenebilir atom sistemleri simülasyonu yapma imkanı sunar.

Bu arayüz, NVIDIA, Los Alamos Ulusal Laboratuvarı ve Sandia Ulusal Laboratuvarı bilim insanları arasındaki işbirliği ile geliştirilmiştir. Geliştiricilerin topluluk modellerini kolaylıkla bağlayabilmesi için tasarlanmış olan bu arayüz, MLIP’leri LAMMPS ile entegre ederek ölçeklenebilir moleküler dinamik simülasyonları sağlamaktadır.

ML-IAP-Kokkos arayüzü, mesajlaşma MLIP modellerini destekler ve LAMMPS’ın yerleşik iletişim yeteneklerinden yararlanarak GPU’lar arasında verimli veri transferini kolaylaştırır. Bu, çoklu GPU’ların hesaplama gücünden faydalanan büyük ölçekli simülasyonları mümkün kılar. Arayüz, Python ile C++/Kokkos LAMMPS arasında köprü kurmak için Cython kullanarak, simülasyon iş akışını optimize eder ve GPU hızlandırması sağlar.

Gerekli Araçları ve Ortamı Hazırlamak

Simülasyonları başlatmadan önce aşağıdaki gereksinimlerin karşılandığından emin olun:

  • LAMMPS veya diğer MD simülasyon araçları hakkında deneyim.
  • Python ile birlikte PyTorch ve makine öğrenimi modelleri (Cython kullanımı pluss).
  • Opsiyonel: Ek hızlandırma için cuEquivariance destekli modeller ve NVIDIA GPU’ları.

ML-IAP-Kokkos Arayüzünü Kurmak

Aşama 1: Ortamı Kurmak

Tüm gerekli yazılımların kurulu olduğundan emin olun:

  • LAMMPS Kokkos, MPI ve ML-IAP ile birlikte çalışacak şekilde derlenmiş olmalıdır.
  • Python ve PyTorch ile birlikte.
  • Eğitilmiş bir PyTorch MLIP modeline sahip olun (isteğe bağlı olarak cuEquivariance destekli).

Aşama 2: LAMMPS’i ML-IAP-Kokkos/Python Desteği ile Kurmak

LAMMPS 2025 Eylül sürümünü veya daha güncel bir sürümünü Kokkos, MPI, ML-IAP ve Python desteği ile indirin ve kurun.

Python desteği ile LAMMPS’i kurduğunuzda, lammps modülü Python ortamınıza da yüklenmiş olacaktır.

Kolaylık olması açısından, önceden derlenmiş bir konteyner hazırladık.

Aşama 3: Tercih Edilen Model için ML-IAP-Kokkos Arayüzünü Geliştirmek

ML-IAP-Kokkos arayüzünü kullanırken, LAMMPS simülasyon sırasında Python yorumlayıcısını çağıracak ve yerel Python modelinizi çalıştıracaktır. Böylece, modelinizi derleme zorluğu olmadan kullanabileceksiniz. Bunun için, mliap_unified_abc.py dosyasındaki MLIAPUnified soyut sınıfını uygulamanız gerekecek. Özellikle, compute_forces fonksiyonunu, LAMMPS’tan geçen verilerle çift yönlü kuvvetleri ve enerjileri tahmin edecek şekilde tanımlamanız gerekiyor.

Basit bir örnek ile başlayarak, LAMMPS’ın sınıfınıza geçirdiği verileri anlamaya çalışalım:

from lammps.mliap.mliap_unified_abc import MLIAPUnified
import torch


class MLIAPMod(MLIAPUnified):
    def __init__(
        self,
        element_types = None,
    ):
        super().__init__()
        self.ndescriptors = 1
        self.element_types = element_types
        self.rcutfac = 1.0 # Yarı çap kesme
        self.nparams = 1


    def compute_forces(self, data):
        print(f"Toplam atom sayısı: {data.ntotal}, Yerel atom sayısı: {data.nlocal}")
        print(f"Atom indeksleri: {data.iatoms}, Atom türleri: {data.elems}")
        print(f"Komşu çiftleri: {data.npairs}")
        print(f"Çift indekleri ve yer değiştirme vektörleri: ")
        print("n".join([f"   ({i}, {j}), {r}" for i,j,r in zip(data.pair_i, data.pair_j, data.rij)]))

Yukarıdaki __init__ fonksiyonu, modelimizin ele alacağı öğe türlerini ve komşu listesinin kesim yarıçapını tanımlamaktadır. Bu parametreler, LAMMPS üzerinde modelin kurulabilmesi için gereklidir.

Model nesnemizi PyTorch’un save işlevi ile kaydediyoruz:

mymodel = MLIAPMod(["H", "C", "O"])
torch.save(mymodel, "my_model.pt")

Bu, LAMMPS’ın yükleyebileceği, bir “pickle” Python nesnesi olan my_model.pt dosyasını oluşturur. Gerçek bir senaryoda bu, modelimizin ağırlıklarını içerecektir fakat burada sınıfımızın bir seri hale getirilmesidir. (PyTorch’un son sürümlerinde, bir sınıfı yükleme güvenlik koruması olarak varsayılan olarak devre dışı bırakılmıştır; TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 ortam değişkenini ayarlamak, sınıfı yüklemenizi sağlar. Sadece güvenilen *.pt dosyalarını kullanın).

Bir CO2 molekülü sistemini tanımlamak için sample.pos adlı bir dosya oluşturuyoruz:

#Zorunlu olarak ilk satır bir yorumdur ve atlanır!
3 atom
2 atom türü

0.0 100.0 xlo xhi
0.0 100.0 ylo yhi
0.0 100.0 zlo zhi

Kütleler

1 12.0
2 16.0

Atomlar

1 2 1.0 10.0 10.0
2 1 2.0 10.0 10.0
3 1 -0.1 10.0 10.0

Bu molekül, x-ekseni boyunca lineer olarak dizilmiş bir karbon atomu ve iki oksijen atomundan oluşmaktadır; karbon, x = 1.0 Å konumundayken, oksijen atomları sırasıyla x = -0.1 Å ve x = 2.0 Å’dadır. Bu minimal yapı, LAMMPS’ın yerel ve hayalet atomlar, çift etkileşimler ve periyodik sınır koşullarını nasıl ele aldığına dair net bir gözlem yapmamıza olanak tanır; bu, arayüzde veri akışını anlamak için kritik unsurlardır.

Aşağıdaki LAMMPS giriş dosyası sample.in, konum dosyasını yükler, potansiyelimizi yükler ve bir moleküler dinamik adımı çalıştırır:

units           metal
atom_style      atomic
atom_modify     map array
boundary p p p
read_data   sample.pos

# ML-IAP modelini yükle
pair_style      mliap unified my_model.pt 0
pair_coeff      * * C O

dump dump0 all custom 1 output.xyz id type x y z fx fy fz
run 0

pair_style kelimesini kullanarak, mliap ile “unified” anahtar kelimesi, LAMMPS’a Python arayüzünü kullanmasını belirtir, ardından yüklemek için model dosyanız gelir ve bu durumu etkilemeyen 0 değeri eklenir.

İlgilendiğimiz atom türlerini (C ve O) ve bunların konum dosyasındaki indekslerle eşleştirilmesini tanımlarız.

LAMMPS’ı bir GPU ile Kokkos ile çalıştırmak için aşağıdaki komutu kullanırız:

lmp -k on g 1 -sf kk -pk kokkos newton on neigh half -in sample.in

Ve çıktıdaki ilgili kısmı kontrol ettiğimizde:

Toplam atom sayısı: 6, Yerel atom sayısı: 3
Atom indeksleri: [0 1 2], Atom türleri: [2 1 1 2 1 1]
Komşu çiftleri: 4
Çift indeksleri ve yer değiştirme vektörleri:
   (0, 5), [-1.1  0.   0. ]
   (0, 1), [1. 0. 0.]
   (1, 0), [-1.  0.  0.]
   (2, 3), [1.1 0.  0. ]

Sınıfımız ile karşılaştırdığımızda, veri yapısını anlamış oluyoruz:

  • Atomlar iki kategoriye ayrılır: güncellemeye ihtiyaç duyan fiziksel atomlar, [0, nlocal) aralığındadır ve nlocal ve ntotal arasındaki ghost atomlardır; bu atomlar periyodik kopyaları veya başka bir işlemci/GPU’dan gelen fiziksel atomlardır.
  • Toplamda altı atom vardır; ancak nlocal olan üç fiziksel atom bulunur.
  • iatoms, yerel atomların indekslerini (ilk nlocal) listeler.
  • elems, tüm atomların türlerini listeler.
  • Bir çiftin en az bir yerel atom içermesi durumunda geçerli npairs sayısını belirtir.
    Not: İlişkili (i,j) ve (j,i) çiftleri ayrı olarak rapor edilmektedir.

Her çiftin (i,j) indekslerine ve göreceli yer değişimine pair_i, pair_j ve rij ile ulaşabiliriz.
nlocal ağırlıklı indeksi >= nlocal olanlar hayalet atomlardır (örneğin 3, 0’ın periyodik hayaletidir, 5 ise 2’nin periyodik hayaletidir).

Basit bir potansiyel uygulaması olarak, yerel atomlar üzerinde bazı “düğüm özellikleri” hesaplayıp, komşu özellikleri ve bağ uzunluklarına göre enerjiyi güncelleyebiliriz; bu, MLIP’ler için en yaygın düzenektir. Örnek olarak, tüm özelliklerin 1 kabul edildiği ve radial fonksiyonun atomik mesafe olduğu kondisyonda, tam uygulama olmayabilir.

Aşağıdaki basit potansiyeli içeren bir kodu inceleyelim:

   def compute_forces(self, data):
        rij = torch.as_tensor(data.rij)
        rij.requires_grad_()
        features = torch.zeros((data.ntotal,1)).cuda() 
        for i in range(data.nlocal):
            features[i,0] = 1.0
        Ei = torch.zeros(data.nlocal).cuda()

        # Yerel mesaj geçişi adımı (j → i)
        for r,i,j in zip(rij, torch.as_tensor(data.pair_i), torch.as_tensor(data.pair_j)):
            Ei[i] += features[j,0] * torch.norm(r)

        # Toplam enerji ve kuvvetleri hesapla
        Etotal = torch.sum(Ei)
        Fij = torch.autograd.grad(Etotal, rij)[0]

        # Toplam enerji ve kuvvetleri güncelle
        data.energy = Etotal
        data.update_pair_forces_gpu(Fij)

Bu kodda dikkat edilmesi gereken birkaç unsur var:

  • Bütün veri yapılarının fp64 formatında ve cihazda olduğunu unutmayın; ancak tip geçişleri ve host<->cihaz hareketi model içerisinde izinlidir.
  • Toplam enerji doğrudan modelimizde güncellenir.
  • Kuvvetler, rij’ye göre otomatik türev hesabıyla (autograd) hesaplanır, yani kuvvetler “her bağ” için hesaplanır ve “her atom” için değil. Atomik kuvvetler doğrudan hesaplayan potansiyellerin, bu formata yeniden düzenlenmesi gerekir.

Kuvvetleri LAMMPS’ta update_pair_forces_gpu fonksiyonu aracılığıyla güncelliyoruz; burada fp64 GPU tensörü (npairs, 3) kullanılır.

Bu kodu CO2 molekülü üzerine test ediyoruz.

Referans olsun diye, sadece iki çift ve iki mesafeyi hesaplayarak, toplam enerji 4.2 olmalı, her çift 1 kuvvet katkıda bulunmalı ve karbon atomları için 2’lik kuvvet ve oksijen için ise 0 kuvvet elde edilmelidir.

Bununla birlikte bu kodu çalıştırdığımızda, toplam enerjinin 2 (çıkışta görüldüğü gibi) ve kuvvetin de 2 olduğunu görüyoruz, bir karbonda ve oksijende böyle bir durum söz konusu. Sorunun kaynağı hayalet atomların bulunmasıdır; bu atomların özellikleri, gerçek atomlar üzerinde güncellenmediği için yalnızca (0, 1) bağı toplam enerji ve kuvvetlere katkıda bulunur.

Bu basit senaryoda, hayalet atomları 1 ile başlatabilirdik; ancak gerçek bir senaryoda, bu daha karmaşık bir çevre fonksiyonuna dayalı olabilir. Bu nedenle, ML-IAP-Kokkos’ta yer alan mesaj geçişi araçlarını kullanarak hayalet özelliklerini karşılık gelen fiziksel atomların değerine güncellememiz gerekmektedir. Bu, periyodik sınır koşullarındaki ve tamamen farklı bir GPU’da bulunan hayalet parçacıkları için önlem alır.

Bu başlatma sürecine yardımcı olacak iki rutin vardır: forward_exchange, bu fiziksel atomların ghost atomlarını ayarlamakta ve reverse_exchange, ghost atomları üzerinden geri gelirken toplam katkıları fiziksel atomlara geri kazandırmaktadır.

Aşama 4: ML-IAP-Kokkos için Mesaj Geçişi Desteğini Uygulamak

Modelimizi güncelleyerek mesaj geçişi adımını (aşağıdaki kod) kullanarak ekleyelim; forward_exchange fonksiyonu, ntotal ve feat_size boyutundaki iki tensörü yükleyerek, features[0,nlocal) aralığını new_features[0,ntotal) ile ayarlayacaktır; bu özelliklerin genişliği buradaki feat_size ile aktarılır. Bu fonksiyon, hayalet değerlerini, fiziksel karşıtlarının özellikleri ile doldurarak bir kopya oluşturur.

 def compute_forces(self, data):
        rij = torch.as_tensor(data.rij)
        rij.requires_grad_()
        features = torch.zeros((data.ntotal,1)).cuda()
        for i in range(data.nlocal):
            features[i,0] = 1.0
        Ei = torch.zeros(data.nlocal).cuda()

        # Özelliklerin güncellenmesi
        new_features = torch.empty_like(features)
        data.forward_exchange(features, new_features, 1)
        features = new_features

        # Yerel mesaj geçişi adımı
        for r,i,j in zip(rij, torch.as_tensor(data.pair_i), torch.as_tensor(data.pair_j)):
            Ei[i] += features[j,0] * torch.norm(r)

        # Toplam enerji ve kuvvetleri hesapla
        Etotal = torch.sum(Ei)
        Fij = torch.autograd.grad(Etotal, rij)[0]

        # Toplam enerji ve kuvvetleri güncelle
        data.energy = Etotal
        data.update_pair_forces_gpu(Fij)

Bu kodu tekrar çalıştırdığımızda, doğru enerji ve kuvvetleri şimdi elde ediyoruz. Ancak, birden fazla GPU ile çalıştığımızda kuvvet hesaplama doğru çalışmayacaktır, çünkü gradyanlar farklı cihazlar arasında hesaplanmamaktadır.

Bu durumu düzeltmek için, LAMMPS_MP.forward() fonksiyonunun gradyanını PyTorch ile tanımlayıp kaydetmemiz gerekiyor. forward_exchange‘in gradyanı, kullanılan reverse_exchange ile sağlanır. Bu, gradyanları gerçek atomlara hapseder ve tüm hayalet atomlarında sıfırlanır.

Mesaj geçişinin tamamlanmış bir uygulamasının basit bir modeli aşağıda verilmiştir:

import torch 
class LAMMPS_MP(torch.autograd.Function):
    @staticmethod
    def forward(ctx, *args):
        feats, data = args  # unpack
        ctx.vec_len = feats.shape[-1]
        ctx.data = data
        out = torch.empty_like(feats)
        data.forward_exchange(feats, out, ctx.vec_len)
        return out


    @staticmethod
    def backward(ctx, *grad_outputs):
        (grad,) = grad_outputs  # unpack
        gout = torch.empty_like(grad)
        ctx.data.reverse_exchange(grad, gout, ctx.vec_len)
        return gout, None

Yeni bu sınıf ile birlikte, son versiyon pluginimiz şöyle olacaktır:

 def compute_forces(self, data):
        rij = torch.as_tensor(data.rij)
        rij.requires_grad_()
        features = torch.zeros((data.ntotal,1)).cuda()
        for i in range(data.nlocal):
            features[i,0] = 1.0
        Ei = torch.zeros(data.nlocal).cuda()

        # Özelliklerin güncellenmesi
        features = LAMMPS_MP().apply(features, data)

        # Yerel mesaj geçişi adımı
        for r,i,j in zip(rij, torch.as_tensor(data.pair_i), torch.as_tensor(data.pair_j)):
            Ei[i] += features[j,0] * torch.norm(r)

        # Toplam enerji ve kuvvetleri hesapla
        Etotal = torch.sum(Ei)
        Fij = torch.autograd.grad(Etotal, rij)[0]

        # Toplam enerji ve kuvvetleri güncelle
        data.energy = Etotal
        data.update_pair_forces_gpu(Fij)

Artık doğru bir şekilde çoklu-GPU mesaj geçişini destekleyebiliyoruz ve autograd’dan faydalanabiliyoruz.

Bu yapı, mesaj geçişi MLIP modellerinden bağımsız olarak çalışacaktır.

Herhangi bir mesaj geçişinden önce LAMMPS_MP‘yi çağırmak, özelliklerin senkron kalmasını garantiler ve modelimizin doğru çalışmasını sağlar. Ancak bazı iletişim adımlarının atlanabileceği optimizasyonlar mevcut olabilir. LAMMPS, önceki adımlarda hayaletlerin toplanmasını sağladığı için ve kuvvetleri hayaletlere yaymayı sağlayarak, bu durumda forward_exchange ve son reverse_exchange adımlarını sıklıkla atlayabiliriz.

Verimlilik sağlamak için, kuvvet hesaplaması bağlar ile ilgili enerji gradyanlarının hesaplanmasını gerektirir. Ayrıca tüm terimlerin LAMMPS’a gönderilmesi, LAMMPS’ın stres hesaplamaları için katkıları toplamasını sağlar; stresler plugin tarafından hesaplanmamalıdır.

HIPPYNN Modeli ile Performans Değerlendirmesi

Mesaj geçişi araçlarının sağladığı performans iyileştirmelerini göstermek amacıyla, HIPPYNN (Chigaev 2023) modelleri ile birden fazla 512 NVIDIA H100 GPU kullanarak zayıf ölçeklenebilirlik testlerini gerçekleştirdik. Her bir model, MLIP’de h=1 ile h=4 katmanlı etkileşimlere sahiptir ve testlerde toplamda her bir GPU’da yaklaşık 203 atom bulunur.

Modeller, Alüminyum için sağlam bir etkileşimsel potansiyelin otomatik keşfi adlı çalışmadan yararlanarak eğitim sırasında kullanıldı. Doğru simülasyon sağlamak için, r MLIAP katmanının etkileşim yarıçapını temsil ederken, hayalet kopya uzunluğunu r+epsilon ile köklenmiş katman özelliklerinin kullanıldığı durumda aynı işlemler uygulanır.

The top chart shows atom-steps per second per GPU executed (higher is better) versus number of GPUs used in weak scaling tests of LAMMPS running HIPPYNN ML-IAP models. Orange, blue, green, and pink lines indicate tests using 1, 2, 3, 4 layer MLIPs respectively  and solid and dashed lines indicate tests using and not using inter-layer communication hooks. The bottom chart shows the speedup from using the communication hooks versus the number of GPUs.
Şekil 2. Yukarıdaki grafik, daha fazla GPU sayısına (y ekseni) karşı atom adımı başına saniye (x ekseni) değerlerini gösteriyor. Farklı renkler, farklı katman testlerini (1, 2, 3, 4) temsil eder, bütün kütle kuvvetlerinin sağladığı hızlanmaları karşılaştırır.

İletişim araçlarının etkinleştirilmesinden kaynaklanan hızlanma, hayalet atomların yüzde oranından doğrudan etkilenmiştir. Bir, iki, üç, ve dört katmanlı hayalet atomlarının toplam atom sayısına oranı %54, %38, %26 ve %18 olmuştur. İletişim araçları aktif olduğunda, her durumda 54% düzeyinde gerçek atomlar kullanılarak tek bir radyo yeterli olduğundan, dolayısıyla iki, üç ve dört katman için 1.4x, 2.1x ve 3x toplam atom azalması sağlanmıştır; bu gözlemlenen hızlanmalarla doğrudan ilişkilidir.

MACE Entegrasyonunu Kullanarak Sonuçları Karşılaştırma

Bir başka gösterim olarak, popüler MLIP MACE kullanarak, LAMMPS’taki orijinal MACE pairstyle’ı ile karşılaştırma yapılmıştır; daha çok libtorch kullanarak C++ üzerinden model ağırlıklarını yüklemektir. cuEquivariance, geometrik sinir ağlarını hızlandırmak için NVIDIA’nın Python kütüphanesidir.

Şekil 3. Farklı LAMMPS eklentileri ve farklı hassasiyetler (fp32/fp64) için sistem boyutuna bağlı simülasyon süresi grafiği.

Bu test için, GPU bellek sınırlarına kadar artan su kutuları simüle edilmiştir; medium versiyonunu kullanarak MACE-OFF23 potansiyeli ile tek bir A100 GPU’da 80GB bellekte simülasyon yapılmıştır. Orijinal MACE pairstyle’ın hızı ise referans olarak gösterilmiştir. Yeni eklenti, çok daha hızlı ve hafıza açısından daha verimlidir. Bu hızlanmanın kaynağı, cuEquivariance ile sağlanan model hızlandırması ve MLIAP-Kokkos arayüzündeki mesaj geçiş sıkıştırmalarıdır. Hızlandırılmış MLIP’ler, LAMMPS içerisinde ML-IAP-Kokkos arayüzünden geçerek sorunsuz bir şekilde kullanılabilir.

Sonuç

LAMMPS ML-IAP-Kokkos arayüzü, şu anda çoklu-GPU ve çoklu-düğüm moleküler dinamik simülasyonları için en önemli araçtır ve modern ML-tabanlı kuvvet alanları ile yüksek performanslı işleme altyapıları arasındaki boşluğu kapatmaktadır.

PyTorch tabanlı MLIP’leri LAMMPS’a entegre ederek, geliştiriciler ölçeklenebilir ve verimli moleküler dinamik simülasyonları gerçekleştirebilirler. Bu kılavuz, simülasyonları başlatmak ve yürütmek için gereken adımları özetler ve LAMMPS’daki GPU hızlandırması ve mesaj geçişi yeteneklerinin faydalarını vurgular. Benchmark sonuçları, büyük ölçekli simülasyonlardaki arayüzün potansiyelini daha da öne çıkarmaktadır; bu nedenle araştırmacılar ve geliştiriciler için değerli bir araçtır.

Simülasyonlarınızı hızlandırmaya hazır mısınız? Bugün kılavuzu deneyin ve sonuçlarınızı toplulukla paylaşın. NVIDIA haberlerine abone olun ve NVIDIA AI‘yı LinkedIn, X, Discord ve YouTube‘da takip edin.

Kaynak

Nvdia Blog

Exit mobile version