NVIDIA представила BioNeMo Inference Runtime (BioIR): у 2,90 раза вища пропускна здатність згортання Boltz-2 і 58,5 тис. залишків на GPU-годину на 8xH100

Прогнозування біомолекулярної структури перейшло від запусків для окремих цілей до роботи зі списками завдань масштабом протеому. Вузьким місцем більше не є питання, чи може модель згорнути білок. Питання в тому, наскільки швидко вся черга незалежних цілей проходить через аналіз, формування ознак, GPU-інференс і запис результатів. Новий технічний огляд NVIDIA розповідає про BioNeMo Inference Runtime (BioIR) — бібліотеку Python, яка прискорює підтримувані моделі прогнозування структур на GPU NVIDIA, зберігаючи стандартний робочий процес PyTorch. BioIR уже працював у виробничому масштабі. Він забезпечив нещодавнє розширення бази даних AlphaFold, згенерувавши структури білкових комплексів для 4 777 протеомів — близько 31 мільйона кандидатних комплексів, з яких 1,81 мільйона були опубліковані як високонадійні прогнози.

Чи придатна вона для розгортання? Так. BioIR уже доступна як відкритий репозиторій на GitHub із wheel-пакетом, що містить попередньо скомпільовані CUBIN. Для роботи потрібні Python 3.12+, сумісні GPU та драйвер NVIDIA, підготовлений чекпойнт моделі й A3M MSA для кожного ланцюга. nvcc, вихідний код CUDA, CMake або CUDA Toolkit не потрібні.

Що таке BioIR

BioIR оптимізує операції, які стеки інференсу загального призначення не оптимізують повною мірою. До них належать стеки Pairformer і Evoformer, трикутні операції, попарна увага, дифузійні трансформери та модулі на рівні атомів. Моделі залишаються звичайними об’єктами torch.nn.Module. Між чекпойнтом і прямим проходом немає ані побудови рушія, ані етапу експорту, ані окремого артефакту.

Існує 2 способи використання. Процесор повного циклу переміщує InputRequest через аналіз, токенізацію, формування ознак, GPU-інференс і запис у PDB або mmCIF. Пряма інтеграція з PyTorch дає змогу розробникам створювати підтримувану модель або повторно використовувати вибрані оптимізовані модулі у власному коді. У навчальному посібнику демонструється шлях через процесор із Boltz-2 (model_source="boltz-2"). Для кожного білкового ланцюга потрібна A3M MSA. Для вхідних даних із кількома неідентичними білковими ланцюгами приймаються парні або непарні MSA. Шаблони можна надавати вручну, оскільки BioIR не запускає HHsearch або HMMsearch. Процесор підтримує прогнозування структури лігандів, але не прогнозування афінності лігандів.

Три рівні прискорення

BioIR оптимізує роботу на 3 окремих рівнях, кожен із яких націлений на своє вузьке місце:

1. Вибір ядер: підтримувані операції обирають сумісні власні реалізації BioIR, cuEquivariance або резервні реалізації PyTorch на основі конфігурації моделі, GPU, типу даних і форми тензора.
2. Оптимізація модулів: окремий механізм optimize() уможливлює захоплення CUDA Graph для сумісних модулів, зменшуючи накладні витрати на запуск.
3. Масштабування конвеєра: виконавець Ray розміщує 1 повну репліку моделі на кожному видимому GPU вузла та розподіляє між ними незалежні вхідні дані. Етапи на CPU (аналіз, формування ознак, запис) виконуються паралельно зі згортанням на GPU.

Примітка: Ray не розподіляє один прямий прохід між GPU. Режим реплік масштабує списки завдань, а не окремі цілі. Згідно з матрицею підтримки, контекстно-паралельне згортання заплановане, але ще недоступне. Правило розрахунку місткості просте: engine_stage.compute x num_gpus не має перевищувати кількість видимих GPU.

На рівні прямого проходу моделі ранні результати бенчмаркінгу NVIDIA показують середні геометричні прискорення порівняно з базовою реалізацією OSS torch.compile: 1,55x (OpenFold3), 1,78x (Boltz2) і 2,56x (мономер OpenFold2) на H100. Показники для H200 подібні: 1,54x, 1,75x і 2,61x. Вимірювання проводилися на 17 вхідних даних, що охоплювали від 29 до 1 734 залишків.

Бенчмарк: 1 000 людських димерів на 8xH100

Щоб кількісно оцінити доставку результатів від початку до кінця, команда NVIDIA провела порівняльний бенчмарк на 1 000 людських димерних цілей із сумарною довжиною послідовностей меншою за 2 800 залишків. У межах порівняння BioIR-прискорений Boltz-2 зіставляли з реалізацією Boltz-2 із відкритим кодом, скомпільованою за допомогою torch, на 8 GPU H100 по 80 ГБ. В обох випадках використовувалися ідентичні цілі, підготовлені MSA, рецепт інференсу (3 повторні цикли, 200 кроків семплювання, 5 дифузійних семплів) і конфігурація GPU.

Результати:

• BioIR завершила всі 1 000 цілей і забезпечила 58,5 тис. успішно згорнутих залишків на виділений GPU-годину.
• Публічна реалізація забезпечила 20,2 тис. залишків на GPU-годину і вичерпала пам’ять для 29 цілей.
• Підсумок: 2,90-кратне покращення пропускної здатності, нормалізованої за кількістю залишків.

Ці показники стосуються етапу згортання, конкретного набору даних і апаратного забезпечення. Вони не враховують генерацію MSA, виділення ресурсів CPU для попередньої обробки, сховище, передавання даних і повторні спроби. У блозі прямо застерігають від узагальнення цих результатів на всі моделі або набори даних, що підтримуються BioIR.

Енергія для одного мільйона цілей

Якщо лінійно екстраполювати результати бенчмарку на 1 мільйон аналогічних цілей, за оцінками, BioIR потребує 11 МВт·год проти 35 МВт·год для публічної реалізації, якщо використовувати еквіваленти TDP для 8 GPU. За використання еквівалентів максимальної потужності всього вузла оцінка становить 21 МВт·год проти 64 МВт·год. Це оцінки за номінальною потужністю, призначені лише для етапу згортання ІТ-обладнання, а не вимірювання за лічильником; вони не враховують накладні витрати дата-центру, зокрема PUE. Водночас економія від 23 до 43 МВт·год на мільйон цілей є суттєвою цифрою для кампаній масштабом протеому.

Ключові висновки

• BioIR прискорює інференс Boltz-2, OpenFold2 і OpenFold3 на GPU NVIDIA, зберігаючи роботу у звичайному PyTorch.
• Порівняльний бенчмарк на 8xH100: 58,5 тис. проти 20,2 тис. згорнутих залишків на GPU-годину, приріст пропускної здатності у 2,90 раза.
• Режим реплік Ray масштабує незалежні списки завдань; він ніколи не розподіляє 1 прямий прохід між GPU.
• Оціночне енергоспоживання для 1 млн цілей зменшується з 35 МВт·год до 11 МВт·год за еквівалентів TDP для 8 GPU.
• Уже доведено працездатність у масштабі: для розширення бази даних AlphaFold згенеровано 31 млн кандидатних комплексів.

---

Ознайомтеся з технічним блогом, репозиторієм на GitHub, документацією і BioNeMo Agent Toolkit для агентської оркестрації. Також стежте за нами у Twitter і не забудьте приєднатися до нашого субреддиту про машинне навчання зі 150 тис.+ учасників та підписатися на нашу розсилку. Стривайте! Ви є в Telegram? тепер ви також можете приєднатися до нас у Telegram.

Потрібно співпрацювати з нами для просування вашого репозиторію на GitHub, сторінки на Hugging Face, релізу продукту, вебінару тощо? Зв’яжіться з нами