Предсказание структуры белков и кофолдинг (co-folding) перешли из разряда узкоспециализированных академических задач в категорию основных рабочих нагрузок, питающих современные исследования в области разработки лекарств и дизайна белков. Сегодня эти процессы все чаще управляются полностью автономно с помощью ИИ-агентов. Однако для того чтобы такой агент мог эффективно выполнять свою работу, каждый этап конвейера — от генерации множественного выравнивания последовательностей (MSA) до инференса кофолдинга и масштабирования на множество GPU — должен быть не только быстрым, но и масштабируемым. Узкое место на любом из этих этапов мгновенно ограничивает общую пропускную способность всей системы.
Скорость и эффективность использования памяти критически важны для ключевых рабочих процессов, таких как виртуальный скрининг и предсказание крупных молекулярных ансамблей. В сценарии виртуального скрининга миллионы или даже миллиарды соединений проверяются на взаимодействие с одной или несколькими белковыми мишенями. Хотя модели кофолдинга часто предоставляют наилучшие предсказанные структуры, их запуск может быть крайне дорогим и ресурсоемким, что делает их практически неприменимыми для скрининга больших библиотек соединений. Именно здесь ускорение от NVIDIA становится ключевым фактором, позволяющим развернуть модели OpenFold3 и аналогичные методы в масштабах крупных библиотек соединений.
Скорость также играет решающую роль при предсказании крупных молекулярных ансамблей, состоящих из нескольких белков и тысяч аминокислотных остатков, поскольку время работы модели кофолдинга масштабируется кубически с количеством остатков. Но еще更大的 вызов представляет собой использование памяти: объем памяти одного GPU часто ограничен, что ставит жесткий потолок на размер комплексов, которые можно предсказать за один проход. Методы, снижающие требования к памяти и распределяющие задачи предсказания между несколькими GPU, открывают качественно новые возможности, которые сегодня просто невозможны.
01Устранение узкого места MSA с помощью GPU-ускорения
Для моделей кофолдинга построение множественного выравнивания последовательностей (MSA) традиционно было этапом, ограниченным производительностью процессора (CPU), который мог доминировать во времени выполнения всего конвейера. Технология MMseqs2-GPU переносит поиск гомологии на GPU NVIDIA, устраняя это узкое место и масштабируясь в зависимости от длины последовательности как на архитектуре NVIDIA Hopper, так и на Blackwell.
Последняя версия GPU-ускоренного решения добавляет специфические оптимизации для архитектур Hopper и Blackwell. Это включает эффективную поддержку поиска по базам данных, размер которых превышает объем памяти GPU, на системах NVIDIA Grace, а также дополнительные ускорения благодаря улучшенным инструкциям DPX Blackwell, доступным начиная с CUDA 13.2. Эти GPU-вклады были интегрированы обратно в основной репозиторий MMseqs2, чтобы вся научная сообщество могло воспользоваться этими ускорениями.
Сервис MSA Search NIM использует MMseqs2-GPU. Согласно статье в журнале Nature Methods, этот подход обеспечивает ускорение выравнивания до 177 раз по сравнению с классическим CPU-инструментом JackHMMER на одном GPU L40S. В наших бенчмарках этот этап масштабируется плавно даже для последовательностей длиной более 10 000 токенов на GPU H100 и B300 (см. рисунок 2). MSA Search NIM можно вызывать напрямую, размещать самостоятельно или использовать как инструмент в агентном рабочем процессе.
Ниже приведен пример того, как добавить навык MSA Search NIM в вашего агента и выполнить запрос:
# Add the MSA Search NIM skill to your agent (browse all: add --list)
npx skills add NVIDIA-BioNeMo/bionemo-agent-toolkit --skill msa-search-nim --agent claude-code
# Use hosted API on build.nvidia.com (nothing to download)
# You can also use the skill to download the NIM container and provide self-hosted API endpoint. We don't cover that in that tutorial.
export NVIDIA_API_KEY=
# Just prompt the agent:
# You have to download a sample target.fasta file. You can prompt the agent to download it for you or point to an already existing file.
"Build an MSA for the sequence in target.fasta with the MSA Search NIM." 
02Кофолдинг со скоростью SOTA с помощью cuEquivariance и OpenFold3 NIM
Библиотека cuEquivariance от CUDA-X представляет собой набор геометрических примитивов обучения для атомистического моделирования. Она предоставляет ускоренные версии ядер Triangle Attention, Triangle Multiplication и Attention Pair Bias, которые доминируют в вычислениях при кофолдинге. На GPU B300 cuEquivariance сокращает время задержки (latency) до 3 раз, как показано в таблице 1.
Ядра cuEquivariance интегрированы непосредственно в модели с открытым исходным кодом, такие как OpenFold3 (предоставляются как опциональная зависимость), OpenFold2, RosettaFold3, Protenix и Boltz. Поскольку ускорения интегрированы в эти модели с открытым исходным кодом, исследователь получает прирост скорости автоматически, просто запуская модель, которую он уже использует, на GPU NVIDIA. Кроме того, ядра cuEquivariance расширяют максимальную длину последовательности до ~5,9 тысяч токенов, тогда как PyTorch выдает ошибку нехватки памяти (OOM) при длинах свыше ~1,5–2,5 тысяч токенов.
На базе cuEquivariance сервис OpenFold3 NIM применяет дополнительные оптимизации инференса, которые суммируют эффект ускорения. Это позволяет достигать длин последовательностей до ~6 400 токенов на одном GPU B300. Эти дополнительные ускорения доставляются через NIM; для получения наилучших результатов «из коробки» разработчики могут вызывать эндпоинт NIM напрямую или интегрировать его в агентный рабочий процесс.
Пример вызова OpenFold3 NIM через агента:
# Add the OpenFold3 NIM skill to your agent
npx skills add NVIDIA-BioNeMo/bionemo-agent-toolkit --skill openfold3-nim --agent claude-code
# Hosted API on build.nvidia.com
# You can also use the skill to download the NIM container and provide self-hosted API endpoint. We don't cover that in that tutorial.
export NVIDIA_API_KEY=
# Prompt the agent:
# You have to download a sample target.fasta file. You can prompt the agent to download it for you or point to an already existing file.
"Fold target.fasta with OpenFold3 using the MSA from the previous step; return the ranked structures with confidence scores." 03Масштабирование за пределы одного GPU с помощью Fold-CP
Память одного GPU исторически ограничивала модели кофолдинга несколькими тысячами остатков. На GPU NVIDIA B300 (Blackwell Ultra) больший объем HBM и эффективность поколения Blackwell, в сочетании с cuEquivariance и продвинутыми оптимизациями инференса, значительно повышают этот потолок без необходимости изменения самой модели.
Во многих случаях ограничения одного устройства недостаточно. Технология Fold-CP вводит новый метод параллелизма, при котором требования к памяти на каждом устройстве масштабируются как O(N²/P), где N — количество токенов, а P — количество GPU. Это позволяет достигать длины в 32 000 токенов на 64 GPU B300 с использованием модели Boltz-2, что представляет собой скачок в 12 раз по сравнению с пределом одного GPU.
Чтобы попробовать Fold-CP, достаточно направить вашего агента на кодовую базу Boltz-CP и попросить выполнить много-GPU инференс. Ниже приведен пример запуска контекстно-параллельного инференса на 4 GPU:
# Context-parallel inference across 4 GPUs with Fold-CP (boltz-cp)
git clone https://github.com/NVIDIA-Digital-Bio/boltz-cp && cd boltz-cp
# Install dependencies and then run the command below
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
04Ускорение всего конвейера кофолдинга
Производительность предсказания структуры теперь является системной задачей «от конца к концу» (end-to-end). Для OpenFold3 практический рабочий процесс охватывает генерацию MSA, инференс кофолдинга, развертывание и ограничения памяти, которые определяют, насколько большой биологический ансамбль можно смоделировать.
NVIDIA ускоряет каждый слой этого конвейера: MSA Search NIM ускоряет поиск гомологии в 177 раз, cuEquivariance и OpenFold3 NIM снижают задержку инференса до 4 раз на GPU Blackwell, а Fold-CP демонстрирует, как контекстный параллелизм может расширить возможности кофолдинга за пределы одного GPU, позволяя моделировать комплексы длиной до 32 000 токенов на 64 GPU NVIDIA B300.
Вместе эти инструменты делают предсказание структуры быстрее, более масштабируемым и проще интегрируемым в агентные рабочие процессы открытия, помогая исследователям переходить от предсказаний моделей к более крупным и полезным биологическим системам.
05Что это значит на практике
Эти улучшения открывают классы проблем структурной биологии, которые ранее были недоступны. В виртуальном скрининге более быстрый инференс кофолдинга означает, что методы, основанные на структуре, которые исторически зарезервировались для финальных стадий кампании по разработке лекарств, теперь могут применяться на гораздо более ранних этапах и в масштабах гораздо больших библиотек, улучшая качество и разнообразие кандидатов, продвигающихся через конвейер.
Для крупных биомолекулярных ансамблей комбинация расширенной емкости одного GPU на B300 и фреймворка контекстного параллелизма Fold-CP меняет то, что можно моделировать: комплексы масштаба рибосомы, сплайсосомы или крупных сигнальных ансамблей были структурно недоступны для моделей кофолдинга, и эти ускорения начинают менять это положение дел. Чтобы сделать это конкретным: сворачивание комплекса из ~10 000 остатков, примерно масштаба бактериальной рибосомы, было бы prohibitively expensive или просто недоступно на одном GPU; на B300 с Fold-CP такие предсказания становятся выполнимыми в рамках много-GPU узла.
Это качественные сдвиги в вопросах, которые может задавать структурная биология вычислительными методами, а не просто улучшения пропускной способности. Сделав эти инструменты доступными через интеграции с открытым исходным кодом и агентные API, мы ускорим темп, с которым вычислительные предсказания превращаются в биологические инсайты и, в конечном итоге, в новые лекарства.
06Заключение
Интеграция GPU-ускорения на всех этапах конвейера предсказания структуры белков представляет собой значительный прорыв. От устранения узких мест в поиске гомологии до расширения возможностей моделирования гигантских молекулярных комплексов, NVIDIA BioNeMo Agent Toolkit предоставляет исследователям мощный набор инструментов. Это не просто ускорение существующих процессов, но и открытие новых горизонтов в понимании биологических систем и разработке лекарств.
Авторы статьи: Kyle Tretina, Roy Tal, Zoey Zhang, Emine Kucukbenli, Jared Wilber и Mohammed AlQuraishi. Особая благодарность команде NVIDIA и сообществам OpenFold и OMSF за сотрудничество и вклад в разработку этих технологий.
Источник: NVIDIA Developer ↗