[論文レビュー] BioNeMo Framework: a modular, high-performance library for AI model development in drug discovery

Peter C. St. John, Dejun Lin+91 more|arXiv (Cornell University)|Nov 15, 2024

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 batcherとbucket batch sampler）を組み込む。
• ウェブデータセットとLightningDatamoduleを統合してデータ処理を簡素化するWebDataModuleを提供する。

実験結果

リサーチクエスチョン

• RQ1 BioNeMoは標準のPyTorch/Transformers実装と比較して、どのようにして訓練のスループットを高めているのか？
• RQ2 多数のGPUで大規模な生体分子モデルをトレーニングする際のBioNeMoのスケーラビリティはどの程度か？
• RQ3 BioNeMoは実際に多様なデータ型（タンパク質配列、シングルセルデータ）とメモリを考慮したバッチ処理を効率的に扱えるか？
• 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 batchingは、MiDi/ベースライン手法と比較してほぼ均一なデータサイズ分布と極めて小さなパディングを生み出し（メモリパディングを削減）。
• コミュニティ主導の貢献とクラウド規模のデプロイ（AWS）を実証し、推論を高速化し、より大規模な探索ワークフローを可能にした。