NVIDIA presento un conjunto de herramientas que aceleran de extremo a extremo las tuberias de prediccion de estructuras biomoleculares, integradas en el BioNeMo Agent Toolkit. La propuesta cubre tres cuellos de botella clasicos del co-folding (plegamiento conjunto de proteinas): la generacion del alineamiento multiple de secuencias (MSA), la inferencia del modelo y el escalado a varias GPU.
Los numeros que la compania pone sobre la mesa son concretos:
- MMseqs2-GPU logra generar el MSA hasta 177 veces mas rapido que JackHMMER en CPU, sobre arquitecturas NVIDIA Hopper y Blackwell.
- cuEquivariance recorta hasta 3 veces el tiempo de la pasada hacia adelante de OpenFold3 y extiende el limite de secuencia a unos 5.900 tokens.
- El OpenFold3 NIM alcanza hasta 6.400 tokens en una sola GPU B300.
- Fold-CP habilita inferencia con paralelismo de contexto y permite modelar ensamblajes de hasta 32.000 tokens en 64 GPU B300, cerca de un salto de 12 veces frente al limite de una sola GPU.
¿Por que importa acelerar el co-folding?
La prediccion de estructuras y el co-folding con modelos como OpenFold3 ya son cargas de trabajo masivas y de uso corriente que impulsan el descubrimiento de farmacos y el diseno de proteinas. Cada vez mas, agentes de IA las ejecutan de punta a punta, y para que ese agente corra bien la tuberia, cada paso debe ser rapido y escalable. Un cuello de botella en cualquier etapa limita el rendimiento total.
La velocidad y la eficiencia de memoria son criticas en flujos como el virtual screening, donde se evaluan millones o miles de millones de compuestos contra uno o pocos blancos proteicos. Los modelos de co-folding suelen dar las mejores estructuras predichas, pero resultan caros de ejecutar, lo que hasta ahora los volvia poco practicos para tamizados a gran escala.
El otro gran desafio es el tamano de los ensamblajes. El tiempo de ejecucion de estos modelos crece de forma cubica respecto al numero de residuos, y la memoria de una sola GPU impone un techo duro al tamano del complejo que se puede predecir de una vez.
Eliminar el cuello de botella del MSA
Para los modelos de co-folding, construir el MSA fue tradicionalmente un paso limitado por la CPU que podia dominar el tiempo total. MMseqs2-GPU traslada la busqueda de homologia a las GPU de NVIDIA y escala con la longitud de la secuencia en arquitecturas Hopper y Blackwell.
La ultima version acelerada por GPU suma optimizaciones especificas para Hopper y Blackwell, incluido el soporte eficiente para bases de datos mayores que la memoria de la GPU en sistemas Grace, y aceleraciones adicionales gracias a las instrucciones DPX mejoradas de Blackwell disponibles desde CUDA 13.2. Estas contribuciones se integraron de vuelta al repositorio principal de MMseqs2, de modo que toda la comunidad se beneficie.
El MSA Search NIM usa MMseqs2-GPU, cuyo articulo en Nature Methods reporta hasta 177 veces mas velocidad de alineamiento que JackHMMER en CPU sobre una sola L40S. En los benchmarks de NVIDIA, la etapa escala sin sobresaltos mas alla de los 10.000 tokens en GPU H100 y B300.
Plegar a velocidad de estado del arte
cuEquivariance es una biblioteca CUDA-X de primitivas de aprendizaje geometrico para el modelado atomistico. Ofrece versiones aceleradas de los kernels de Triangle Attention, Triangle Multiplication y Attention Pair Bias que dominan el co-folding. En B300 recorta la latencia hasta 3 veces.
Sus kernels se integran directamente en modelos de codigo abierto como OpenFold3 (como dependencia opcional), OpenFold2, RosettaFold3, Protenix y Boltz. Como las aceleraciones estan incorporadas rio arriba, el investigador obtiene la mejora de forma automatica con solo correr en una GPU NVIDIA el modelo que ya usa. Ademas, cuEquivariance extiende la longitud maxima de secuencia a unos 5.900 tokens, mientras que PyTorch se queda sin memoria mas alla de los 1.500 a 2.500 tokens.
Sobre cuEquivariance, el OpenFold3 NIM aplica optimizaciones de inferencia adicionales que suman al resultado y alcanzan longitudes de hasta 6.400 tokens en una sola B300.
Escalar mas alla de una GPU con Fold-CP
La memoria de una sola GPU limito historicamente los modelos de co-folding a unos pocos miles de residuos. Fold-CP introduce una tecnica de paralelizacion en la que la memoria por dispositivo escala como O(N2/P), donde N es el numero de tokens y P la cantidad de GPU. Asi alcanza 32.000 tokens en 64 GPU B300 con el modelo Boltz-2, cerca de un salto de 12 veces sobre el limite de una sola GPU.
Para probar Fold-CP, basta apuntar el agente al repositorio de Boltz-CP y pedirle que ejecute inferencia multi-GPU:
# Inferencia con paralelismo de contexto en 4 GPU con Fold-CP (boltz-cp)
git clone https://github.com/NVIDIA-Digital-Bio/boltz-cp && cd boltz-cp
# Instalar dependencias y luego ejecutar
torchrun --nnodes 1 --nproc_per_node 4 \
src/boltz/distributed/main.py predict /path/to/preprocessed_data \
--out_dir ./predictions \
--size_dp 1 --size_cp 4 \
--recycling_steps 3 --sampling_steps 200 --diffusion_samples 5¿Que problemas nuevos habilita esto?
Las mejoras abren clases de problemas de biologia estructural antes inalcanzables. En el virtual screening, una inferencia de co-folding mas rapida permite aplicar metodos basados en estructura en etapas mucho mas tempranas y sobre bibliotecas de compuestos mucho mas grandes, lo que mejora la calidad y diversidad de los candidatos que avanzan.
Para los grandes ensamblajes biomoleculares, la combinacion de mayor capacidad en una sola B300 y el marco Fold-CP cambia lo que es modelable: complejos a escala del ribosoma, el espliceosoma o grandes ensamblajes de senalizacion, que eran intratables para estos modelos. Plegar un complejo de unos 10.000 residuos, cerca del tamano del ribosoma bacteriano, habria sido prohibitivamente caro o directamente imposible en una sola GPU; con B300 y Fold-CP, esas predicciones se vuelven abordables en un nodo multi-GPU.
Para los laboratorios de LatAm que trabajan con presupuestos acotados de computo, la clave es que estas aceleraciones llegan rio arriba en modelos de codigo abierto: el investigador que ya corre OpenFold3 o Boltz sobre una GPU NVIDIA hereda la mejora sin reescribir su codigo, lo que baja la barrera para sumarse al diseno computacional de farmacos.
