NVIDIA BioNeMo Inference Runtime가 구조 예측 처리량을 높이는 방법
생체분자 구조 예측은 소수의 표적을 넘어 프로테옴 단위로 확장되고 있다. 이 규모에서는 단일 예측의 속도뿐 아니라 대량의 작업을 얼마나 안정적으로 완료할 수 있는지가 중요하다. NVIDIA BioNeMo Inference Runtime(BioIR)은 NVIDIA GPU에서 생체분자 구조 예측 모델의 추론과 작업 분산을 가속하는 Python 라이브러리다.
모델 최적화와 작업 확장을 결합한다
BioIR은 입력 형태와 GPU에 맞는 사용자 정의 커널을 선택하고, 호환되는 모듈에는 CUDA Graph 캡처를 적용한다. 여러 입력을 처리할 때는 Ray를 이용해 GPU마다 전체 모델 복제본 하나를 배치하고 작업을 분산한다.1
이 방식은 하나의 예측을 여러 GPU로 나누는 모델 병렬화와 다르다. 각 GPU가 서로 독립적인 표적을 처리하므로 대규모 작업 목록에서 처리량을 높이는 데 적합하다. 실제 효과를 평가할 때는 GPU 추론 속도만이 아니라 입력 크기 분포, CPU 전처리, 결과 저장 시간까지 포함한 전체 완료 시간을 측정해야 한다.
8개 H100과 인간 이합체 표적 1,000개를 사용한 벤치마크에서 BioIR 기반 Boltz-2는 GPU 시간당 5만 8,500개 잔기를 성공적으로 폴딩했다. 같은 조건에서 torch.compile로 컴파일한 오픈소스 구현은 GPU 시간당 2만 200개를 기록해, BioIR 기반 Boltz-2의 처리량이 2.90배 높았다. 다만 이 수치는 특정 데이터와 추론 설정에서 나온 결과이므로 다른 워크로드에 그대로 적용된다고 볼 수는 없다.1
두 가지 통합 방식 중 선택할 수 있다
BioIR은 종단 간 프로세서와 직접 PyTorch 모듈 통합이라는 두 가지 사용 방식을 제공한다.1
종단 간 프로세서는 다음 단계를 처리한다.
입력 파싱
토큰화와 특징 생성
GPU 추론
PDB 또는 mmCIF 형식으로 결과 저장
기존 파이프라인을 유지해야 한다면 직접 통합 방식을 사용할 수 있다. 지원 모델을 torch.nn.Module로 구성하거나, 필요한 모듈만 사용자 코드에서 재사용할 수 있다. 전자는 빠르게 완성된 워크플로를 구축하는 데 적합하고, 후자는 기존 시스템에 필요한 기능만 선택적으로 결합할 때 유용하다.
설치 환경도 비교적 단순하다. BioIR 휠에는 사전 컴파일된 CUBIN이 포함돼 런타임에 nvcc, CUDA 소스, CMake 또는 CUDA Toolkit이 필요하지 않다. 다만 Python 3.12 이상과 호환되는 NVIDIA GPU 및 드라이버는 필요하다.1
입력 준비와 지원 범위를 먼저 확인해야 한다
추론 런타임이 빨라도 필요한 입력이 준비되지 않으면 전체 워크플로는 완료되지 않는다. Boltz-2 워크플로에서는 각 단백질 사슬에 A3M 형식의 다중서열정렬(MSA)이 필요하다. 서로 다른 단백질 사슬이 여러 개 포함된 입력에는 페어드 MSA와 언페어드 MSA를 모두 사용할 수 있다.1
템플릿 검색도 BioIR의 역할에 포함되지 않는다. BioIR은 HHsearch나 HMMsearch를 실행하지 않으므로 템플릿이 필요하면 사용자가 직접 제공해야 한다. 또한 종단 간 프로세서는 리간드 구조 예측을 지원하지만 리간드 친화도 예측은 지원하지 않는다.1
따라서 도입 전에는 다음 사항을 확인해야 한다.
각 사슬에 필요한 MSA를 생성하거나 확보할 수 있는가?
템플릿 검색과 공급을 별도 파이프라인에서 처리할 수 있는가?
필요한 결과가 리간드 구조인지, 리간드 친화도까지 포함하는지 명확한가?
종단 간 프로세서와 직접 통합 방식 중 기존 시스템에 맞는 경로는 무엇인가?
성능 배수보다 실제 워크로드가 중요하다
NVIDIA는 AlphaFold Database 확장에서 4,777개 프로테옴에 걸쳐 약 3,100만 개의 후보 단백질 복합체를 생성하는 데 BioIR을 사용했다. 이는 독립적인 예측 작업을 대규모로 확장할 수 있음을 보여준다.1
그러나 도입 여부를 결정할 때는 공개된 2.90배라는 수치보다 자체 워크로드를 기준으로 검증해야 한다. 대표 입력을 사용해 단일 GPU 실행과 여러 Ray 복제본 실행을 비교하고, GPU 추론 시간과 전체 작업 완료 시간을 각각 측정하는 것이 핵심이다.
요약
NVIDIA BioNeMo Inference Runtime(BioIR)의 가치는 빠른 커널에만 있지 않다. CUDA Graph와 사용자 정의 커널로 모델 실행을 최적화하고, Ray 기반 복제본 확장으로 독립적인 표적을 여러 GPU에 분산하며, 전처리부터 결과 저장까지 연결할 수 있다는 점이 핵심이다. 실제 적용 시에는 MSA와 템플릿 준비, 지원 기능의 범위, 전체 파이프라인 처리량을 함께 검토해야 한다.
1High-Throughput Structure Prediction with BioNeMo Inference Runtime
이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.