NVIDIA BioNeMo における文脈並列化を用いた生体分子モデリングの拡張

何十年にわたり、計算生物学は還元主義的な妥協の下で運営されてきました。複雑な生物学的システムを単一の GPU の限られたメモリに収めるために、研究者たちはそれらを孤立した断片—単一タンパク質や小さなドメイン—へと分解せざるを得ませんでした。これにより文脈のギャップが生じ、GPU ハードウェアのメモリ制約のため、より大きなタンパク質や複合体はゼロショットで折りたたむことができませんでした。 現在、NVIDIA BioNeMo チームによる新しい [コンテキスト並列化 (CP)](https://arxiv.org/abs/2603.14806) フレームワークが構造生物学のメモリバリアを打ち破り、システム全体のホリスティックなモデリングを可能にしています。 この投稿では、標準的な Transformer とは異なるバイオ分子アーキテクチャにおいて CP を実現する方法について説明します。大規模なバイオ分子複合体をグローバルな文脈を犠牲にすることなくモデル化したい構造生物学者、計算化学者、または機械学習エンジニアの方は、ぜひお読みください。 この投稿で概説されたソリューションを使用するには、以下の準備が必要です: - AlphaFold3 や Boltz-2 などの幾何学的深層学習基盤モデルへの習熟度。 - PyTorch Distributed (DTensor) オペレーションとカスタム自動微分関数の理解。 - NVIDIA H100 または B200 GPU クラスターへのアクセス。このフレームワークは、エクサスケールタスクにおいてその相互接続帯域幅と Transformer Engine の加速に大きく依存しているためです。 詳細については、[Fold-CP: A Context Parallelism Framework for Biomolecular Modeling](https://arxiv.org/abs/2603.14806) をご覧ください。 ## Sharding a single large molecular system across multiple GPUs コンテキスト並列性（Context Parallelism）が存在しない場合、1,000〜3,000 残基を超えるような巨大な複合体を折りたたむには、還元主義的なアプローチが必要となります。これはシステムを実質的または計算的に管理可能な断片に分解する手法です。これらの方法は研究者が単一 GPU の厳格な VRAM（ビデオメモリ）の制限内に留まらせることを可能にしますが、しばしばグローバルな構造精度を犠牲にします。 巨大タンパク質に対する最も一般的な回避策は、配列をより小さく重なり合うセグメントにスライスすることです。局所的な二次構造が分割点全体で整合性を持つようにするためには、断片は大幅に重なり合う必要があります。この手法は長距離情報を破壊してしまいます。例えば、研究者は複合体全体にわたるアロステリーやシグナル伝達をモデル化することができません。 もう一つの一般的な回避策はチャンキング（chunking）です。これは物理的な配列の断片化とは異なり、推論中の VRAM を節約するためにモデルアーキテクチャ内で発生します。Boltz などのモデルは、大きな行列をより小さなタイルで処理するために積極的なチャンキングを採用しています。また、FastFold のような他の技術では、自動チャンキング（autochunking）を用いて動的にチャンキング戦略を調整し、ピークメモリ使用率を改善しています。詳細については、[FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours](https://ar5iv.labs.arxiv.org/html/2203.00854) をご覧ください。 これらの手法はすべて、特にトレーニング中にグローバルコンテキストの欠如という本質的な課題を抱えています。NVIDIA BioNeMo の CP（Context Parallelism）フレームワークは、単一の巨大な分子システムを複数の GPU にシャードすることでこれらの限界を克服します。従来のデータ並列性（Data Parallelism）が各 GPU に異なるタンパク質の折りたたみタスクを割り当てるのに対し、CP は単一の膨大なサンプルを GPU 間で分割します。 ## BioNeMo コンテキスト並列性の実装 NVIDIA BioNeMo の CP 実装は、GPU 間通信のための Torch distributed APIs を基盤として構築されています。このアーキテクチャはボトムアップアプローチで設計されており、低レベルの通信プロトコルから始まり、高レベルのモデル固有ワークフローへと段階的に拡張されます。本稿では、Boltz をコードベースの例として使用します。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/04/nvidia-bionemo-context-parallelism-framework-software-architecture.webp)*図 1. CP フレームワークのソフトウェアアーキテクチャ設計は、モジュールレベルの演算子をコンテキスト並列性（CP）対応とし、段階的にモデル固有のワークフローへと構築していくボトムアップアプローチを採用しています* システムの能力が GPU 数に対して線形に成長する線形キャパシティスケーリングを実現するために、このフレームワークは多次元シャード戦略を実装しています。これにより、生体分子の完全なグローバル状態を単一のデバイスが保持してしまうことが防がれ、CP のメモリ目標が損なわれることを回避します。異なるモジュールに対してカスタムのレゴブロックのような部品が構築されており、実装の効率化と他のアーキテクチャへの移植容易性が確保されています。 ## ペア表現の 2D タイルリング このフレームワークは、グローバル (N x N) 行列をブロックのグリッドに分割します。10,000 残基からなる複合体（1 億回の相互作用に相当）の場合、各 GPU は特定のサブブロックのみを管理します。これにより、デバイスあたりのメモリフットプリントが O(N^2) から O(N^2/P) に局所化されます。 ## 計算と通信の重なり合わせ このフレームワークは、非同期のピアツーピア転送とローカル計算を調整するための分散プリミティブを実装しています。ある GPU がローカル更新を計算している間、行および列リング内の隣接ノードとの間でデータの送受信も同時に行われます。生物学的な問題サイズが大きくなるにつれて、計算から通信への比率が向上し、大規模スケールにおいてシステムはより効率的になります。 ## 効率的な原子シーケンスローカルアテンション AlphaFold3 のシーケンスローカルアテンションでは、原子アテンションを 32 x 128 サイズのローカルアテンションウィンドウに制限し、これらのウィンドウスタックがバッチ処理でアテンション計算を実行します。前述の原子特徴量のタイルリングは、このウィンドウバッチベースのアトムアテンションの分散バージョンのために再分割する必要があります。NVIDIA チームは、ハローエクスチェンジに基づく分散プリミティブを実装して原子特徴量を分割し、その後のウィンドウバッチアテンションで GPU 間通信が不要となるようにしました。 ## トライアングル乗算に対するコンテキスト並列化の実装 以下の例は、CP（コンテキスト並列）対応のレイヤーを作成し、三角形乗算を分散処理する方法を示しています。 DistributedManager.initialize(device_type="cuda") manager = DistributedManager() size_ring = math.isqrt(manager.world_size) DistributedManager.create_grid_group({"dp": 1, "cp": (size_ring, size_ring)}) ring_comm = Ring2DComm(manager.group["cp"], manager.subgroups["cp"][0], manager.layout_subgroups["cp"]) x_dtensor, mask_dtensor = … layer_serial = TriangleMultiplicationOutgoing(size_input_embed) layer_serial.load_state_dict(layer_state_dict) layer_serial = layer_serial.to(manager.device) layer = DistributedTriangleMultiplication(Outgoing, layer_serial, manager.device_mesh_subgroups, ring_comm) result_dtensor = layer(x_dtensor, mask_dtensor) DistributedManager.cleanup() このコードは、DistributedManager を使用して正方形の 2D デバイスメッシュを設定することで CP 環境を初期化します。これは行方向および列方向の通信パターンが対称かつ効率的に保たれるようにするための特定のアーキテクチャ要件です。次に、Ring2DComm ハンドルがインスタンス化され、データブロックを連続ループで循環させるための特殊なピアツーピア通信を管理します。このリングアプローチは極めて重要であり、ローカル計算とデータ転送のオーバーラップを可能にすることで、pair 表現テンソルのサイズが単一の GPU のメモリ容量を超えないように保証します。 スクリプトの後半部分は、標準的なシリアルモデルから DTensor ベースのモデルへの移行を処理します。TriangleMultiplicationOutgoing などの標準レイヤーは、NVIDIA BioNeMo CP フレームワークに適応したアルゴリズムの分散版を実装する DistributedTriangleMultiplication にラップされる前に CPU にロードされます。入力を分散テンソル（DTensors）として処理することで、モデルは大規模な活性化テンソルをグリッド全体にシャードします。 ## 構造生物学におけるトークンスケーリングの解放 図 2 は、CP の導入により、生体分子アーキテクチャに対して [token](https://blogs.nvidia.com/blog/ai-tokens-explained/)（トークン）容量のスケーリング法則が解き放たれたことを示しています。Boltz の予測は最大約 20,000 トークンを 256 台の GPU で実行可能であり、NVIDIA H100 GPU では最大トークン長をスケーリングでき、NVIDIA B300 GPU では加速されたスケーリングが可能です。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/04/token-capacity-scaling-nvidia-bionemo-cp-framework.webp)*図 2. NVIDIA BioNeMo CP フレームワークを用いて達成したトークン容量のスケーリング* 追加のトレーニングや長いクロップ長でのファインチューニングを行わずに、チームは 3,605 個のアミノ酸残基からなる TTC7A/PI4KA/FAM126A/EFR3A(700–823) システムを折りたたみました。これは 4 つの鎖にまたがる構造であり、Boltz-2 のトレーニング用クロップサイズである 768 残基や、単一の GPU のメモリ容量を大幅に上回っています。CP（文脈並列処理）を活用することで、4 枚の NVIDIA H100 GPU を使用して、モデル内のコンテキストウィンドウ内ですべての長距離サブユニット間接触を維持したまま、5 つの構造サンプルを 5 分未満（1 サンプルあたり約 54 秒）で生成することが可能になりました。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/04/boltz-2-structural-visualization.webp)*図 3. TTC7A/PI4KA/FAM126A/EFR3A(700–823)複合体（3,605 残基）に対する Boltz-2 の予測。PI4KA（緑）、TTC7A（シアン）、FAM126A（マゼンタ）、そして本質的に無秩序な EFR3A C 末端（黄色）* 並行して、チームは CP フレームワークの開発に深く貢献した Rezo Therapeutics、Proxima、Earendil Labs などの協力者と連携し、この次世代の最前線を開拓しています。 Rezo Therapeutics は CP フレームワークを統合し、最大 6,500 残基にわたる大規模なタンパク質間相互作用（PPIs）の予測を実現しました。これにより、既知のタンパク質複合体の绝大多数に対する構造予測が可能となり、新規複合体の迅速な発見が実現しています。実際、公的データベースで高い信頼性を持つ PPI のみを用いて行われた予測と比較して、CP で解決された高品質な新規タンパク質複合体は 3 倍以上に富んでいます。 Proxima は、すべての原子からなる生成基盤モデル Neo に文脈並列処理（Context Parallelism）を組み込み、独自の大質量分光法プラットフォームによってマッピングされたプロテオーム全体にわたる治療的に重要な相互作用を構造的に解明するために、最大 4,000 トークンまでのアセンブリに対する推論を可能にしました。これにより、分子接着剤やその他の近接ベースの創薬アプローチの開発における創薬研究が支援されています。 Earendil Labs はこのフレームワークを独自の生体分子基盤モデルに統合し、従来は計算コストが高すぎて不可能だった複雑な多タンパク質システムをモデル化するために、入力シーケンス長を拡張することに成功しました。Earendil Labs は、文脈並列処理（CP）が、シーケンスの複雑さがスケールしても高い忠実度の構造予測を維持する可能性があり、次世代バイオ医薬品の発見期間を短縮できることを示しました。 ## 生体分子モデリングのための文脈並列処理の活用 初期の実証実験により、文脈並列処理（CP）が以前のメモリ制約を打破し、無制限サイズの構造モデル化を可能にすることが示されましたが、物理的な容量だけでは生物学的な正確さが保証されるわけではありません。現在のモデルは、小断片でトレーニングされているため、スケールした高忠実度の折りたたみ性能を発揮するのが難しい傾向があります。長距離相互作用の創発するロジックを正確に捉えるためには、より大きなクロップサイズでのファインチューニングが不可欠です。 この取り組みの一環として、この分野における主要なボトルネックの一つであるデータ不足の問題は、[AlphaFold Protein Structure Database](https://alphafold.ebi.ac.uk/) への貢献を通じて解決されています。NVIDIA の加速コンピューティングソフトウェアである [NVIDIA cuEquivariance](https://docs.nvidia.com/cuda/cuequivariance/) や [NVIDIA TensorRT](https://developer.nvidia.com/tensorrt) を活用して、大規模なホモマーおよびヘテロマー複合体の高スループット予測を AFCDB に追加することで、生物学におけるより大きなシステムを表すファウンデーションモデルのトレーニングに有用となる可能性のある合成複合体データの構築に向けた基盤が整えられています。 詳しくは、[Boltz CP code open-source documentation](https://github.com/NVIDIA-Digital-Bio/boltz-cp) をご覧ください。また、[Fold-CP: A Context Parallelism Framework for Biomolecular Modeling](https://arxiv.org/abs/2603.14806) もご参照ください。

NVIDIA BioNeMo における文脈並列化を用いた生体分子モデリングの拡張

背景や根拠まで確認しますか？

調べる

選ぶ

サイト