NVIDIA BioNeMo : le contexte parallèle pour modéliser les biomolécules
L'équipe de NVIDIA BioNeMo a présenté une nouvelle approche nommée parallélisme contextuel (CP) pour surmonter les limitations de mémoire qui entravent la modélisation biomoléculaire depuis des décennies. Traditionnellement, la nécessité de faire entrer des systèmes biologiques complexes dans la mémoire vive limitée d'un seul GPU obligeait les chercheurs à réduire ces systèmes à des fragments isolés, comme des protéines uniques ou de petits domaines. Cette méthode de réduction créait un écart contextuel, rendant impossible le repliement précis de grandes protéines ou de complexes entiers sans sacrifier la précision structurelle globale. Le cadre de parallélisme contextuel de BioNeMo permet désormais de modéliser ces systèmes de manière holistique en répartissant une seule grande molécule sur plusieurs GPUs. Contrairement au parallélisme de données classique qui affecte à chaque GPU une protéine différente, cette technique divise l'échantillon massif en plusieurs parties. L'architecture repose sur des APIs de distribution Torch et utilise une stratégie d'fragmentation multidimensionnelle pour atteindre une mise à l'échelle linéaire de la capacité. Cela garantit qu'aucun dispositif individuel ne détient l'état global de la biomolécule, localisant ainsi l'empreinte mémoire et évitant les goulots d'étranglement. Le système met en œuvre plusieurs optimisations techniques. Il utilise un tuilage bidimensionnel des matrices de paires pour diviser les interactions globales en blocs gérés localement, réduisant la complexité mémoire par dispositif. De plus, il superpose le calcul et la communication : alors qu'un GPU effectue une mise à jour locale, il transfère simultanément des données avec ses voisins pour optimiser l'efficacité. Des primitives de communication distribuées permettent également de gérer l'attention locale sur les atomes sans interagir entre les GPUs, et des adaptations spécifiques du module de multiplication triangulaire assurent que les tensions de paire ne dépassent jamais la capacité mémoire d'une seule carte. Les résultats de cette technologie sont significatifs. Avec 256 GPUs, le système peut traiter jusqu'à 20 000 tokens, permettant à des modèles comme Boltz de prédire des structures bien au-delà des limites précédentes. L'équipe a réussi à replier un complexe contenant 3 605 résidus sur quatre chaînes en moins de cinq minutes, en maintenant tous les contacts inter-subunitaires à longue portée. Cette capacité a déjà été intégrée par des partenaires industriels. Rezo Therapeutics l'utilise pour prédire des interactions protéine-protéine massives jusqu'à 6 500 résidus, Proxima l'a embarquée dans son modèle génératif Neo pour l'inférence sur des assemblages de 4 000 tokens, et Earendil Labs l'a appliquée à son modèle fondamental biomoléculaire pour étendre les longueurs de séquence. Bien que le parallélisme contextuel brise les barrières physiques de la mémoire, la précision biologique dépend toujours de la qualité des données d'entraînement. Les modèles actuels étant souvent entraînés sur de petits fragments, un réajustement avec des coupes plus grandes est nécessaire pour capturer la logique des interactions à longue portée. Pour répondre à la pénurie de données, NVIDIA accélère la création de données synthétiques complexes pour la base de données AlphaFold, en utilisant des outils comme cuEquivariance et TensorRT. Cette évolution ouvre la voie à une découverte plus rapide de nouveaux complexes et de thérapies de prochaine génération.
