NVIDIA представи BioNeMo Inference Runtime (BioIR): 2,90 пъти по-висока пропускателна способност при нагъване с Boltz-2 и 58,5 хил. остатъка на GPU-час при 8xH100

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

Може ли да бъде внедрен? Да. BioIR е наличен още сега като отворено GitHub хранилище с wheel пакет, съдържащ предварително компилирани CUBIN файлове. За работа по време на изпълнение са необходими Python 3.12 или по-нова версия, съвместими NVIDIA GPU и драйвер, подготвен контролен пункт на модела и A3M MSA за всяка верига. Не са необходими nvcc, CUDA source, 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 monomer) на 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, върху 8xH100 GPU с 80 GB. И двата варианта използваха идентични мишени, предварително подготвени MSA, рецепта за инференс (3 рециклирания, 200 стъпки на семплиране, 5 дифузионни семпли) и конфигурация на GPU.

Резултатите:

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

Тези числа представляват измервания само на етапа на нагъване и са специфични за този набор от данни и този хардуер. Те изключват генерирането на MSA, разпределените CPU ресурси за предварителна обработка, съхранението, трансфера на данни и повторните опити. В блога изрично се предупреждава да не бъдат обобщавани за всички модели или набори от данни, поддържани от BioIR.

Енергия при един милион мишени

При линейно екстраполиране на бенчмарка към 1 милион сравними мишени се изчислява, че BioIR ще се нуждае от 11 MWh спрямо 35 MWh за публичната реализация, използвайки еквиваленти на TDP на 8 GPU. При използване на еквиваленти на максималната мощност на целия възел оценката е 21 MWh спрямо 64 MWh. Това са оценки на базата на номиналната мощност, само за нагъването и за IT оборудването, а не измервания от електромер; освен това те изключват режийните разходи на центъра за данни, като PUE. Въпреки това спестяването от 23 до 43 MWh на всеки милион мишени е съществено число за кампании в мащаба на протеома.

Основни изводи

• BioIR ускорява инференса на Boltz-2, OpenFold2 и OpenFold3 върху NVIDIA GPU, като остава в стандартния PyTorch.
• Сравнителен тест с 8xH100: 58,5 хил. спрямо 20,2 хил. нагънати остатъка на GPU-час — 2,90-кратно увеличение на пропускателната способност.
• Режимът с реплики на Ray мащабира независимите работни списъци; той никога не разделя 1 предно преминаване между GPU.
• Изчислената енергия за 1 млн. мишени спада от 35 MWh на 11 MWh при еквиваленти на TDP на 8 GPU.
• Вече е доказан в мащаб: генерирани са 31 млн. кандидат-комплекса за разширяването на базата данни AlphaFold.

---

Разгледайте техническия блог, GitHub хранилището, документацията и BioNeMo Agent Toolkit за агентна оркестрация. Също така можете да ни последвате в Twitter и не забравяйте да се присъедините към нашия ML SubReddit с над 150 хил. членове и да се абонирате за нашия бюлетин. Чакайте! В Telegram ли сте? сега можете да се присъедините към нас и в Telegram.

Имате нужда от партньор за популяризирането на вашето GitHub хранилище ИЛИ страница в Hugging Face ИЛИ продуктова премиера ИЛИ уебинар и т.н.? Свържете се с нас