생체 분자 구조 예측은 이제 단백체(Proteome) 규모로 활발히 수행되고 있으며, 이 분야의 핵심 과제는 대량의 작업 목록을 파이프라인으로 효율적으로 처리하는 것입니다.
NVIDIA BioNeMo Inference Runtime은 익숙한 PyTorch 워크플로우를 유지하면서 NVIDIA GPU상에서 주요 생체 분자 구조 예측 모델의 가속을 지원합니다. 최적화된 커널과 CUDA Graphs를 활용해 모델 실행 속도를 높이며, 대용량 배치 처리가 필요한 경우 Ray를 통해 단일 노드의 GPU마다 독립적인 모델 복제본을 할당함으로써 전체 처리량을 극대화합니다.
실제로 BioIR은 AlphaFold 데이터베이스 확장 작업(AFDB) 등 대규모 단백체 연구에 적용되었습니다. 총 4,777개 단백체에 달하는 약 3,100만 개의 후보 복합체 구조 조립 및 생성을 가속화했으며, 이 가운데 181만 개가 고신뢰도 예측 구조로 공식 공개되었습니다.
BioIR은 아래 그림 1과 같이 두 가지 방식으로 활용할 수 있습니다.
- 엔드투엔드 프로세서: InputRequest 파싱부터 토큰화, 피처 생성, GPU 추론, PDB 또는 mmCIF 파일 생성까지 전 과정을 일괄 처리합니다.
- 직접 PyTorch 통합: 지원되는 torch.nn.Module 기반 모델을 직접 실행하거나, 필요한 모듈만 선별하여 커스텀 코드에 재사용합니다.
본 튜토리얼에서는 입력 데이터 준비부터 GPU 추론, 최종 파일 출력에 이르는 BioIR 엔드투엔드 프로세서의 활용법을 다루며, 시간당 구조 예측 수와 자원 사용 효율을 추적하는 방법까지 함께 소개합니다.
사전 요구사항
- Python 3.12 이상
- 호환 가능한 NVIDIA GPU 및 드라이버, BioIR 휠 파일 또는 지원되는 개발 환경
- 스테이징된 Boltz-2 모델 체크포인트 및 필수 화학 메타데이터
- 각 단백질 체인에 필요한 A3M MSA. 동일하지 않은 다중 단백질 체인 입력의 경우 페어링 또는 언페어링 MSA 수용 가능
- Ray 처리량 확장을 위해 동일 노드 내 복수의 가용 GPU를 확보하고, 복제본 수보다 많은 독립 레코드를 준비할 것
제공되는 휠 파일에는 사전 컴파일된 CUBIN이 포함되어 있으므로 런타임 실행 시 nvcc, CUDA 소스 코드, CMake, CUDA 툴킷 등이 별도로 필요하지 않습니다.
1단계. 지원되는 구조 예측 워크플로우 선택
이하 내용은 BioIR 환경에서 Boltz-2 기반 엔드투엔드 워크플로우를 구현하는 방안을 다룹니다. model_source="boltz-2" 옵션을 적용하여 실행합니다. 단백질 체인에는 MSA가 필수적이며, 동일하지 않은 다중 단백질 체인 입력의 경우 페어링 또는 언페어링 MSA를 사용할 수 있습니다. BioIR은 HHsearch나 HMMsearch를 직접 실행하지 않으므로 사용자가 템플릿을 임의로 지정할 수 있습니다. 엔드투엔드 프로세서는 리간드 구조 예측을 지원하지만, 리간드 결합 친화도 예측은 지원하지 않습니다.
from bionemo_ir.data.schemas import InputRequest, MSARecord, Polymer
request = InputRequest(
input_id="demo",
polymers=[
Polymer(
polymer_type="protein",
chain_id=["A"],
sequence="GSHMSL...",
msas=[MSARecord(path="msa.a3m", format="a3m")],
paired_msas=[],
templates=None,
)
],
)
rows = [{"record": request, "__record_id": request["input_id"]}]
생략된 서열 및 MSA 경로를 유효한 값으로 교체하십시오. Ray 테스트의 경우 복제본 수보다 많은 레코드를 포함하는 실제 작업 목록으로 행을 구성해야 하며, 유용한 스케일링의 증거로 동일한 행을 반복 사용하지 마세요.
2단계. 직렬 프로세서로 단일 예측 검증
BioIR은 엔드투엔드 프로세서 워크플로우를 위해 직렬(Serial)과 Ray라는 두 가지 실행기 백엔드를 제공합니다. 직렬 백엔드는 단일 입력에 대해 각 단계를 순차적으로 실행하며, 다음 입력으로 이동하기 전에 전체 워크플로우를 완료합니다. 이는 3단계에서 Ray 백엔드를 사용해 독립된 입력을 동시 처리하기 전, 설정을 검증하는 데 유용합니다.
import json
from bionemo_ir.pipeline.processor.engine_proc import (
EngineProcessorConfig,
build_processor,
)
from bionemo_ir.pipeline.stages.configs import (
FeatureGeneratorStageConfig,
WriterStageConfig,
)
serial_config = EngineProcessorConfig(
model_source="boltz-2",
executor_backend=None, # None은 직렬 프로세서를 의미함
runtime_args={
"recycling_steps": 3,
"num_sampling_steps": 50,
"diffusion_samples": 1,
},
feature_generator_stage=FeatureGeneratorStageConfig(
init_context={"random_seed": 42},
),
writer_stage=WriterStageConfig(
output_path="output/serial",
format="cif",
),
engine_kwargs={"profile_inference": True},
)
serial_processor = build_processor(serial_config)
serial_outputs = serial_processor(rows)
for row in serial_outputs:
scores = json.loads(row["scores"])
print(row["output_path"])
print(row["model_inference_time"])
print(scores.get("confidence_score"))
model_inference_time은 BioIR의 CUDA 동기화 기반 폴딩 모델 순전파 측정값입니다. 이 측정값에서는 파싱, 토큰화, 피처 생성, 후처리, 파일 작성이 제외됩니다. 시드 값은 피처 생성기의 init_context를 통해 설정합니다. scores 필드는 JSON 문자열 형태이므로 반드시 디코딩해야 합니다.
3단계. Ray 복제본을 통한 독립 입력 스케일링
기본 복제본 모드 구성을 위한 Ray 백엔드는 다음과 같이 선택할 수 있습니다.
from bionemo_ir.pipeline.processor.engine_proc import EngineProcessorConfig
ray_config = EngineProcessorConfig.create_default_replica_mode_config(
model_source="boltz-2",
output_dir="output/ray",
output_format="cif",
)
이 구성은 현재 노드의 가용 GPU마다 하나의 완전한 모델 복제본을 배치하고, torch.cuda.device_count()를 기반으로 CPU 단계의 규모를 결정합니다. 아래는 4대의 GPU를 명시적으로 제어하는 대안 구성 예시입니다.
import ray
from bionemo_ir.pipeline.processor.engine_proc import (
EngineProcessorConfig,
build_processor,
)
from bionemo_ir.pipeline.stages.configs import (
EngineStageConfig,
FeatureGeneratorStageConfig,
ParallelismMode,
ParserStageConfig,
TokenizerStageConfig,
WriterStageConfig,
)
ray_config = EngineProcessorConfig(
model_source="boltz-2",
executor_backend="ray", # ray 백엔드 선택
parser_stage=ParserStageConfig(compute=4),
tokenizer_stage=TokenizerStageConfig(compute=4, num_cpus=2),
feature_generator_stage=FeatureGeneratorStageConfig(
compute=8,
num_cpus=4,
init_context={"random_seed": 42},
),
engine_stage=EngineStageConfig(
parallelism_mode=ParallelismMode.REPLICA,
compute=4,
num_gpus=1.0,
num_cpus=4,
),
writer_stage=WriterStageConfig(
compute=4,
output_path="output/ray",
format="cif",
),
)
ray_processor = build_processor(ray_config)
dataset = ray.data.from_items(rows)
ray_outputs = list(ray_processor(dataset).materialize().iter_rows())
수용 용량 규칙은 engine_stage.compute × engine_stage.num_gpus ≤ 가용 GPU 수입니다. 이 4개 복제본 예시는 4대의 가용 GPU를 가정하는 단일 노드 구성입니다. 본 튜토리얼에서는 다중 노드 배포를 다루지 않습니다. 여기서 Ray는 4개의 엔진 액터를 생성하고 각 액터마다 1대의 GPU를 할당합니다. 모든 액터는 전체 모델을 로드합니다. build_processor는 필요한 경우 Ray를 자동으로 초기화합니다. 실제 처리량은 입력 데이터의 분포, 단계별 균형, 스토리지, 스케줄링, 오류 발생 여부에 따라 달라지므로 직접 측정해야 합니다.
4단계. 5개 프로세서 단계의 균형 조정
Ray 엔드투엔드 프로세서에서 5개 프로세서 단계는 파서 → 토큰화 생성기 → 피처 생성기 → 폴딩 엔진 → 작성기 순서의 의존성에 따라 입력을 소비합니다. 각 단계를 해당하는 *StageConfig로 구성하십시오. 3단계 예시에서 관련 필드를 확인할 수 있습니다. enabled 필드는 공용 건너뛰기 제어 옵션이 아닙니다. 각 단계는 compute를 제공하며, 관련 단계는 num_cpus, memory, batch_size도 제공합니다. Ray 엔진은 max_concurrent_batches, accelerator_type, num_gpus를 추가로 제공합니다.
Ray 파이프라인을 튜닝하려면 EngineProcessorConfig.create_default_replica_mode_config(...)로 시작하십시오. 작업자(Worker)를 추가하려면 특정 단계의 compute를 늘리고, 자원 할당을 설정하려면 num_cpus, memory, 그리고 엔진 액터의 경우 num_gpus를 활용하십시오. 엔진이 입력을 기다리는 병목이 발생하면 파서, 토큰화 생성기, 피처 생성기 작업자를 추가하십시오. GPU나 오브젝트 스토어 메모리로 인해 오류가 발생하는 경우 동시성을 낮추거나 대용량 입력을 분리하세요.
Ray는 CPU 단계와 추론을 오버랩하도록 설계되었으나, 이것이 타깃 워크로드를 실제로 개선하는지 여부는 특정 하드웨어 구성 및 입력 데이터에 대한 파싱, 피처 생성, 출력 작성 단계의 런타임 비용에 따라 달라집니다. OpenFold의 전처리 시간 다양성에 대한 자세한 내용은 ScaleFold의 그림 4를 참조하세요.
5단계. 복제본별 가속과 파이프라인 스케일링의 분리
BioIR은 세 가지 서로 다른 레이어에서 최적화를 제공합니다.
- 커널 선택: 지원되는 연산은 모델 구성, GPU, 데이터 타입, 텐서 형상에 따라 호환되는 BioIR 커스텀 구현, cuEquivariance, 또는 PyTorch 폴백 구현을 선택합니다.
- 모듈 최적화: 지원되는 환경에서 별도의
optimize()메커니즘을 통해 호환 가능한 모듈에 대한 CUDA Graph 캡처를 활성화합니다. - 파이프라인 스케일링: Ray 실행기는 GPU에 완전한 모델 복제본을 배치하고 독립적인 입력들을 복제본 사이에 분배합니다.
이 레이어들은 서로 다른 병목 지점을 타깃으로 합니다. 커널 및 모듈 최적화는 복제본 내부의 모델 순전파 시간을 단축합니다. Ray는 CPU 단계와 GPU 폴딩을 오버랩하고, 독립된 입력에 대해 별도의 GPU에서 전체 모델 복제본을 실행함으로써 전체 작업 목록의 처리량을 높일 수 있습니다. Ray는 단일 모델 순전파 연산을 여러 GPU에 분할하여 처리하지 않습니다. 상단 그림 1은 프로세서 경로와 직접 통합 경로를 구별하여 보여주며, Ray 스케얼링은 프로세서 경로에만 적용됩니다.
BioNeMo Inference Runtime을 활용한 초기 벤치마킹에서 추정된 모델 순전파 가속 효과는 다음과 같습니다.
속도 향상은 29~1,734개 잔기 범위의 17개 입정을 대상으로 1회 워밍업 실행(제외됨) 및 1회 측정 호출을 사용해 측정되었습니다. OpenFold3 및 Boltz2 오픈소스 베이스라인은 dynamic=None, fullgraph=False, recompile_limit=128, accumulated_recompile_limit=256, fail_on_recompile_limit_hit=True 조건의 torch.compile을 사용했습니다. Boltz2 오픈소스는 cuEq를 사용했고, OpenFold3 오픈소스는 use_cuequivariance=True 및 use_deepspeed=True를 사용했습니다.
이 결과는 단일 모델 복제본 내부의 가속을 정량화한 것입니다. 파싱, 피처 생성, 출력 작성, Ray 스케줄링, 다중 GPU 처리량, 또는 전체 작업 목록의 경과 시간(Wall time)은 측정에 포함되지 않았습니다. 하단 그림 3은 모델 순전파 측정과 엔드투엔드 측정을 반드시 분리해야 하는 이유를 보여줍니다.
실제 배포 시 추가 GPU가 제공하는 효과를 확인하려면, 단일 노드에서 1개, 2개, 4개의 Ray 복제본을 사용해 동일한 대표 작업 목록을 측정해보십시오. 6단계에서 필요한 지표와 비교 방법을 정의합니다.
6단계. 폴딩 단계 효율성 및 엔드투엔드 처리 성능 벤치마크: AFDB 사례 연구
측정을 구체화하기 위해 결합 서열 길이가 2,800개 잔기 미만인 1,000개의 인간 이량체(Dimer) 타깃을 대상으로 비교 벤치마크를 실행했습니다. 이는 최근 AlphaFold 데이터베이스에 추가된 데이터셋과 유사한, 독립된 생체 분자 구조 예측 작업의 대규모 집합을 대표합니다.
이 대표적인 폴딩 단계 벤치마크는 8x H100 GPU 환경에서 BioIR로 가속된 Boltz-2와 torch.compile이 적용된 오픈소스 Boltz-2 구현을 비교합니다. 두 구현 모두 동일한 타깃, 스테이징된 MSA, 추론 레시피, GPU 구성을 사용했습니다. 처리량 지표 및 기타 결과는 해당 구성에 특화된 것이므로 모든 BioIR 지원 모델, 데이터셋, 하드웨어로 일반화해서는 안 됩니다.
워크플로우는 타깃당 3회의 리사이클, 200회의 샘플링 단계, 5회의 디퓨전 샘플을 사용했습니다. BioIR은 1,000개 타깃을 모두 완료하고 할당된 GPU 시간당 5만 8,500개의 폴딩 성공 잔기를 처리한 반면, 공개 구현은 2만 200개에 그쳐 잔기 정규화 처리량에서 2.90배의 성능 향상을 보였습니다. 오픈소스 구현은 29개 타깃에서 메모리 부족 오류가 발생했습니다.
상단 그림 3의 좌측 패널은 BioIR과 torch.compile 기반 오픈소스 구현의 모델 순전파 시간을 비교하며 BioIR이 달성한 더 낮은 모델 순전파 시간을 직접적으로 보여줍니다. 그림 3의 우측 패널은 BioIR과 오픈소스 구현이 제공하는 처리량을 비교하며, 여기서 처리량은 예측된 구조의 총 잔기 수를 할당된 GPU 시간으로 정규화한 값입니다. 그림 3의 우측 패널은 BioIR의 커널 레벨, 모듈 레벨, 파이프라인 레벨 최적화가 제공하는 속도 향상을 보여주며, 좌측 패널은 커널 레벨 및 모듈 레벨 구현이 제공하는 속도 향상을 보여줍니다.
Boltz2 가속과 마찬가지로, BioIR은 OpenFold2 및 OpenFold3와 같은 다른 생체 분자 공동 폴딩 모델에 대해서도 더 빠른 추론을 가능하게 합니다. 초기 버전의 BioIR은 NVIDIA 내부 버전의 OpenFold2-MM에 가속 모듈을 제공했으며, 이를 통해 3,100만 개의 단백질 복합체 구조를 포함하는 AFDB용 대규모 단백질 구조 예측을 구현했습니다.
그림 3의 1,000개 타깃 비교 벤치마크 결과를 8x H100 80GB HBM3 노드의 정격 전력 등급을 적용하여 유사한 100만 개 타깃으로 선형 추정했습니다(하단 그림 4 참조).
BioIR은 8-GPU 열 설계 전력(TDP) 등급 기준으로 공개 구현의 35MWh 대비 11MWh가 소요되고, 단일 노드 최대 전력 등급 기준으로는 공개 구현의 64MWh 대비 21MWh가 소요될 것으로 추정됩니다. 이는 IT 장비에 대한 폴딩 전용 추정치이며 실제 측정된 전력량이 아니므로 전력 사용 효율성(PUE)과 같은 데이터센터 오버헤드는 포함되지 않습니다.
통제된 비교 분석은 각 구현에 대해 동일한 입력 및 MSA를 사용하여 폴딩 처리량을 측정하며, MSA 생성, 전처리 CPU 할당, 스토리지, 데이터 전송, 재시도, 엔지니어링 오버헤드는 제외합니다. 완료된 시간당 구조 수, GPU 및 CPU 사용률, 피크 GPU 메모리, 완료율, 실패, 재시도 등을 포함한 엔드투엔드 파이프라인 성능 지표는 모델 순전파 지표와 명확히 구별하여 보고하세요.
문제 해결
- 단 1대의 GPU만 활성화됨: Ray, REPLICA 모드 설정, 1개 초과의 복제본 지정, 복수의 가용 GPU 존재 여부, 그리고 충분한 수의 독립 레코드가 확보되었는지 확인하세요.
- 프로세서 생성 시
ValueError발생:compute * num_gpus값이 가용 GPU 수를 초과하지 않는지 점검하세요. - 단백질 입력 실패: 필수 페어링되지 않은 A3M 파일 및 작업자가 접근 가능한 경로인지 확인하세요.
- GPU 대기 상태 발생: 복제본을 추가하기 전에 CPU 단계, CPU 할당 자원, 큐잉 상태, Ray 오브젝트 스토어 용량을 점검하세요.
- 추론 후 행 데이터 대기 발생: 작성기 동시성 및 대상 스토리지의 처리량을 점검하세요.
- Ray가 액터를 배치하지 못함: CPU, GPU, 메모리,
accelerator_type라벨을 확인하세요. - 단 1개의 레코드로 인해 전체 작업이 중단됨: 기본 패일 패스트(Fail-fast) 설정은
FoldingPredictionError를 발생시킵니다. 행 레벨에서 계속 실행하려는 경우에만should_continue_on_error=True로 설정한 후__inference_error__를 점검하세요.
시작하기
NVIDIA BioNeMo Inference Runtime을 살펴보고 대규모 구조 예측 워크플로우에 통합해 보세요: http://github.com/NVIDIA-BioNeMo/BioNeMo-Inference-Runtime
에이전틱 조율을 통해 신약 개발 워크플로우를 더욱 가속화하려면 NVIDIA BioNeMo Agent Toolkit을 확인하세요.
