אנבידיה משיקה את BioIR: פי 2.9 תפוקה ב-Boltz-2 ו-58.5 אלף רזידואים לשעת-GPU על 8xH100

חיזוי מבנה ביומולקולרי עבר מעבודה על חלבון בודד לריצות בקנה מידה פרוטאומי, וצוואר הבקבוק כבר לא טמון ביכולת לקפל חלבון יחיד אלא בקצב שבו תור שלם של מטרות עצמאיות עובר דרך פירסינג, פיצ'ריזציה, אינפרנס על GPU וכתיבת פלט. NVIDIA פרסמה השבוע פירוט טכני של BioNeMo Inference Runtime (BioIR), ספריית פייתון שמאיצה מודלים נתמכים לחיזוי מבנה על GPU של החברה תוך שמירה על זרימת עבודה סטנדרטית של PyTorch. הכלי כבר הוכח בפרודקשן: הוא הניע את ההרחבה האחרונה של AlphaFold Database, שיצרה מבני קומפלקסים על פני 4,777 פרוטאומים, כ-31 מיליון קומפלקסים מועמדים, מתוכם שוחררו 1.81 מיליון תחזיות ברמת ביטחון גבוהה.

זמין עכשיו, בלי CUDA Toolkit

BioIR זמין כבר עכשיו כמאגר גיטהאב פתוח עם wheel שמכיל CUBINs מקומפלים מראש. הריצה דורשת פייתון 3.12 ומעלה, GPU תואם של אנבידיה עם דרייבר מתאים, צ'קפוינט מודל מוכן (staged) ו-MSA בפורמט A3M לכל שרשרת. מה שלא צריך זה nvcc, קוד מקור של CUDA, CMake או ערכת הכלים המלאה של CUDA, נקודה משמעותית למעבדות שרוצות להריץ בלי לנהל סביבת בנייה מורכבת.

שתי דרכי שימוש, מודל אחד נשאר torch.nn.Module

הספרייה מציעה שני מסלולים: מעבד מקצה-לקצה שמעביר InputRequest דרך פירסינג, טוקניזציה, יצירת פיצ'רים, אינפרנס GPU וכתיבת PDB או mmCIF; ואינטגרציה ישירה של PyTorch שמאפשרת לבנות מודל נתמך או לשלב מודולים מואצים נבחרים בתוך קוד מותאם. המודלים נשארים אובייקטי torch.nn.Module רגילים, אין שלב בניית מנוע, ייצוא או ארטיפקט נפרד בין הצ'קפוינט ל-forward pass. המדריך מדגים את נתיב המעבד עם Boltz-2 (model_source="boltz-2"). כל שרשרת חלבון דורשת MSA בפורמט A3M; MSA מזווגים או לא-מזווגים מתקבלים עבור קלטים עם מספר שרשראות חלבון לא-זהות. טמפלטים אפשר לספק ידנית כי BioIR לא מריץ HHsearch או HMMsearch. המעבד תומך בחיזוי מבנה ליגנד אך לא בחיזוי אפיניות ליגנד.

שלוש שכבות האצה, ומה עדיין חסר

BioIR מייעל בשלוש שכבות נפרדות, שכל אחת מכוונת לצוואר בקבוק אחר. בחירת קרנל (kernel selection): אופרציות נתמכות בוחרות בין מימושים מותאמים של BioIR, cuEquivariance או נפילה ל-PyTorch לפי קונפיגורציית המודל, ה-GPU, סוג הנתונים וצורת הטנזור. אופטימיזציית מודול: מנגנון optimize נפרד מאפשר לכידת CUDA Graph למודולים תואמים, חותך אוברהד של השקות. סקיילינג פייפליין: מריץ Ray מציב רפליקה מלאה אחת של המודל על כל GPU גלוי בצומת ומחלק קלטים עצמאיים ביניהן; שלבי CPU (פירסינג, פיצ'ריזציה, כתיבה) חופפים לאינפרנס GPU. חשוב: Ray לא מפצל forward pass בודד על פני מספר GPU, מצב רפליקה מגדיל תפוקת worklists, לא מטרות בודדות. לפי מטריצת התמיכה, קיפול context-parallel מתוכנן אך עדיין לא זמין. כלל הקיבולת פשוט: engine_stage.compute כפול num_gpus לא יעלה על מספר ה-GPU הגלויים.

בנצ'מרק ברמת המודל: עד פי 2.6 על torch.compile

ברמת ה-forward pass של המודל, הבנצ'מרקים המוקדמים של אנבידיה מדווחים על האצה גיאומטרית ממוצעת מול בסיס OSS של torch.compile: פי 1.55 ב-OpenFold3, פי 1.78 ב-Boltz2, ופי 2.56 ב-OpenFold2 מונומר, כולם על H100. המספרים על H200 דומים: פי 1.54, פי 1.75 ופי 2.61 בהתאמה. המדידות בוצעו על 17 קלטים באורכים של 29 עד 1,734 רזידואים.

בנצ'מרק קצה-לקצה: 1,000 דימרים אנושיים על 8xH100

כדי לכמת אספקה מקצה-לקצה, הצוות הריץ בנצ'מרק תואם על 1,000 מטרות דימר אנושיות עם אורך רצף משולב מתחת ל-2,800 רזידואים. ההשוואה העמידה Boltz-2 מואץ ב-BioIR מול מימוש Boltz-2 פתוח מקומפל ב-torch על 8 GPU מדגם H100 80GB. שני הצדדים השתמשו באותן מטרות, MSA מוכנים, מתכון אינפרנס זהה (3 מחזורי מחזור, 200 צעדי דגימה, 5 דגימות דיפוזיה) וקונפיגורציית GPU זהה. התוצאה: תפוקת קיפול גבוהה פי 2.90 וקצב של 58.5 אלף רזידואים לשעת-GPU, נתון שהופך ריצות פרוטאומיות שלמות למעשיות בזמני המתנה סבירים.