NVIDIA ra mắt BioNeMo Inference Runtime: Tăng tốc gấp 2.9 lần hiệu suất dự đoán cấu trúc protein

Dự đoán cấu trúc phân tử sinh học đã chuyển dịch từ các lượt chạy đơn mục tiêu sang danh sách công việc ở quy mô hệ protein (proteome). Điểm nghẽn hiện nay không còn nằm ở việc liệu một mô hình có thể gập protein hay không, mà là tốc độ xử lý toàn bộ hàng đợi các mục tiêu độc lập thông qua các bước phân tích cú pháp, tạo đặc trưng, suy luận trên GPU và ghi kết quả đầu ra. Bài phân tích kỹ thuật chuyên sâu mới của NVIDIA giới thiệu BioNeMo Inference Runtime (BioIR), một thư viện Python giúp tăng tốc các mô hình dự đoán cấu trúc được hỗ trợ trên GPU NVIDIA mà vẫn giữ nguyên quy trình làm việc tiêu chuẩn của PyTorch. BioIR đã được vận hành ở quy mô sản xuất. Nó là nền tảng cho sự mở rộng gần đây của Cơ sở dữ liệu AlphaFold, tạo ra cấu trúc phức hợp protein trên 4.777 hệ protein, tương đương khoảng 31 triệu phức hợp ứng viên, với 1,81 triệu cấu trúc được công bố dưới dạng dự đoán có độ tin cậy cao.

Nó có thể triển khai được không? Có. BioIR hiện đã có sẵn dưới dạng kho lưu trữ GitHub mã nguồn mở với một wheel chứa các CUBIN đã biên dịch sẵn. Việc sử dụng runtime yêu cầu Python 3.12+, GPU và driver NVIDIA tương thích, checkpoint mô hình đã được dàn dựng (staged) và các tệp A3M MSA cho từng chuỗi. Nó không yêu cầu nvcc, mã nguồn CUDA, CMake hoặc bộ công cụ CUDA.

BioIR là gì

BioIR nhắm vào các thao tác mà các ngăn xếp suy luận (inference stacks) đa năng chưa tối ưu hóa hoàn toàn. Các thao tác này bao gồm các ngăn xếp Pairformer và Evoformer, các phép toán tam giác, cơ chế chú ý theo cặp (pairwise attention), diffusion transformer và các mô-đun ở cấp độ nguyên tử. Các mô hình vẫn là các đối tượng torch.nn.Module thông thường. Không cần xây dựng engine, bước xuất file hay tạo artifact riêng biệt nào giữa checkpoint và quá trình truyền tiến (forward pass).

Có 2 cách để sử dụng nó. Bộ xử lý end-to-end di chuyển một InputRequest qua các bước phân tích cú pháp, mã hóa (tokenization), tạo đặc trưng, suy luận trên GPU và ghi tệp PDB hoặc mmCIF. Việc tích hợp trực tiếp vào PyTorch cho phép các nhà phát triển xây dựng một mô hình được hỗ trợ hoặc tái sử dụng các mô-đun đã tối ưu hóa được chọn lọc bên trong mã tùy chỉnh. Hướng dẫn minh họa đường dẫn bộ xử lý với Boltz-2 (model_source="boltz-2"). Mỗi chuỗi protein yêu cầu một tệp A3M MSA. Các MSA theo cặp hoặc không theo cặp đều được chấp nhận cho các đầu vào có nhiều chuỗi protein không giống nhau. Các khuôn mẫu (templates) có thể được cung cấp thủ công vì BioIR không chạy HHsearch hoặc HMMsearch. Bộ xử lý hỗ trợ dự đoán cấu trúc phối tử (ligand) nhưng không hỗ trợ dự đoán ái lực phối tử.

Ba lớp tăng tốc

BioIR tối ưu hóa tại 3 lớp riêng biệt, mỗi lớp nhắm vào một điểm nghẽn khác nhau:

Lưu ý: Ray không chia nhỏ một lần truyền tiến (forward pass) đơn lẻ trên các GPU. Chế độ Replica mở rộng danh sách công việc, không phải từng mục tiêu riêng lẻ. Theo ma trận hỗ trợ, việc gập song song theo ngữ cảnh (context-parallel folding) đã được lên kế hoạch nhưng chưa khả dụng. Quy tắc về công suất rất đơn giản: engine_stage.compute x số lượng GPU không được vượt quá số GPU hiển thị.

Ở cấp độ truyền tiến của mô hình, các báo cáo đo lường sớm của NVIDIA cho thấy tốc độ tăng trung bình hình học so với baseline OSS torch.compile là 1,55x (OpenFold3), 1,78x (Boltz2) và 2,56x (OpenFold2 monomer) trên H100. Các con số trên H200 tương tự ở mức 1,54x, 1,75x và 2,61x. Các kết quả này được đo lường trên 17 đầu vào với độ dài từ 29 đến 1.734 gốc axit amin.

Điểm chuẩn: 1.000 dimer ở người trên 8xH100

Để định lượng hiệu suất phân phối end-to-end, nhóm NVIDIA đã chạy một bài kiểm tra đối chứng trên 1.000 mục tiêu dimer ở người với tổng độ dài chuỗi dưới 2.800 gốc axit amin. So sánh này đặt Boltz-2 được tăng tốc bởi BioIR đối đầu với triển khai Boltz-2 mã nguồn mở đã được biên dịch bằng torch trên 8 GPU H100 80GB. Cả hai đều sử dụng các mục tiêu giống hệt nhau, MSA đã dàn dựng, công thức suy luận (3 lần tái chế, 200 bước lấy mẫu, 5 mẫu khuếch tán) và cấu hình GPU.

Kết quả:

Những con số này là các phép đo ở giai đoạn gập protein dành riêng cho tập dữ liệu và phần cứng này. Chúng không bao gồm việc tạo MSA, phân bổ CPU tiền xử lý, lưu trữ, truyền dữ liệu và các lần thử lại. Blog cảnh báo rõ ràng về việc không nên khái quát hóa chúng cho tất cả các mô hình hoặc tập dữ liệu được BioIR hỗ trợ.

Năng lượng cho một triệu mục tiêu

Ngoại suy tuyến tính bài kiểm tra cho 1 triệu mục tiêu tương đương, BioIR ước tính cần 11 MWh so với 35 MWh cho triển khai công cộng khi sử dụng mức tiêu thụ điện năng (TDP) tương đương của 8 GPU. Sử dụng mức tiêu thụ điện năng tối đa của toàn node, ước tính là 21 MWh so với 64 MWh. Đây là các ước tính về công suất định mức, chỉ dành cho giai đoạn gập protein đối với thiết bị CNTT, không phải là các phép đo thực tế và không bao gồm chi phí vận hành trung tâm dữ liệu như PUE. Tuy nhiên, mức tiết kiệm từ 23 đến 43 MWh cho mỗi triệu mục tiêu là một con số đáng kể cho các chiến dịch ở quy mô hệ protein.

Những điểm chính cần lưu ý

Hãy xem blog kỹ thuật, kho lưu trữ GitHub, tài liệu và BioNeMo Agent Toolkit để điều phối tác nhân. Ngoài ra, đừng ngần ngại theo dõi chúng tôi trên Twitter và đừng quên tham gia cộng đồng 150k+ ML SubReddit của chúng tôi cũng như đăng ký nhận bản tin. Khoan đã! Bạn có dùng Telegram không? Bây giờ bạn cũng có thể tham gia cùng chúng tôi trên Telegram.

Bạn cần hợp tác với chúng tôi để quảng bá kho lưu trữ GitHub, trang Hugging Face, bản phát hành sản phẩm hoặc hội thảo trực tuyến, v.v. của bạn? Hãy kết nối với chúng tôi.