Прорыв в эффективности: 2.9x ускорение
NVIDIA BioNeMo Inference Runtime (BioIR) — это специализированное программное обеспечение, предназначенное для ускорения инференса моделей предсказания структур биомолекул на GPU NVIDIA. Ключевое отличие BioIR от стандартных решений заключается в трехуровневой оптимизации: выбор оптимальных ядер (kernels), использование CUDA Graphs для захвата вычислительных графов и масштабирование через реплики Ray.
В бенчмарке на 1000 человеческих димерных мишенях при использовании 8 GPU H100, модель Boltz-2, ускоренная BioIR, достигла показателя 58.5K успешно сложенных остатков на GPU-час. Для сравнения, открытая реализация с torch-компиляцией показала лишь 20.2K остатков. Это дает прирост пропускной способности в 2.90 раза.
Энергоэффективность и реальные кейсы
Помимо скорости, BioIR демонстрирует значительное снижение энергопотребления. Для складывания миллиона аналогичных мишеней требуется всего 11 MWh энергии (в эквиваленте rated power), тогда как публичная реализация потребляет около 35 MWh. Это критически важно для протеомного масштаба, где обрабатываются миллионы структур.
Технология уже внедрена в производство: именно BioIR ускорил генерацию структур белковых комплексов для 4,777 протеом в рамках расширения AlphaFold Database (AFDB). Было обработано около 31 миллиона кандидатных комплексов, из которых 1.81 миллиона были опубликованы как высококонфиденциальные предсказания.
Технические детали и интеграция
BioIR поддерживает два режима работы:
- End-to-end процессор: автоматизирует весь пайплайн от парсинга входных данных и токенизации до генерации признаков, GPU-инференса и записи результатов в PDB или mmCIF.
- Прямая интеграция PyTorch: позволяет использовать оптимизированные модули (torch.nn.Module) в пользовательском коде, сохраняя знакомый рабочий процесс разработчиков.
Для масштабирования на больших батчах независимых входных данных используется бэкенд Ray, который размещает полную реплику модели на каждом GPU узла, позволяя перекрывать этапы обработки на CPU с этапами сворачивания на GPU.
Сравнение производительности
| Метрика | BioIR (Boltz-2) | Open Source (Torch-compiled) |
|---|---|---|
| Пропускная способность (остатков/GPU-час) | 58,500 | 20,200 |
| Коэффициент ускорения | 2.90x | 1.0x (база) |
| Энергозатраты (на 1 млн мишеней) | 11 MWh | 35 MWh |
| Аппаратная база бенчмарка | 8x NVIDIA H100 | |
Как начать работу
Для использования BioIR требуется Python 3.12+, совместимые GPU NVIDIA и предварительно скомпилированные CUBIN-файлы (компиляция из исходного кода не требуется). Поддерживаются модели, такие как Boltz-2, с возможностью подачи входных данных в формате A3M MSA. Разработчики могут интегрировать BioIR как через готовый процессор, так и через API для кастомных пайплайнов.
Источник: NVIDIA dev blog ↗