BioNeMo Inference Runtime: Ускорение предсказания структур белков в 2.9 раза

Предсказание структур белков (protein structure prediction) перешло от единичных экспериментов к протеомному масштабу. Для обработки миллионов мишеней критически важна не только точность модели, но и эффективность использования GPU. NVIDIA BioNeMo Inference Runtime (BioIR) решает эту задачу, предоставляя оптимизированный движок для инференса биомолекулярных моделей на базе PyTorch, сохраняя совместимость с привычным кодом разработчиков.

Ключевое преимущество BioIR — трехуровневая оптимизация: выбор оптимальных ядер (kernels), захват CUDA Graphs для модулей и масштабирование через реплики Ray. В бенчмарке на 8xH100 ускорение модели Boltz-2 составило 2.90x по сравнению с открытой torch-compiled реализацией, а энергопотребление снизилось с 35 MWh до 11 MWh на миллион предсказаний.

01Архитектура и режимы работы

BioIR предлагает два способа интеграции:

1. End-to-End Processor: Полный пайплайн от парсинга входных данных (A3M, последовательности) до генерации PDB/mmCIF файлов. Включает токенизацию, генерацию фич, инференс на GPU и запись.
2. Direct PyTorch Integration: Прямое использование оптимизированных torch.nn.Module в кастомных скриптах.

Для пакетной обработки больших списков (worklists) используется бэкенд Ray. Он размещает полную копию модели на каждой доступной GPU узла, позволяя перекрывать CPU-стадии (парсинг, генерация фич) с GPU-стадией (фолдинг).

02Бенчмарки: Пропускная способность и Энергоэффективность

Сравнение проводилось на наборе из 1,000 димерных мишеней человека. Результаты демонстрируют существенный прирост производительности:

Метрика model_inference_time в BioIR измеряет только время прямого прохода модели на GPU (CUDA-synchronized), исключая накладные расходы на CPU-стадии, что дает более честную картину вычислительной эффективности.

03Практическая реализация: От валидации к масштабу

Для начала работы требуется Python 3.12+, совместимая NVIDIA GPU и wheel-пакет BioIR. Предкомпилированные CUBIN-файлы исключают необходимость в nvcc или сборке из исходников.

Шаг 1: Подготовка входных данных

Каждая цепь белка требует MSA (Multiple Sequence Alignment) в формате A3M. BioIR не запускает HHsearch/HMMsearch автоматически, поэтому MSA должны быть подготовлены заранее. Поддерживаются как парные, так и непарные MSA для комплексов.

from bionemo_ir.data.schemas import InputRequest, MSARecord, Polymer

request = InputRequest(
    input_id="demo",
    polymers=[
        Polymer(
            polymer_type="protein",
            chain_id=["A"],
            sequence="GSHMSL...",  # Замените на реальную последовательность
            msas=[MSARecord(path="msa.a3m", format="a3m")],
            paired_msas=[],
            templates=None,
        )
    ],
)
rows = [{"record": request, "__record_id": request["input_id"]}]

Шаг 2: Валидация через Serial Processor

Перед запуском масштабного пайплайна рекомендуется проверить конфигурацию на одном примере. Serial executor выполняет стадии последовательно.

from bionemo_ir.pipeline.processor.engine_proc import (
    EngineProcessorConfig,
    build_processor,
)
from bionemo_ir.pipeline.stages.configs import (
    FeatureGeneratorStageConfig,
    WriterStageConfig,
)

serial_config = EngineProcessorConfig(
    model_source="boltz-2",
    executor_backend=None,  # Serial mode
    runtime_args={
        "recycling_steps": 3,
        "num_sampling_steps": 50,
        "diffusion_samples": 1,
    },
    feature_generator_stage=FeatureGeneratorStageConfig(
        init_context={"random_seed": 42},
    ),
    writer_stage=WriterStageConfig(
        output_path="output/serial",
        format="cif",
    ),
    engine_kwargs={"profile_inference": True},
)

serial_processor = build_processor(serial_config)
serial_outputs = serial_processor(rows)

for row in serial_outputs:
    import json
    scores = json.loads(row["scores"])
    print(f"Path: {row['output_path']}")
    print(f"Inference Time: {row['model_inference_time']}")
    print(f"Confidence: {scores.get('confidence_score')}")

Шаг 3: Масштабирование через Ray

Для обработки тысяч мишеней используется Ray. По умолчанию создается одна реплика модели на каждую видимую GPU. Можно явно задать распределение ресурсов CPU/GPU для каждой стадии.

import ray
from bionemo_ir.pipeline.processor.engine_proc import (
    EngineProcessorConfig,
    build_processor,
)
from bionemo_ir.pipeline.stages.configs import (
    EngineStageConfig,
    FeatureGeneratorStageConfig,
    ParallelismMode,
    ParserStageConfig,
    TokenizerStageConfig,
    WriterStageConfig,
)

# Конфигурация для 4 GPU с явным распределением ресурсов
ray_config = EngineProcessorConfig(
    model_source="boltz-2",
    executor_backend="ray",
    parser_stage=ParserStageConfig(compute=4),
    tokenizer_stage=TokenizerStageConfig(compute=4, num_cpus=2),
    feature_generator_stage=FeatureGeneratorStageConfig(
        compute=8,
        num_cpus=4,
        init_context={"random_seed": 42},
    ),
    engine_stage=EngineStageConfig(
        parallelism_mode=ParallelismMode.REPLICA,
        compute=4,
        num_gpus=1.0,  # 1 GPU на реплику
        num_cpus=4,
    ),
    writer_stage=WriterStageConfig(
        compute=4,
        output_path="output/ray",
        format="cif",
    ),
)

ray_processor = build_processor(ray_config)
dataset = ray.data.from_items(rows)
ray_outputs = list(ray_processor(dataset).materialize().iter_rows())

04Кому подойдёт / что запустится

• Исследователи и биоинформатики: Для протеомного масштаба (тысячи/миллионы структур). BioIR уже использовался для расширения AlphaFold Database (AFDB), обработав ~31 млн кандидатных комплексов.
• Разработчики ML-инфраструктуры: Для интеграции в пайплайны drug discovery. Поддерживается прямое подключение модулей PyTorch.
• Железо: Требуется NVIDIA GPU с поддержкой CUDA. Для максимального ускорения рекомендуются H100/A100. CPU-часть (парсинг MSA) также требует ресурсов, особенно при высокой параллельности.
• Модели: На момент выхода поддерживается Boltz-2. Другие модели добавляются через поддержку матрицы совместимости.

Для начала работы клонируйте репозиторий BioNeMo Inference Runtime и изучите API reference. Для агентных оркестраций стоит также рассмотреть BioNeMo Agent Toolkit.

Источник: NVIDIA Developer ↗