NVIDIA optimise l'entraînement MoE des modèles biologiques

NVIDIA a dévoilé des optimisations logicielles permettant d'entraîner efficacement les modèles biologiques fondés sur l'architecture Mixture-of-Experts (MoE). Alors que la montée en puissance des modèles denses devient prohibitivement coûteuse en calcul et en mémoire, les architectures MoE offrent une alternative en distribuant le modèle en plusieurs sous-réseaux, ou experts, et en n'activant qu'un sous-ensemble restreint pour chaque donnée d'entrée. Cette approche gagne en popularité, mais son déploiement efficace exige de surmonter des défis techniques liés à l'orchestration du matériel et à la gestion mémoire. Pour répondre à ces contraintes, l'ingénierie logiciel de NVIDIA, via Transformer Engine et la recette BioNeMo MoE, a introduit trois améliorations majeures. La première consolide l'exécution des experts : au lieu de les traiter séquentiellement, le module GroupedLinear regroupe leurs calculs en un seul appel GPU, réduisant ainsi la surcharge d'ordonnancement. La seconde optimise l'utilisation de la mémoire grâce au format MXFP8. En réduisant la précision de 16 à 8 bits par paramètre, tout en appliquant des facteurs d'échelle par bloc pour préserver la stabilité numérique, cette méthode diminue de moitié l'empreinte mémoire. Elle est directement supportée par les unités Tensor des GPU NVIDIA Blackwell. La troisième avancée fusionne plusieurs étapes computationnelles, dont la quantification et l'activation SwiGLU, au sein d'un noyau GroupedMLP unique. En évitant les conversions répétées entre formats et la sauvegarde de résultats intermédiaires, ce fuseau accélère significativement les passages avant et arrière de l'entraînement. Les benchmarks exécutés sur huit GPU NVIDIA B200 confirment l'efficacité de cette approche, avec un débit de formation jusqu'à 2,21 fois supérieur à une implémentation standard de référence. Ces gains de performance rendent viable l'entraînement à grande échelle de modèles biologiques et génomiques, là où les méthodes traditionnelles peinent à passer à l'échelle. La recette est mise à disposition de la communauté des développeurs, accompagnée de guides de déploiement adaptés au parallélisme d'experts et aux architectures matérielles disponibles. En standardisant ces techniques, NVIDIA positionne Transformer Engine comme un composant essentiel pour l'avenir de l'intelligence artificielle scientifique, où l'efficacité computationnelle déterminera directement la capacité à analyser des données biologiques complexes.