NVIDIA detalla el tiempo de ejecución de inferencia BioNeMo (BioIR): rendimiento de plegado Boltz-2 2,90 veces mayor y 58,5 000 residuos por hora de GPU en 8xH100

La predicción de estructuras biomoleculares ha pasado de ejecuciones de un solo objetivo a listas de trabajo a escala de proteoma. El cuello de botella ya no es si un modelo puede plegar una proteína. Es la rapidez con la que se mueve una cola completa de objetivos independientes a través del análisis, la caracterización, la inferencia de GPU y la escritura de resultados. La nueva inmersión técnica profunda de NVIDIA recorre BioNeMo Inference Runtime (BioIR), una biblioteca de Python que acelera los modelos de predicción de estructuras admitidos en las GPU de NVIDIA mientras mantiene el flujo de trabajo estándar de PyTorch. BioIR ya se ha ejecutado a escala de producción. Impulsó la reciente expansión de la base de datos AlphaFold, generando estructuras de complejos proteicos en 4.777 proteomas, alrededor de 31 millones de complejos candidatos, de los cuales 1,81 millones se publicaron como predicciones de alta confianza.

¿Es desplegable? Sí. BioIR ya está disponible como un repositorio abierto de GitHub con una rueda que contiene CUBIN precompilados. El uso del tiempo de ejecución necesita Python 3.12+, una GPU y un controlador NVIDIA compatibles, un punto de control de modelo por etapas y MSA A3M por cadena. No requiere nvcc, fuente CUDA, CMake ni el kit de herramientas CUDA.

¿Qué es BioIR?

BioIR apunta a las operaciones que las pilas de inferencia de propósito general no optimizan completamente. Estos incluyen pilas Pairformer y Evoformer, operaciones triangulares, atención por pares, transformadores de difusión y módulos a nivel atómico. Los modelos siguen siendo objetos de módulo de antorcha ordinarios. No hay construcción de motor, paso de exportación ni artefacto separado entre un punto de control y un pase hacia adelante.

Hay 2 maneras de usarlo. El procesador de un extremo a otro mueve una solicitud de entrada mediante análisis, tokenización, generación de funciones, inferencia de GPU y escritura PDB o mmCIF. La integración directa de PyTorch permite a los desarrolladores construir un modelo compatible o reutilizar módulos optimizados seleccionados dentro de código personalizado. El tutorial demuestra la ruta del procesador con Boltz-2 (model_source=”boltz-2″). Cada cadena de proteínas requiere un MSA A3M. Se aceptan MSA emparejadas o no emparejadas para entradas con múltiples cadenas de proteínas no idénticas. Las plantillas se pueden proporcionar manualmente porque BioIR no ejecuta HHsearch o HMMsearch. El procesador admite la predicción de la estructura del ligando, pero no la predicción de la afinidad del ligando.

Tres capas de aceleración

BioIR se optimiza en 3 capas distintas, cada una de las cuales apunta a un cuello de botella diferente:

Selección de kernel: las operaciones admitidas seleccionan implementaciones de respaldo personalizadas de BioIR, cuEquivariance o PyTorch compatibles según la configuración del modelo, GPU, tipo de datos y forma del tensor. Optimización del módulo: un mecanismo de optimización independiente permite la captura de CUDA Graph para módulos compatibles, lo que reduce la sobrecarga de lanzamiento. Escalado de canalización: un ejecutor de Ray coloca 1 réplica de modelo completa en cada GPU visible en un nodo y distribuye entradas independientes entre ellas. Las etapas de la CPU (análisis, caracterización, escritura) se superponen con el plegado de la GPU.

Nota: Ray no divide un solo pase hacia adelante entre las GPU. El modo réplica escala listas de trabajo, no objetivos individuales. Según la matriz de soporte, el plegado paralelo al contexto está previsto, pero aún no está disponible. La regla de capacidad es simple: Engine_stage.compute x num_gpus no debe exceder las GPU visibles.

A nivel de modelo avanzado, las primeras evaluaciones comparativas de NVIDIA informan aceleraciones de media geométrica sobre una línea base de compilación de antorcha OSS de 1,55x (OpenFold3), 1,78x (Boltz2) y 2,56x (monómero OpenFold2) en H100. Los números del H200 son similares a 1,54x, 1,75x y 2,61x. Estos se midieron en 17 entradas que abarcaban entre 29 y 1734 residuos.