Проблема масштабируемости в структурной биологии
Предсказание структур белков перешло от единичных запусков к протеомным задачам. Главная проблема теперь — не способность модели свернуть белок, а скорость обработки очередей независимых целей. NVIDIA выпустила BioNeMo Inference Runtime (BioIR) — Python-библиотеку, оптимизирующую работу моделей на GPU без потери совместимости с PyTorch. BioIR уже доказал свою эффективность: он обеспечил генерацию 31 миллиона кандидатных комплексов для расширения базы AlphaFold (4 777 протеомов, 1.81 млн высоконадежных предсказаний).
Архитектура ускорения: три уровня оптимизации
BioIR не требует сборки движков или экспорта моделей. Модели остаются обычными torch.nn.Module объектами. Ускорение достигается за счет трех слоев:
- Выбор ядер (Kernel selection): Автоматический выбор между кастомными ядрами BioIR, cuEquivariance и PyTorch-фоллбэками на основе конфигурации модели и типа данных.
- Оптимизация модулей: Использование CUDA Graph capture через механизм
optimize()для снижения накладных расходов на запуск. - Масштабирование пайплайна: Executor на базе Ray распределяет реплики моделей по GPU узла, позволяя CPU-этапам (парсинг, запись) работать параллельно с GPU-сворачиванием.
Бенчмарк: 8xH100 и Boltz-2
Для оценки производительности NVIDIA провела тест на 1 000 димеров человека (общая длина последовательностей < 2 800 остатков). Сравнивались BioIR-accelerated Boltz-2 и open-source Boltz-2 с torch.compile на кластере из 8 GPU H100 80GB. Параметры инференса были идентичны: 3 рецикла, 200 шагов сэмплирования, 5 диффузионных сэмплов.
| Метрика | BioIR (Boltz-2) | Open Source (torch.compile) |
|---|---|---|
| Пропускная способность (остатков/GPU-час) | 58.5K | 20.2K |
| Ускорение (Throughput) | 2.90x | |
| Ошибки памяти (OOM) | 0 | 29 целей |
Энергоэффективность и масштабирование
Экстраполяция результатов на 1 миллион целевых структур показывает значительную экономию энергии. При использовании эквивалентов TDP 8-GPU ноды, BioIR требует 11 MWh против 35 MWh у публичной реализации. Даже с учетом максимального энергопотребления узлов разница составляет 21 MWh против 64 MWh. Это экономия от 23 до 43 MWh на миллион предсказаний, что критично для протеомных кампаний.
Технические требования и доступность
BioIR доступен как open-source репозиторий на GitHub. Для работы требуется Python 3.12+, совместимый драйвер NVIDIA GPU и предварительно скомпилированные CUBIN-файлы. Сборка из исходников (nvcc, CUDA toolkit) не требуется. Поддерживаются модели Boltz-2, OpenFold2 и OpenFold3. Важно отметить: Ray-режим масштабирует рабочие списки (worklists), а не разделяет один инференс между GPU. Поддержка контекстного параллелизма (context-parallel) пока находится в планах.
Источник: MarkTechPost ↗