生物分子結構預測如今通常在整個蛋白質組規模上運行，其目標是高效地推動整個工作列表通過處理流程。
NVIDIA BioNeMo推理運行時（BioIR）有助於在NVIDIA GPU上加速受支持的生物分子結構預測模型，同時保持熟悉的PyTorch工作流程。它使用優化的核心，並在適用的情況下使用CUDA Graphs來加速模型執行。對於大批量的獨立輸入，Ray可以在單個節點上的每個GPU上運行一個完整的模型副本，以提高整體吞吐量。
BioIR還已用於實際的蛋白質組規模工作，包括最近擴展的AlphaFold資料庫（AFDB），加速了跨4777個蛋白質組生成蛋白質複合物結構，總計約3100萬個候選複合物，其中181萬個作為高置信度預測被發布。
您可以通過兩種方式使用它（見圖1）：
端到端處理器通過解析、Token化、特徵生成、GPU推理以及PDB或mmCIF寫入來處理InputRequest。
直接PyTorch集成允許您構建受支持的模型（torch.nn.Module）或在自定義代碼中重用選定的模組。
本教學演示了BioIR的端到端處理器，從輸入準備到GPU推理和PDB或mmCIF輸出，並展示了如何跟蹤每小時結構數和資源效率。
先決條件
Python 3.12或更高版本
兼容的NVIDIA GPU和驅動程式，以及BioIR wheel包或受支持的開發環境
暫存的模型檢查點（下面的示例中為Boltz-2）和所需的化學元數據
每個蛋白質鏈需要一個A3M MSA。對於具有多個非相同蛋白質鏈的輸入，接受配對或非配對MSA
對於Ray吞吐量擴展，請在同一節點上使用多個可見GPU，以及比副本數量更多的獨立記錄
該wheel包中包含預編譯的CUBIN，因此運行時使用不需要nvcc、CUDA源代碼、CMake或CUDA工具包。
步驟1. 選擇受支持的結構預測工作流程
下面，我們演示了BioIR中Boltz2的端到端工作流程。使用model_source="boltz-2"。蛋白質鏈需要MSA；對於具有多個非相同蛋白質鏈的輸入，接受配對或非配對MSA。您可以選擇自行提供模板，因為BioIR不運行HHsearch或HMMsearch。端到端處理器支持配體結構預測，但不支持配體親和力預測。
（代碼示例：構建InputRequest，包含蛋白質鏈、序列及MSA記錄）
請用有效值替換截斷的序列和MSA路徑。對於Ray測試，請從真實的工作列表構建數據行，記錄數量要多於副本數量；不要重複使用一行數據來證明有效的擴展性。
步驟2. 使用串行處理器驗證單次預測
BioIR的端到端處理器工作流程有兩種執行器後端：串行和Ray。串行後端為單個輸入依次運行每個階段，在完成整個工作流程後再處理下一個輸入。這使其在使用Ray後端並發處理獨立輸入（步驟3）之前，非常適合用來檢查您的設置。
（代碼示例：配置串行處理器，設置回收步驟、採樣步驟、擴散樣本數等參數，並運行處理器獲取輸出）
model_inference_time是BioIR經CUDA同步的摺疊模型前向計算測量值。它不包括解析、Token化、特徵生成、後處理和寫入。請通過特徵生成器的init_context設置隨機種子。scores欄位必須解碼，因為它是一個JSON字符串。
步驟3. 使用Ray副本擴展獨立輸入
可以按如下方式選擇Ray後端，用於默認的副本布局：
（代碼示例：使用create_default_replica_mode_config創建默認Ray配置）
此配置在當前節點上每個可見GPU上放置一個完整的模型副本，並根據torch.cuda.device_count()調整CPU階段的大小。下面是一個顯式控制四個GPU的替代配置：
（代碼示例：顯式配置解析器、Token化器、特徵生成器、引擎和寫入器階段，指定計算資源、CPU數量、GPU數量等）
容量規則為：engine_stage.compute × engine_stage.num_gpus ≤ 可見GPU數量。這個四副本示例是一個單節點配置，假定有四個可見GPU。本教學不涉及多節點部署。這裡Ray創建了四個引擎actor，每個都預留一個GPU。每個actor加載完整模型。build_processor會在需要時初始化Ray。實際吞吐量取決於輸入分布、階段平衡、儲存、調度和故障情況，因此需要實際測量。
步驟4. 平衡五個處理器階段
在Ray端到端處理器中，五個處理器階段按以下依賴順序處理輸入：解析器→Token化器→特徵生成器→摺疊引擎→寫入器。使用相應的*StageConfig配置每個階段；步驟3的示例展示了相關欄位。enabled欄位不是公開的跳過控制項。每個階段都公開compute；相關階段還公開num_cpus、memory和batch_size。Ray引擎添加了max_concurrent_batches、accelerator_type和num_gpus。
要調整Ray流水線，請從EngineProcessorConfig.create_default_replica_mode_config(...)開始。增加某個階段的compute以添加工作節點；使用num_cpus、memory以及（對於引擎actor）num_gpus來設置資源預留。如果引擎在等待輸入，請添加解析器、Token化器或特徵生成器工作節點。當GPU或對象儲存內存導致故障時，請降低並發度或分離大型輸入。
Ray旨在使CPU階段與推理重疊，但這是否能改善目標工作負載，取決於在給定硬體配置下，給定輸入的解析、特徵生成和輸出寫入階段的運行時間成本。請參閱ScaleFold的圖4，了解OpenFold預處理時間的多樣性。
步驟5. 區分單副本加速與流水線擴展
BioIR在三個不同層面提供優化：
核心選擇：受支持的操作根據模型配置、GPU、數據類型和張量形狀，選擇兼容的BioIR自定義實現、cuEquivariance實現或PyTorch後備實現。
模組優化：在支持的情況下，獨立的optimize()機制為兼容模組啟用CUDA Graph捕獲。
流水線擴展：Ray執行器將完整的模型副本放置在多個GPU上，並在它們之間分配獨立的輸入。
這些層面針對不同的瓶頸。核心和模組優化減少了單個副本內的模型前向時間。Ray可以通過將CPU階段與GPU摺疊計算重疊，以及在獨立的GPU上運行完整模型副本處理獨立輸入，來提高工作列表的吞吐量。Ray不會將單次模型前向傳遞拆分到多個GPU上。上面的圖1區分了處理器路徑和直接集成路徑；Ray擴展僅適用於處理器路徑。
我們對BioNeMo推理運行時的早期基準測試估計了以下模型前向加速效果：
加速比是使用1次預熱運行（丟棄）和1次測量調用，跨越29-1734個殘基的17個輸入進行測量的。OpenFold3和Boltz2開源基線使用了torch.compile，參數為dynamic=None、fullgraph=False、recompile_limit=128、accumulated_recompile_limit=256、fail_on_recompile_limit_hit=True。Boltz2開源版本使用了cuEq；OpenFold3開源版本使用了use_cuequivariance=True和use_deepspeed=True。
這些結果量化了單個模型副本內的加速效果。它們並未測量解析、特徵生成、輸出寫入、Ray調度、多GPU吞吐量或完整工作列表的實際耗時。下面的圖3展示了為什麼模型前向測量和端到端測量必須保持獨立。
要確定額外GPU在實際部署中能帶來什麼提升，請在單個節點上使用一個、兩個和四個Ray副本，測量同一代表性工作列表。步驟6定義了所需的指標和比較方法。
步驟6. 基準測試摺疊階段效率和端到端交付：AFDB案例研究
為了使這些測量更加具體，我們在1000個人類二聚體目標（組合序列長度低於2800個殘基）上進行了匹配基準測試，代表了一大批獨立的生物分子結構預測任務，類似於AlphaFold資料庫中最近添加的數據集。
這個代表性的摺疊階段基準測試在8xH100 GPU上，比較了BioIR加速的Boltz-2與torch編譯的開源Boltz-2實現。兩種實現使用了相同的目標、暫存的MSA、推理方案和GPU配置；吞吐量指標和其他結果是特定於此配置的，不應推廣到所有BioIR支持的模型、數據集或硬體。
該工作流程使用了三次回收、200個採樣步驟，每個目標五個擴散樣本。BioIR完成了全部1000個目標，每分配的GPU小時交付了58.5K個成功摺疊的殘基，而公開實現為20.2K——在殘基歸一化吞吐量方面提升了2.90倍；開源實現在29個目標上出現了內存不足的情況。
上面圖3的左側面板比較了BioIR與torch編譯的開源實現的模型前向時間，直接展示了BioIR實現的更低模型前向時間。圖3的右側面板比較了BioIR與開源實現所交付的吞吐量，其中吞吐量為預測結構中的殘基總數除以分配的GPU小時數。圖3的右側面板展示了BioIR中核心級、模組級和流水線級優化所帶來的加速效果。圖3的左側面板展示了核心級和模組級實現所帶來的加速效果。
與Boltz2的加速效果類似，BIR還為其他生物分子共摺疊模型（如OpenFold2和OpenFold3）提供了更快的推理速度。BIR的早期版本為NVIDIA內部版本的OpenFold2-MM貢獻了加速模組，使得能夠為擁有3100萬個蛋白質複合物結構的AFDB進行大規模蛋白質結構預測。
我們將圖3中1000個目標的匹配基準測試線性外推到100萬個可比目標，使用8×H100 80GB HBM3節點的額定功率等效值（見下面的圖4）。
據估計，使用8-GPU TDP（熱設計功耗）等效值時，BioIR需要11兆瓦時，而公開實現需要35兆瓦時；使用整節點最大功率等效值時，兩者分別為21兆瓦時和64兆瓦時。這些僅是針對IT設備的摺疊估算值，而非實測能耗數據，且不包括數據中心開銷，如電源使用效率（PUE）。
這項受控比較使用每種實現相同的輸入和MSA來測量摺疊吞吐量；不包括MSA生成、預處理CPU分配、儲存、數據傳輸、重試和工程開銷。請將端到端流水線性能指標（包括每小時完成的結構數、GPU和CPU利用率、GPU內存峰值、完成率、故障和重試）與模型前向指標區分開來報告。
故障排除
只有一個GPU處於活動狀態：確認使用了Ray、REPLICA模式、多於一個副本、多個可見GPU，以及足夠的獨立記錄數量。
處理器構建引發`ValueError`：檢查compute * num_gpus是否超過了可見GPU數量。
蛋白質輸入失敗：檢查所需的非配對A3M文件以及工作節點可見的路徑。
GPU處於等待狀態：在添加副本之前，先檢查CPU階段、CPU預留、排隊情況以及Ray對象儲存容量。
推理後數據行處於等待狀態：檢查寫入器並發度和目標儲存的吞吐量。
Ray無法放置actor：檢查CPU、GPU、內存以及accelerator_type標籤。
單條記錄導致任務停止：默認的快速失敗模式會引發FoldingPredictionError。僅在有意進行行級容錯的情況下才設置should_continue_on_error=True，然後檢查__inference_error__。
開始使用
探索BioNeMo推理運行時（BioIR），並將其集成到您大規模的結構預測工作流程中：http://github.com/NVIDIA-BioNeMo/BioNeMo-Inference-Runtime
要通過智能體編排進一步加速藥物發現工作流程，請查看NVIDIA BioNeMo Agent工具包（BAT）。
有關最新的加速數據，請查閱API參考文檔和支持矩陣。
Q&A
Q1：BioNeMo推理運行時（BioIR）是什麼？
A：BioIR是NVIDIA推出的用於加速生物分子結構預測模型的運行時工具，它可以在GPU上加速支持的模型執行，同時保持PyTorch工作流程不變，適用於蛋白質組規模的結構預測任務。
Q2：BioIR相比開源實現能帶來多大的性能提升？
A：在1000個人類二聚體目標的基準測試中，BioIR每GPU小時能交付58.5K個成功摺疊殘基，而公開實現只有20.2K，吞吐量提升了2.90倍，同時能耗也更低。
Q3：使用BioIR進行大規模結構預測需要什麼條件？
A：需要Python 3.12及以上版本、兼容的NVIDIA GPU和驅動、模型檢查點及化學元數據，每個蛋白質鏈還需要A3M格式的MSA文件。若要使用Ray進行吞吐量擴展，需要同一節點上有多個可見GPU。
