NVIDIA ускорила предсказание белков в 2.9 раза: BioIR обрабатывает 58.5K остатков в час на 8xH100

NVIDIA выпустила BioNeMo Inference Runtime (BioIR) — открытую Python-библиотеку, которая ускоряет инференс моделей предсказания биомолекулярных структур на GPU H100/H200. Главная идея: оптимизировать не отдельную модель, а конвейер — от парсинга и фичеризации до вывода PDB-файлов. BioIR уже применялась в продакшене: с её помощью расширили AlphaFold Database, предсказав структуры белковых комплексов для 4777 протеомов (около 31 млн кандидатов, из которых 1.81 млн — с высокой точностью).

Библиотека работает на трёх уровнях: подбор оптимальных ядер (BioIR, cuEquivariance или PyTorch), захват CUDA Graph для снижения оверхэда и распараллеливание через Ray — по одной полной реплике модели на каждый видимый GPU. CPU-этапы (парсинг, фичеризация, запись) перекрываются с GPU-инференсом. При этом модель остаётся обычным torch.nn.Module — не нужны ни экспорт, ни отдельные движки. Доступны два режима: сквозной процессор (InputRequest → PDB) и прямая интеграция в PyTorch.

По бенчмаркам NVIDIA на H100, BioIR даёт геометрическое среднее ускорение 1.55x для OpenFold3, 1.78x для Boltz2 и 2.56x для OpenFold2 monomer относительно torch.compile. На H200 цифры схожи: 1.54x, 1.75x, 2.61x. Тесты проводились на 17 входах с длиной последовательностей от 29 до 1734 остатков. Для Boltz-2 пропускная способность достигла 2.9x при 58.5K остатков на GPU-час на 8xH100.

BioIR доступна на GitHub в виде wheel с предкомпилированными CUBIN. Требования: Python 3.12+, совместимая NVIDIA GPU и драйвер, staged checkpoint и per-chain A3M MSA. Не нужны nvcc, CUDA source, CMake или CUDA toolkit. Пока не поддерживается контекстно-параллельный фолдинг (в планах), но для пакетной обработки протеомов решение уже готово.