NVIDIA több rétegben gyorsította fel az OpenFold3 és az ehhez kapcsolódó co-folding munkafolyamatokat, hogy a fehérje- és biomolekuláris szerkezet-előrejelzés nagy skálán is használható legyen gyógyszerkutatási és tervezési feladatokban. A megoldások köre magába foglalja a GPU-alapú MSA-generálást (MMseqs2-GPU), a cuEquivariance CUDA-könyvtárat a geometriára alapuló kernelgyorsításokhoz, az OpenFold3 NIM-re optimalizált inference-t, valamint a Fold-CP kontextuspárhuzamosítási technikát a több-GPU-s skálázáshoz. Ezeket az eszközöket az NVIDIA BioNeMo Agent Toolkit teszi egyszerűen hívhatóvá agentekből.
Miért fontos ez?
Az MSA-generálás, a co-folding inference, a kiszolgálás és a több-GPU-s skálázás mind kritikus lépések egy end-to-end szerkezet-előrejelző pipeline-ban. Ha bármelyik lépés szűk keresztmetszetet képez, az korlátozza az egész rendszer áteresztőképességét. Gyorsaság és memóriahatékonyság különösen fontos nagy könyvtárakon végzett virtuális szűrésnél (milliók–milliárdok vegyület), vagy nagy többfehérje-összetételű rendszerek (több ezres aminosavszám) előrejelzésénél.
GPU MSA: az MSA-bottleneck oldása
A homológ keresés hagyományosan CPU-kötött lépés volt, ami gyakran dominálja a teljes futási időt. Az MMseqs2-GPU áthelyezi ezt a munkát NVIDIA GPU-kra, csökkentve a falióra-időt és skálázódva a szekvenciahosszal mind a NVIDIA Hopper-, mind a Blackwell-architektúrákon. A legújabb GPU-gyorsított verzió Hopper- és Blackwell-specifikus optimalizációkat tartalmaz, továbbá hatékony támogatást nagyobb, a GPU memóriájánál nagyobb adatbázis-kereséshez NVIDIA Grace rendszereken és további gyorsításokat használ Blackwell DPX utasításokkal (CUDA 13.2).
Az MSA Search NIM MMseqs2-GPU-t használ; az eredeti Nature Methods cikk szerint egyetlen L40S-en a CPU JackHMMER-hez képest akár 177× gyorsabb igazítás érhető el. NVIDIA mérései szerint az MSA-fázis latenciája H100 és B300 GPU-kon folytonosan skálázódik 10k+ tokenig (példaként ~429 s vs ~463 s körül ~10k tokennél a bemutatott ábra szerint). Az MSA Search NIM önállóan hívható, self-hosted vagy agentes eszközként is beágyazható. (NVIDIA megjegyzése: build.nvidia.com hosztolt végpontok nem alkalmasak nagy volumenű lekérésekre; nagy munkaterhelésekhez a self-hosted NIM ajánlott.)
cuEquivariance és OpenFold3 NIM: gyorsabb, hosszabb futtatások
A cuEquivariance egy CUDA-X könyvtár, amely geometriai tanulási primitíveket gyorsít (például Triangle Attention, Triangle Multiplication és Attention Pair Bias kernel-ek), amelyek a co-folding során dominálnak. B300-on a cuEquivariance a latency-t nagyjából 1.8–3.1×-re csökkentette különböző szekvenciahosszakon (az NVIDIA által közölt példák: 1,024 token H100-on 79.4 s → 41.6 s, 1,024 token B300-on 49.1 s → 27.3 s, 1,536 token B300-on 149.0 s → 56.2 s, 2,048 token B300-on 300.1 s → 97.6 s).
A cuEquivariance közvetlenül integrálható nyílt forráskódú modellekbe, mint az OpenFold3 (opcionális függőség), OpenFold2, RosettaFold3, Protenix és Boltz. A gyorsítások upstreamelve vannak az OSS modellekbe, így a kutatók automatikusan profitálnak belőlük NVIDIA GPU-kon futtatva. Emellett a cuEquivariance kibővíti az elérhető maximális szekvenciahosszt körülbelül 5.9k tokenig, míg natív PyTorch sok esetben 1.5k–2.5k token körül kifogy a memóriából.
Az OpenFold3 NIM további inferencia-optimalizációkat ad hozzá a cuEquivariance fölé, amellyel egyetlen B300-on elérhető volt akár ~6,400 tokenes futtatás és jelentős, több mint 3×–4× gyorsulás H100 és B300 viszonyításában a teljesen optimalizált NIM esetén.
Fold-CP: több-GPU-s memória-paralelizáció
A single-GPU memória korlátja korábban a co-folding modellt pár ezer maradékig korlátozta. B300-on a nagyobb HBM és a Blackwell-architektúra hatékonysága, a cuEquivariance és az inferenciaoptimalizálások együtt önmagukban is növelik a single-GPU limitet, de sok esetben ez még nem elegendő.
A Fold-CP egy új párhuzamosítási technikát vezet be, amely szerint az egyes eszközök memóriaigénye O(N²/P)-ben skálázódik (N a token-szám, P a GPU-k száma). NVIDIA példája: 64 B300 GPU-val a Boltz-2 modellel 32,000 tokenig terjedő futtatás érhető el — ez kb. 12×-os növekedés a single-GPU limithez képest. A Fold-CP kód és példaparancsok elérhetők a Boltz-CP repository-ban; agenten keresztül is meghívható.
Végponti pipeline-gyorsítás és gyakorlati hatások
NVIDIA megközelítése mindig a teljes pipeline-re fókuszál: MSA Search NIM a homológ keresést gyorsítja (akár 177× a CPU-hoz képest), cuEquivariance és az OpenFold3 NIM csökkentik az inferencia-latenciát (akár több mint 3–4×), és Fold-CP lehetővé teszi a co-foldinget jóval a single-GPU korlátain túl (32,000 token 64 B300-mal). Ezek együtt gyorsabb, skálázhatóbb és agentekből könnyebben összeállítható munkafolyamatokat eredményeznek, amelyek megkönnyítik a modellszintű előrejelzésektől nagyobb, hasznosabb biológiai rendszerekig való eljutást.
Gyakorlati következmények:
- Virtuális szűrés: a gyorsabb co-folding révén a szerkezet-alapú módszerek korábban és nagyobb könyvtárakon alkalmazhatók, növelve az esélyét a jobb és változatosabb találatoknak.
- Nagy komplexumok modellezése: a megnövelt single-GPU kapacitás és a context-parallel Fold-CP lehetővé teheti olyan rendszerek előrejelzését, amelyek eddig gyakorlatilag elérhetetlenek voltak co-folding modellekkel (például riboszóma-méretű komplexumok ~10,000 maradék felett), ha több GPU áll rendelkezésre.
Hozzáférés és kezdés
A gyorsításokat és eszközöket az NVIDIA BioNeMo Agent Toolkit-en keresztül lehet kipróbálni. Fontos komponensek:
- MSA Search NIM (MMseqs2-GPU)
- cuEquivariance (CUDA-X könyvtár)
- OpenFold3 NIM (optimalizált inference)
- Fold-CP (Boltz-CP repository a több-GPU-s kontextuspárhuzamosításért)
Az NVIDIA közleménye kitér arra is, hogy a NIM-ek letölthetők és self-hosted módon is futtathatók, ami ajánlott nagy volumenű munkákhoz.
Köszönetnyilvánítás
NVIDIA megköszöni a belső csapat tagjainak (például Franco Pellegrini, Lalit Vaidya, Duc Tran, Tien Pham, Maximilian Stadler, Alejandro Chacon, Quan Vu, Simon Chu, Brian Roland, Dejun Lin, Joseph Chang, Hoa La, Jonathan Mitchell, Vishanth Iyer, Timur Rvachov, Christian Dallago, Christian Hundt) munkáját, valamint az OpenFold és OMSF közösségek együttműködését és hozzájárulásait.
Források és további linken elérhető anyagok elérhetők az NVIDIA BioNeMo és a kapcsolódó GitHub-repozitóriumok dokumentációjában.
