NVIDIA BioNeMo エージェント・ツールキットによるエンドツーエンド共折りたたみ性能の加速

OpenFold3 などのモデルを用いた生体分子構造予測や共折りたたみ解析は現在主流となり、創薬やタンパク質設計を支える大規模なワークロードとして定着しています。こうした処理を AI エージェントがエンドツーエンドで駆動するケースも増えています。エージェントがこのパイプラインを円滑に実行するためには、各ステップの高速化とスケーラビリティが不可欠です。具体的には、Multiple Sequence Alignment (MSA) の生成、共折りたたみ推論、サービス提供、そしてマルチ GPU によるスケールアウトなど、あらゆる工程でボトルネックが生じれば全体の処理能力が制限されてしまいます。

創薬ワークロードの要であるバーチャルスクリーニングや大分子集合体の予測においては、速度とメモリ効率性が極めて重要です。バーチャルスクリーニングでは、数百万から数十億もの化合物を 1 つまたは数種類のタンパク質ターゲットに対してスクリーニングします。共折りたたみモデルは最も精度の高い構造予測をもたらすことがありますが、実行コストが高いため、大規模な化合物ライブラリを対象としたスクリーニングには実用的ではありません。ここで NVIDIA のアクセラレーションが鍵となります。これにより、OpenFold3 や関連手法を大規模な化合物ライブラリのスケールで展開することが可能になります。

複数のタンパク質と数千アミノ酸残基を関与する大規模な分子集合体の予測においても、速度は極めて重要です。共折りたたみモデルのランタイムは残基数に対して立方スケーリングするためです。しかし、さらに大きな課題はメモリ使用量です。単一の GPU メモリには限界があり、一度に予測できる複合体のサイズに厳しい上限を設けてしまいます。メモリ要件を削減し、予測タスクを複数の GPU に分散させる手法は、今日では実現不可能な質的に新しい応用を可能にするでしょう。

NVIDIA は、構造予測と共折りたたみのワークフローにおける各ステップの加速と効率化のためのツールを開発しました。NVIDIA BioNeMo Agent Toolkit を利用すれば、エージェントが生物学や化学のワークフローを加速するために必要なツールにシームレスにアクセスできます。本稿では、NVIDIA B300 および H100 GPU における各段階の加速内容について解説し、それらの段階がどのようにしてエージェントを通じて実行されるかを示します（図 1 を参照）。

図 1．NVIDIA が加速した構造予測ワークフロー：GPU MSA、cuEquivariance、最適化された推論、Fold-CP。これらはすべて BioNeMo Agent Toolkit によってオーケストレーションされています。

GPU MSA で MSA のボトルネックを解消する

共折りたたみモデルにおいて、MSA（多重配列アラインメント）の構築は従来、CPU 依存のステップであり、実測時間の大半を占めることが多かった。MMseqs2-GPU は相同性検索を NVIDIA GPU 上で実行することでこのボトルネックを解消し、NVIDIA Hopper および Blackwell アーキテクチャ上でシーケンス長に応じたスケーラビリティを実現する。

最新の GPU アクセラレーション版では、Hopper と Blackwell に特化した最適化が追加された。これには、NVIDIA Grace システム上での GPU メモリを超える大規模データベース検索の効率的なサポートや、CUDA 13.2 から利用可能になった改善された Blackwell DPX 命令によるさらなる高速化が含まれる available from CUDA 13.2。これらの GPU への貢献は、主要な MMseqs2 リポジトリ にアップストリームされ、コミュニティ全体がその高速化の恩恵を受けられるようになった。

MSA Search NIM は MMseqs2-GPU を採用しており、単一の L40S GPU 上で CPU の JackHMMER と比較して最大 177 倍の高速なアラインメントを実現する Nature Methods paper が報告されている。ベンチマークでは、H100 および B300 GPU 上で 10k トークンを超える領域でもスムーズにスケーリングする（以下の Figure 2 を参照）。MSA Search NIM は直接呼び出し可能であり、セルフホストも可能で、エージェントワークフロー内のツールとしてラップすることもできる。

注釈：ここでは、build.nvidia.com にホストされた NIM API エンドポイントの呼び出し方について解説しますが、大規模なワークロードを処理する場合は、セルフホスト型の NIM を使用することが推奨されます。build.nvidia.com のエンドポイントは大量のリクエストを捌くように設計されていないため、タイムアウトが発生する可能性があります。

npx skills add NVIDIA-BioNeMo`/bionemo-agent-toolkit` --skill msa-search-nim --agent claude-code

NVIDIA_API_KEY=<key from build.nvidia.com> を環境変数として設定する。 （原文の技術表記: export）

"Build an MSA for the sequence in target.fasta with the MSA Search NIM."

cuEquivariance と OpenFold3 NIM で SOTA の速度を実現

「cuEquivariance」(https://github.com/NVIDIA/cuEquivariance) は、原子レベルのモデリング向けの幾何学習用プリミティブを備えた CUDA-X ライブラリです。タンパク質共折りたたみ（co-folding）において中心的な役割を果たす Triangle Attention、Triangle Multiplication、Attention Pair Bias の各カーネルに対して高速化版を提供しています。B300 上では、レイテンシが最大で約 3 倍短縮されます（以下の Table 1 を参照）。

表 1. cuEquivariance を使用した場合と使用しない場合の OpenFold3 のフォワードパスレイテンシ比較

cuEquivariance カーネルは、OpenFold3（オプション依存関係として提供）、OpenFold2、RosettaFold3、Protenix、Boltz といったオープンソースモデルに直接統合されています。

これらの加速機能は各 OSS モデルに上流（upstream）されているため、研究者は NVIDIA GPU で既存のモデルを実行するだけで自動的に速度向上を得られます。また、cuEquivariance カーネルにより最大シーケンス長を約 5.9k トークンまで拡張できますが、PyTorch では約 1.5k〜2.5k トークンを過ぎるとメモリ不足になります。

図 3. H100 および B300 上で PyTorch を上回る、cuEquivariance によるレイテンシ低下と最大シーケンス長の拡張

cuEquivariance をベースに、OpenFold3 NIM はさらに推論最適化を施し、その効果が積み重なることで（図 4 参照）、単一の B300 で最大約 6,400 のシーケンス長に対応可能になりました。これらの追加の高速化機能は NIM を通じて提供されており、開発者は SOTA（最先端）の性能をそのまま利用するために、NIM エンドポイントを直接呼び出すか、エージェントワークフローに組み込むことができます。

図 4. 最適化を施した OpenFold3 NIM の H100 および B300 での性能。それぞれ 3 倍以上、4 倍以上の高速化を実現し、単一の B300 で最大 6,400 トークンの折りたたみを可能にしています。

npx skills add NVIDIA-BioNeMo/bionemo-agent-toolkit --skill openfold3-nim --agent claude-code

export NVIDIA_API_KEY=<key from build.nvidia.com>

"Fold target.fasta with OpenFold3 using the MSA from the previous step; return the ranked structures with confidence scores."

GPU 1 台を超えて Fold-CP で拡張する

従来、単一 GPU のメモリ容量の制約により、共折りたたみモデルは数千アミノ酸程度に制限されていました。しかし NVIDIA B300（Blackwell Ultra）では、大容量 HBM と Blackwell 世代の高い効率性に加え、前述の cuEquivariance や高度な推論最適化を組み合わせることで、モデルの変更なしにこの上限を大幅に引き上げることが可能です。

多くの場合、単一デバイスへの依存だけでは不十分なケースがあります。Fold-CP は新しい並列化手法を導入し、デバイスあたりのメモリ必要量が *O(N²/P)* でスケーリングします（N はトークン数、P は GPU 数）。これにより、Boltz-2 モデルを 64 枚の B300 GPU で実行する際、最大 32,000 トークンを処理可能となり、単一 GPU の制限と比較して約 12 倍の性能向上を実現しました。

Fold-CP を試すには、エージェントに Boltz-CP のコードベースを指し示して、マルチ GPU 推論を実行するよう指示するだけです。

git clone https://github.com/NVIDIA-Digital-Bio/boltz-cp && cd boltz-cp

torchrun --nnodes 1 --nproc_per_node 4 \

src/boltz/distributed/main.py predict /path/to/preprocessed_data \
  --out_dir ./predictions \
  --size_dp 1 --size_cp 4
``` （原文の技術表記: ` `、` `）

--recycling_steps 3 --sampling_steps 200 --diffusion_samples 5

## エンドツーエンドの共折りたたみパイプラインを加速する

構造予測のパフォーマンスは、もはやエンドツーエンドのシステム全体の問題となっています。OpenFold3 の実用的なワークフローには、MSA（多重配列アライメント）の生成、共折りたたみの推論、デプロイ、そして生物学的複合体をどの程度の規模までモデル化できるかを決定するメモリ制限が含まれます。

NVIDIA はこのワークフローの各層を加速します。MSA Search NIM は相同性検索を 177 倍高速化し、cuEquivariance と OpenFold3 NIM は Blackwell GPU 上で推論レイテンシを最大 4 倍短縮します。また Fold-CP では、コンテキスト並列化によって共折りたたみを単一 GPU の範囲から超え、64 枚の NVIDIA B300 GPU を用いて 32,000 トークン規模の複合体への拡張が可能となる道筋を示しています。

これらのツールを組み合わせることで、構造予測がより高速化され、スケーラビリティが高まり、エージェントによる発見ワークフローへの統合も容易になります。これにより、研究者は単なるモデルの予測結果から、より大規模で実用的な生物学的システムへと研究を進めることが可能になります。

## これらの加速技術がもたらす可能性向上

これらの改善により、これまで手が届かなかった構造生物学の問題群に挑戦することが可能になりました。バーチャルスクリーニングにおいては、共折りたたみ推論の高速化により、従来は創薬キャンペーンの最終段階に限られていた構造ベース手法を、より早期かつ大規模なライブラリスケールで適用できるようになります。これによって、パイプラインを進む候補化合物の質と多様性が向上します。

大規模な生体分子集合体においては、B300 による拡張された単一 GPU 容量と、コンテキスト並列 Fold-CP フレームワークの組み合わせが、モデル化可能な対象範囲そのものを変えます。リボソームやスプライソソーム、あるいは大規模なシグナル伝達複合体といったスケールの複雑系は、これまで共折りたたみモデルにとって構造的に扱いにくい課題とされてきました。しかし、これらの加速技術によって状況が変化しつつあります。

具体的な例を挙げると、約 10,000 アミノ酸残基からなる複合体（大腸菌のリボソームと同程度の規模）の折りたたみを予測する場合、従来の単一 GPU ではコストが高すぎて現実的ではなく、あるいはそもそも不可能でした。しかし、B300 で Fold-CP を用いれば、マルチ GPU ノード上でそのような予測が現実的なものとして実行可能になります。

構造化生物学が計算機上で問いかけることができる質問の本質的な変化であり、単なる処理速度の向上ではありません。これらのツールをオープンソース統合やエージェント型 API を通じて利用可能にすることで、計算による予測が生物学的な洞察へと、ひいては新たな医薬品へと転換されるまでのスピードを加速させることができます。

## 始め方
NVIDIA BioNeMo Agent Toolkit を用いた高速化された OpenFold3 ワークフローを試すには、以下のツールから始めましょう：

- MSA Search NIM: [https://build.nvidia.com/colabfold/msa-search](https://build.nvidia.com/colabfold/msa-search)

- cuEquivariance: [https://github.com/nvidia/cuequivariance](https://github.com/nvidia/cuequivariance)

- OpenFold3 NIM: [https://build.nvidia.com/openfold/openfold3](https://build.nvidia.com/openfold/openfold3)

- 大規模な複合体向けの Fold-CP: [https://github.com/NVIDIA-BioNeMo/boltz-cp](https://github.com/NVIDIA-BioNeMo/boltz-cp)

## お礼
ベンチマークとツールの開発に尽力いただいた NVIDIA チームの皆様に感謝いたします：Franco Pellegrini 氏、Lalit Vaidya 氏、Duc Tran 氏、Tien Pham 氏、Maximilian Stadler 氏、Alejandro Chacon 氏、Quan Vu 氏、Simon Chu 氏、Brian Roland 氏、Dejun Lin 氏、Joseph Chang 氏、Hoa La 氏、Jonathan Mitchell 氏、Vishanth Iyer 氏、Timur Rvachov 氏、Christian Dallago 氏、Christian Hundt 氏。また、OpenFold および OMSF チームの皆様には、協力と貢献に対して心より感謝申し上げます。