一流の科学者はどうやって AI を習得するのでしょうか? DeepSpeed4Science: 高度な AI システム最適化技術を使用して科学的発見を可能にする

今後 10 年間で、ディープラーニングは自然科学に革命をもたらし、自然現象をモデル化して予測する能力を高めると考えられます。

これは科学的探究の新しい時代の到来を告げるものとなり、医薬品開発から再生可能エネルギーに至るまでのさまざまな分野で大きな進歩をもたらす可能性がある。

これに対応して、Microsoft の DeepSpeed チームは、AI システムの技術革新を通じて、現場の専門家が今日の最大の科学的謎を解明できるよう支援することを目指した DeepSpeed4Science という新しいプログラムを立ち上げました。

DeepSpeed システムは、Microsoft が開発した業界をリードするオープンソース AI システム フレームワークであり、さまざまな AI ハードウェア上でのディープラーニングのトレーニングと推論に前例のないスケールと速度を提供します。

図 1: DeepSpeed4Science アプローチの概要: 科学的発見を加速し、その複雑さに対処するためにカスタマイズされた AI システム テクノロジの開発。

図 1 は、新しい DeepSpeed4Science イニシアチブに対する基本的なアプローチを示しています。

DeepSpeed4Science は、DeepSpeed の現在のテクノロジー ソリューション (トレーニング、推論、圧縮) を基礎テクノロジー イネーブラーとして活用することで、汎用大規模言語モデル (LLM) を加速するために使用される一般的なテクノロジー アプローチを超えて、科学的発見を加速してその独特の複雑さに対処するようにカスタマイズされた AI システム テクノロジーのセットを作成します。

このブログでは、DeepSpeed4Science が構造生物学研究における 2 つの主要な AI システムの課題の解決にどのように役立つかを紹介します。

（１）エボフォーマー中心のタンパク質構造予測モデルにおけるメモリ爆発問題を解決し、

（２）パンデミックを引き起こす病気の進化をより深く理解するために、AIモデルに長期シーケンスのサポートを提供する。

当初の主な協力者

DeepSpeed4Science の新しいシステム テクノロジーは、科学の限界を押し広げ、AI 主導の科学的発見を可能にする多くの象徴的なモデルで使用できます。

現在、DeepSpeed4Science は、Microsoft Research AI4Science、Microsoft WebXT/Bing、米国エネルギー省国立研究所、および複数の大学のいくつかの主要な科学モデルをサポートできることを誇りに思っています。

社内パートナー

科学的基盤モデル (SFM)、Microsoft Research AI4Science

図2: 科学的基礎モデル (SFM) とその現在の調査: 分散グラフォーマー

科学基盤モデル (SFM) は、さまざまな入力、複数の科学分野 (薬物、材料、生物学、健康など)、および計算タスクをサポートし、自然科学の発見をサポートするための統合された大規模な基盤モデルを作成することを目的としています。

DeepSpeed4Science との提携により、SFM チームには、Distributive Graphomer などの新しい生成 AI 手法などのプロジェクトに関する継続的な研究をサポートするための新しいトレーニングおよび推論手法が提供されます。

ClimaX、マイクロソフト リサーチ AI4Science

図3: ClimaXは、幅広い気象および気候モデリングタスクを実行するために設計された最初の基礎モデルです。

気候は変化しており、異常気象の発生頻度が増加しています。悪影響を軽減するためには、これらの出来事がどこで発生するかを予測することがますます重要になります。

ClimaX は、幅広い気象および気候モデリング タスクを実行するために設計された最初の基礎モデルです。さまざまな変数と解像度を持つ多くのデータセットを吸収し、天気予報の精度を向上させることができます。

DeepSpeed4Science は、非常に大規模な高解像度画像データ (数十から数百ペタバイトなど) と長いシーケンスを処理しながら、より大きな基本モデルを効率的に事前トレーニング/微調整するための、ClimaX の新しいシステム サポートと加速戦略を作成しています。

AI を活用した第一原理分子動力学 (AI2MD)、Microsoft Research AI4Science

図 4: 分子動力学シミュレーションの 100 万ステップ: RBD とタンパク質の相互作用とタンパク質阻害剤。

このプロジェクトでは、AI 駆動型力場モデルを使用して大規模な (100 万個の原子) 分子システムのダイナミクスをシミュレートし、古典的な分子動力学の効率性とスケーラビリティを維持しながら、第一原理計算の精度を近似します。シミュレーションは化学的に意味のあるイベントを観察するのに十分な長さの軌道を生成するのに十分な効率性を持っています。

通常、このプロセスには数百万、あるいは数十億の推論ステップが必要になります。これは、グラフ ニューラル ネットワーク (GNN) + LLM モデルの推論速度を最適化する上で大きな課題となりますが、DeepSpeed4Science はこれに対して新しい加速戦略を提供します。

Microsoft 天気、Microsoft WebXT/Bing

図 5: Microsoft の降水量予報 (今後 4 時間、4 分ごとに予報)。

Microsoft Weather は、1 時間に複数回更新される正確な 10 日間の世界予報など、ユーザーがライフスタイル、健康、仕事、活動についてより適切な判断を下せるよう、正確な気象情報を提供します。

以前、Microsoft Weather は DeepSpeed テクノロジを活用して、マルチ GPU トレーニング環境を高速化していました。

現在、DeepSpeed4Science は Microsoft WebXT Weather チームと協力して、Microsoft の天気予報サービスの最新機能と改善をさらに強化しています。

外部協力者

DeepSpeed4Science の旅は、構造生物学研究のための画期的な LLM ベースの AI モデル 2 つから始まりました。コロンビア大学の OpenFold はオープンソースの高忠実度タンパク質構造予測モデルであり、アルゴンヌ国立研究所の GenSLM は SARS-CoV-2 (COVID-19) ゲノムの進化を学習するための ACM ゴードン・ベル賞を受賞した言語モデルです。

このリリースで取り上げられているのは、今日の AI 主導の構造生物学研究が直面している 2 つの一般的な AI システムの課題です。次のセクションでは、DeepSpeed4Science がどのようにしてこれらの科学的研究を可能にするかについて説明します。

さらに、DeepSpeed4Science は最近、より幅広い科学モデルをサポートするように範囲を拡大しました。

たとえば、アルゴンヌ国立研究所と協力して、Aurora Exascale システムで 1 兆パラメータの科学モデルをトレーニングする作業では、DeepSpeed4Science テクノロジがこの重要なミッションに必要なパフォーマンス要件とスケーラビリティの達成に役立ちます。

さらに、オークリッジ国立研究所および国立がん研究所 (NCI) とのがん監視に関する協力を通じて、DeepSpeed4Science は、MOSSAIC プロジェクトで使用するために、構造化されていない臨床テキストから高い忠実度で情報を抽出し、分類するのに役立ちます。

ブルックヘブン国立研究所は、クリーンエネルギー研究のためのより現実的なシミュレーションデータを生成するために、LLM を使用した大規模デジタルツインモデルの開発をサポートするために DeepSpeed4Science テクノロジーも採用する予定です。外部協力者とその科学的ミッションの詳細については、deepspeed4science.ai をご覧ください。

協力展

デモンストレーション（I）：DeepSpeed4Scienceは、DS4Sci_EvoformerAttentionを通じてEvoformer中心の構造生物学モデルのメモリ爆発問題を解消します。

図6: トレーニング中のPDBチェーン7B3A_AのOpenFoldの予測

OpenFold は、DeepMind の AlphaFold2 をオープンソース コミュニティで再現したもので、新しいデータセットで AlphaFold2 をトレーニングしたり微調整したりすることができます。

研究者はこれを使用して、AlphaFold2 を最初から再トレーニングし、新しいモデル パラメーターのセットを生成し、AlphaFold2 の初期トレーニング段階を研究し (図 6)、新しいタンパク質折り畳みシステムを開発しました。

図 7: OpenFold における Multiple Sequence Alignment (MSA) Attention カーネル (バイアスを含む) のバリアントをトレーニングするためのピーク メモリ要件。 (左) AlphaFold2 の EvoformerAttention を使用したオリジナルの OpenFold 実装。このようなタイプのタンパク質構造予測モデルでは、トレーニング/推論におけるメモリ爆発の問題がよく発生します。最先端の FlashAttention では、このような Attention バリアントを効果的にサポートできません。 (右) DeepSpeed4Science の新しいソリューションである DS4Sci_EvoformerAttention は、モデルの品質に影響を与えることなく、OpenFold のピークトレーニングメモリ要件を大幅に削減します (最大 13 倍)。

OpenFold はパフォーマンスとメモリの最適化のために最先端のシステム技術を使用していますが、AlphaFold2 を最初からトレーニングするには依然として計算コストがかかります。現在のモデルは、わずか 9,300 万のパラメータと小さいサイズですが、非常に大きな中間メモリを必要とするいくつかの特殊な Attention バリアントが含まれています。

標準的な AlphaFold2 トレーニングの「微調整」フェーズでは、これらのバリアントの 1 つだけで半精度で 12 GB を超えるテンソルが生成され、ピーク時のメモリ要件は同じサイズの言語モデルの要件をはるかに超えます。

アクティベーション チェックポイントや DeepSpeed ZeRO 最適化などの技術を使用しても、このメモリ爆発の問題により、トレーニング可能なモデルのシーケンスの長さと MSA の深さが依然として大幅に制限されます。

さらに、近似戦略は、図 7 の左側 (オレンジ) に示すように、モデルの精度と収束に大きな影響を与えると同時に、メモリの爆発的な増加につながる可能性があります。

構造生物学研究におけるこの一般的なシステム課題 (タンパク質構造予測や平衡分布予測など) に対処するために、DeepSpeed4Science は、このような科学モデルで広く見られる注意バリアント (EvoformerAttention など) 向けにカスタマイズされた精密注意カーネルを設計することで、このメモリ効率の問題に対処します。

具体的には、より広範な生物学研究コミュニティが使用するための高品質の機械学習モジュールとして、洗練された融合/マトリックス チャンキング戦略と動的メモリ削減方法で構成される、メモリ効率の高い DS4Sci_EvoformerAttention カーネルのセットを設計します。

OpenFold に組み込まれたこれらのカスタム カーネルは、トレーニング中に大幅な高速化を実現し、モデルのトレーニングと推論の両方のピーク メモリ要件を大幅に削減します。

これにより、OpenFold はより長いシーケンスを使用し、より広範囲のハードウェア上で、より大規模で複雑なモデルを実験できるようになります。この技術に関する詳細情報は、こちらをご覧ください。

デモンストレーション（II）：DeepSpeed4Scienceは、体系的かつアルゴリズム的なアプローチを通じて、ゲノムベースのモデル（GenSLMなど）の長いシーケンスのサポートを提供します。

図 8: GenSLM: 2022 ACM ゴードン・ベル賞を受賞した COVID ゲノムモデル (GPT-NeoX に基づく 25B/33B モデル)。これは、SARS-CoV-2 ゲノムの生物学的意義を記述する潜在空間を学習するために使用されます。この GIF は、重要なタンパク質ファミリーであるリンゴ酸脱水素酵素の潜在空間の投影を示しており、配列の長さや GC 含有量 (核酸のグアニンとシトシンとアデニンとチミンの比率。DNA 鎖の熱耐性を測定します) などの重要な特徴によって色分けされています。

アルゴンヌ国立研究所の 2022 ACM ゴードン・ベル賞を受賞したゲノムモデル GenSLMs は、大規模言語モデル (LLM) を使用してゲノムデータでトレーニングすることで、SARS-CoV-2 (COVID-19) ゲノムの進化を学習できます。これは、パンデミックを引き起こすウイルス、特にSARS-CoV-2の新しい変異体を特定し分類する方法を変えることを目的としています。

GenSLM は、他の予測タスクに一般化できる最初のゲノムベースのモデルです。潜在空間を適切に理解することで、GenSLM はウイルス配列だけでなく新しい領域に取り組むことができ、細菌病原体や真核生物をモデル化する機能を拡張することができます (機能、経路のメンバーシップ、進化関係などを理解するなど)。

この科学的目標を達成するために、GenSLM や類似のモデルでは、トレーニングと推論の両方で非常に長いシーケンスのサポートが必要であり、これは FlashAttention などの汎用 LLM の長いシーケンス戦略の範囲を超えています。

DeepSpeed4Science の新しい設計により、科学者は大幅に長いコンテキスト ウィンドウを使用してモデルを構築およびトレーニングできるようになり、これまでアクセスできなかった関係を探索できるようになりました。

図 9: 異なるスケールの異なるフレームワークでサポートされる 2 つの GenSLM モデルの最大シーケンス長。 NVIDIA DGXを使用すると、各ノードには8つの40G A100 GPUが搭載されます。

具体的にはシステムレベルでは、長いシーケンスのサポートやその他の新しい最適化を含む最新の Megatron-DeepSpeed フレームワークをリリースしました。

科学者は、新たに追加されたメモリ最適化技術 (アテンション マスクの非同期処理や位置コード分割など)、テンソル並列処理、パイプライン並列処理、シーケンス並列処理、ZeRO ベースのデータ並列処理、モデル状態の非同期処理の相乗的な組み合わせにより、より長いシーケンスで GenSLM などの大規模な科学モデルをトレーニングできるようになりました。

図 9 は、新しいバージョンでは、以前の Megatron-DeepSpeed バージョンと比較して、GenSLM の 25B モデルと 33B モデルの最長シーケンス長がそれぞれ 12 倍と 14 倍に増加していることを示しています。

この新しい Megatron-DeepSpeed フレームワークは、サポートされるシーケンス長の点でも NVIDIA の Megatron-LM を大幅に上回っています (それぞれ 25B モデルと 33B モデルで最大 9.8 倍と 9.1 倍)。

たとえば、アルゴンヌチームの GenSLMs 25B モデルは、元々 64 個の GPU で 42K のシーケンス長を持っていましたが、現在は 512K のヌクレオチド シーケンスでトレーニングできます。これにより、精度を犠牲にすることなく、モデルの品質と科学的発見の範囲が大幅に向上します。

相対位置エンコーディング技術などのアルゴリズム戦略を好むドメイン科学者のために、この新しいバージョンではそれらも統合されています。

Microsoft DeepSpeed Groupの公式Zhihuアカウントより転載：

ディープスピード