[논문 리뷰] BioNeMo Framework: a modular, high-performance library for AI model development in drug discovery

Peter C. St. John, Dejun Lin외 91명|arXiv (Cornell University)|2024. 11. 15.

Computational Drug Discovery Methods인용 수 11

한 줄 요약

BioNeMo Framework는 NVIDIA NeMo Megatron 위에 구축된 오픈 소스 모듈형 라이브러리로, 수백 대의 GPU에서 생체분자 AI 모델을 학습하고 확장하며, 처리량과 메모리 효율성의 검증된 향상을 제공합니다.

ABSTRACT

Artificial Intelligence models encoding biology and chemistry are opening new routes to high-throughput and high-quality in-silico drug development. However, their training increasingly relies on computational scale, with recent protein language models (pLM) training on hundreds of graphical processing units (GPUs). We introduce the BioNeMo Framework to facilitate the training of computational biology and chemistry AI models across hundreds of GPUs. Its modular design allows the integration of individual components, such as data loaders, into existing workflows and is open to community contributions. We detail technical features of the BioNeMo Framework through use cases such as pLM pre-training and fine-tuning. On 256 NVIDIA A100s, BioNeMo Framework trains a three billion parameter BERT-based pLM on over one trillion tokens in 4.2 days. The BioNeMo Framework is open-source and free for everyone to use.

연구 동기 및 목표

• 생체분자 AI 모델을 대규모로 효율적으로 학습하고 미세 조정할 수 있도록 한다.
• 기존 워크플로우와 통합되는 데이터 로더, 모델 및 유틸리티 등 모듈형 구성요소를 제공한다.
• 기준 PyTorch 구현 대비 처리량 및 확장성 향상을 입증한다.
• 전문 데이터 로딩(단백질 서열, 단일세포 데이터) 및 메모리 인지 배칭을 지원한다.
• 약물 발견 활용 사례를 위한 커뮤니티 기여 및 클라우드 규모 배포를 장려한다.

제안 방법

• 핵심 인터페이스를 bionemo-core에 두고 PyTorch와 Lightning 위에 구축한다.
• NVIDIA NeMo Megatron을 활용하여 대형 생체분자 BERT-스타일 모델(ESM-2, Geneformer)을 구성한다.
• 사전 구성된 학습, 미세 조정 및 추론을 위한 모듈형 하위 패키지(예: bionemo-esm2, bionemo-geneformer)를 제공한다.
• 단백질 서열 및 단일세포 데이터용 고성능 데이터 로더(BioNeMo-SCDL)를 구현한다.
• 그래프 및 가변 크기 입력의 메모리 사용 최적화를 위해 크기 인식 배치(size-aware batching, size-aware batcher 및 bucket 배치 샘플러를 포함)를 도입한다.
• WebDataset을 LightningDatamodule과 통합하기 위한 WebDataModule을 제공한다.

실험 결과

연구 질문

• RQ1 BioNeMo가 표준 PyTorch/Transformers 구현에 비해 더 높은 학습 처리량을 달성하는 방법은 무엇인가?
• RQ2 다수의 GPU에서 대형 생체분자 모델을 학습시킬 때 BioNeMo의 확장성은 어느 수준인가?
• RQ3 생체분자 서열(단백질 서열) 및 단일세포 데이터와 같은 다양한 데이터 유형 및 메모리 인식 배치가 실제로 효율적으로 처리될 수 있는가?
• RQ4 BioNeMo의 특화된 로더와 배치 전략으로 메모리 활용 및 데이터 로딩 성능에서 어떤 실질적 이점이 있는가?

주요 결과

• 256 NVIDIA A100에서 3B 매개변수 ESM-2 유사 pLM을 1조 개가 넘는 토큰에 대해 4.2일 만에 학습시켰다.
• BioNeMo는 단일 A100에서 650M 매개변수 모델에 대해 Hugging Face Transformers 대비 최대 1.47x의 단일 장치 처리량 향상을 달성하며, MFU가 59.2%로 기준선의 40.1% MFU보다 높다.
• 분산 학습 전반에서 BioNeMo의 3B 매개변수 모델은 256 GPUs에서 외삽된 단일 노드 처리량의 96.9%에 도달한다(16 A100에서 40% MFU, 256 GPUs에서 60% MFU).
• BioNeMo SCDL은 데이터를 메모리에 로드할 필요 없이 비교 가능한 AnnData 로더보다 1.1–2.75x 더 빠른 데이터 로딩을 제공한다.
• Bucket size-aware 배칭은 MiDi/기준 방법에 비해 거의 균일한 데이터 크기 분포와 매우 적은 패딩을 만들어 메모리 패딩을 줄인다.
• 커뮤니티 주도 기여와 클라우드 규모 배포(AWS)를 통해 더 빠른 추론과 더 큰 탐색 워크플로를 가능하게 함을 입증했다.