世界のスーパーコンピュータとそれを動かす人々


10月 18, 2021

Anton 3は「火を噴く」分子シミュレーションの怪物

HPCwire Japan

John Russell

ライフサイエンスの計算機研究は、長い間、従来のHPCよりも統計と並列処理が主流であった。例えば、遺伝子配列の解析やバリアントコールなどだ。メカニズムのシミュレーションやモデリングの役割ははるかに小さいものだったが、それも変わりつつある。例外として、D.E. Shaw Research(DESRES)が設計・製造したAntonシリーズのスーパーコンピュータは、分子動力学モデリングに特化したシステムだ。先日のHot Chipsでは、DESRES社が「Anton 3」のカバーを外し、設計チームが「分子動力学シミュレーションのための火を噴く怪物」と呼んだことを紹介した。

 
   

Anton 3は、Anton 2からの大きな進歩であるだけでなく、AIとHPCの計算能力を融合させてスーパーコンピュータに搭載するという現在のトレンドとは異なる興味深いものである。Antonは、後者を真正面から見据えており、分子システムのメカニズムの詳細を解明することで、複雑な生物学を解明し、より良い薬や治療法を開発することを目指している。そのためには、レイテンシーが大きな課題となる。DESRES社は、静電力と結合エネルギーの計算に特化したコアを搭載した強力なASICと、合理的な通信手段を採用した。

このアプローチは、決して安価ではないが、うまく機能している。今年の夏の初め、DESRES社は最初の「フルサイズ512ノード」のAnton 3マシンを完成させた。このマシンは、4つのラックにまたがり、水冷式で、基本的なトポロジーは3Dトーラスである。基本的なトポロジーは3次元のトーラスである。目の前の課題のためには何よりも大切なことだ。

DESRES社のAdam Buttsは、講演の中で「Anton 3の最高性能のライバルは、比較的後期モデルのGPUであるNvidia A100です。」と述べている。「自社で最適化したMDコードであるDesmondをチップ・バイ・チップで動作させると、Anton 3はA100の約20倍の速度を示します。これは、Nvidia社を非難しているわけではありません。A100は非常に優れた計算エンジンです。むしろこれは、特化することで得られる多大な利益を示しているのです。」

 

一般的に、分子シミュレーションの評価基準は、分子システムのリアルタイムな相互作用をどれだけシミュレートできるか、そしてどれだけ大きな分子システムをシミュレートできるかということだ。

「Anton 3は、1つのノードが科学的作業に役立つサイズになりそうな最初のAntonマシンです。例えば、冒頭で紹介した113,000個の原子を持つACE2(CoV2スパイクの標的となるタンパク質)のシミュレーションは、シングルノードのAnton 3にすっぽりと収まり、GPUよりも1桁以上高い性能を発揮します。」とButtsは語る。

「Anton社の特徴は、ノード数が増えるほど性能が向上することです。Anton 3のピーク時の性能は、原子数が約10万個以下の場合、わずか64台のノードで1日あたり200マイクロ秒を超えます。つまり、1週間でミリ秒単位のシミュレーションが可能なのです。512ノードのフルサイズ機では、原子数100万個以上のシミュレーションで1日100マイクロ秒以上を維持し、原子数5,000万個以上のシミュレーションにも対応しています。」とButtsは語る。

また、少なくとも今のところは、複数のGPUを使ってスケールアップしても意味がないと付け加えている。「システムが非常に大きくなるまで、複数のGPUにシミュレーションを分割しても、現状ではインターノード通信のコストに見合うだけの利益は得られません。」と述べている。

 

HPC関係者の多くにとって、DESRESの誕生秘話はよく知られている。コロンビア大学の教授からヘッジファンドマネージャーに転身したDavid E Shawは、初期の「quant」であり、ウォールストリートでの華々しい成功により、その才能を生物医学研究に向けることができた。彼は2001年にD.E. Shaw Researchを設立し、自分がチーフサイエンティストを務めた。その結果、分子モデリングを目的とした初のスーパーコンピューター「Anton 1」(~2008年)と、高速MDシミュレーション用ソフトウェアパッケージ「Desmond」が誕生した。続いてAnton 2(~2014年)では、「数百万個の原子を持つシステムに対して、1日あたり数マイクロ秒の物理的時間のシミュレーション速度を達成した初めてのプラットフォーム(IEEE抄録)」となっている。

Antonシリーズの核となるのはASICエンジンであり、世代を重ねるごとに段階的な進歩と劇的な革新を遂げてきた。下のスライドに示すように、Antonの容量は着実に増加している。今回の講演では、「Anton 2」のASICのアーキテクチャを検証した上で、「Anton 3」の改良点を紹介した。DESRES社もSC21で論文[i](Anton 3: Twenty Microseconds of Molecular Dynamics Simulation Before Lunch)を発表する予定である。おそらく、SC21での講演もあると思う。

 

大規模な分子システムのシミュレーションは、当然ながら計算上の困難を伴う。そのため、近似法がよく使われる。原子間のエネルギーポテンシャルを記述した物理ベースのモデルに従って、数千から数百万個の原子間の力が計算される。これらの力は、わずか数フェムト秒の離散的なタイムステップで統合され、原子の新しい位置と速度が決定される。このプロセスを何十億回も繰り返すことで、関心のあるタイムスケールの軌道が生成されるのである。

ButtsがDESRES/Antonのアプローチについて説明しているのが興味深い。

「まず、電荷を帯びた粒子間の静電力について説明します。シミュレーション内の原子には、モデルのパラメータである電荷が割り当てられており、粒子のペア間の力は、電荷の積を分離度の2乗で割った値に比例します。すべてのペアの粒子間の相互作用を計算するとスケールが小さくなるため,力の合計は,カットオフ半径までの明示的なペアの相互作用の合計と,その先の電荷からの遠方への寄与とに書き換えられます.後者の寄与は畳み込みで表現でき、フーリエ空間での乗算によりグリッド上で効率的に計算できます。」と述べている。

「もちろん、力場には静電相互作用以外の要素もあります。量子力学的な効果は、化学結合が意味する位相的なつながりを表す力や、原子のペアの間に働くヴァンデルワールス力によって近似されますが、距離が離れるとすぐに落ちてしまうため、基礎となる量子力学方程式を解くよりも、範囲を限定した計算の方が十分に正確です。このようなモデルは、長い時間スケールのシミュレーションの問題を、難解なものから馬鹿げたものに変えてくれます。」とButtsは語っている。

Anton社が火を噴く仕組み

Antonシステムのパワーは、専用のASICに依存している。Anton 2のASICは、主に2つの計算タイルで構成されてる(下のスライドを参照)。「1つはフレックスタイルで、共有メモリに接続されたプライベートキャッシュを持つ4つのコアからなる一般的なマルチコアプロセッサのようなものです。MDコードで頻出するベクトル幾何学計算のための特別な機能を持つことから、これらをジオメトリコア(GC)と呼んでいます。」と述べている。ディスパッチユニットが細かい粒度の同期を処理し、ネットワークインターフェースがフレックスタイルを他のチップやマシンに接続する。

ハイスループット・インタラクション・サブシステム(HTIS)と呼ばれるもう一つのタイルは、「ペアワイズ・ポイント・インタラクション・モジュール(PPIM)の配列で占められています。各PPIMには、相互作用する原子のペア間の力を計算する一対のアンロールされた算術パイプライン(PPIP)が含まれています。PPIPは、グリッドポイント上の電荷を拡散し、グリッドの力を原子に補間することで、遠方の力の計算にも参加しています。」

「相互作用制御ブロック(ICB)は、PPIMアレイへの位置のストリーミングと、蓄積された力のアンロードを整理します。最後に、フレックスタイルのミニチュア版がコマンドとコントロールの機能を果たします。チップの周辺部には、Anton2のASICを相互に接続するSerDesチャンネル、ホストマシンに接続するためのIOインターフェイス、そしてオンチップのロジックアナライザが搭載されています。」とButtsは語る。

 

Buttsは、Anton 3の主な目標を挙げた。「力の計算を高速化することで明らかになったパフォーマンスのボトルネックにも対処しなければなりません。ボンドやGCコードの計算がボトルネックの一つであり、オフチップの通信帯域の拡張性が限られていることもボトルネックの一つです。また、高速化だけでなく、大規模なシミュレーションへの対応、プログラミングの容易化、演算パイプラインでの新しい力場機能のサポートなど、マシンの機能を向上させたいと考えました。最後に、我々の設計チーム(40人程度)はムーアの法則に従って拡張しているわけではないので、設計と実装の複雑さをコントロールする必要があります。」と述べている。

新しいAnton 3コアタイルは、MDの計算に必要なすべての主要コンポーネントを、いくつかのユニークなブロックに集約している、とButtsは言う。中央のルーターは、このようなタイルの大規模な配列の間でチップ上のネットワーク接続を提供する。GCとPPIMは「Anton 2」でおなじみだが、重要な進化を遂げている。例えばPPIMSでは、より広範な力場をサポートするために新しい機能形態が実装されている。GCのメモリ容量は2倍になり、「より大規模なシミュレーションとより柔軟なソフトウェアを可能にした。」また、GCの命令セットも最適化され、新しい命令と高密度のエンコーディングにより、GCの命令キャッシュの有効容量が増加した。

 

「これらの進化に加えて、より重要な変化がここにもありました。それは、柔軟性のある特化したコンピュートリソースを同じタイルにコロケーションしたことです。Anton 3は、GCとPPIMの間の双方向通信をサポートしており、PPIMの高スループットなペア選択とGCのプログラマビリティを細かく連携させた新しいユースモデルを可能にしています。」とButtsは説明する。「また、ボンドコンピュテーションのボトルネックは、ボンドフォースコンピュテーションに特化した新しいパイプラインを導入することで解消されました。ペアワイズインタラクションの単位面積当たりのスループットは、範囲制限のあるインタラクションを新たに分解することで2倍になりました。さらに、同期機能がメモリーとネットワークに分散されたことで、別のディスパッチユニットが不要になりました。」

Anton 3 ASICのクロックは2.8ギガヘルツで、スループットとレイテンシーの両方に貢献し、rawチャネル帯域幅は2倍以上になった(Anton 2は1.5ギガヘルツ)。これらの重要なパラメーターはいずれも、さらなる改善の余地があるとButtsは述べている。トランジスタ数は、新しい7nm FinFETプロセスの採用により16倍に増加した。Anton 2と比較した場合、ダイサイズの増加はわずかだ。最後に、「ASICの消費電力は360WとAnton 2の2倍弱になっていますが、ダイ面積と周波数で正規化するとほぼ同じです。」とButtsは主張している。

 

Hot Chipsの発表会である以上、チップの詳細に注目するのは当然のことだ。Buttsは、配電技術の革新について説明し、ボンドの改良などいくつかの分野についてはより深く掘り下げた。Anton 3の初期の研究課題についてはほとんど語られなかったが、初期のシリコンは問題なく稼働し、512ノードへの拡張も比較的簡単にできたと述べた。

今回のアップグレードの大きな成果は、より大きなシステムをより長い時間シミュレーションできるようになったことだ。これは、このような粒度の高いMDシミュレーションを実用化するための鍵となりる。Anton 3では、それが実現できるはずだ。これまでのところ、AntonではML/DLの使用は限られているようだが、今後は変わる可能性がある。

 

講演後の質疑応答を担当したDESRES社のエンジニアであるBrannon Batsonは、データの再利用や、Antonへの投機的なAIを組み込むことについて質問され、「分子動力学では、古いデータを使用することは非常に困難です。分子動力学では、古いデータを使用することは非常に困難です。時間的に積分可能なシステムの重要な特性のために、シミュレーションでは最新のデータのみを使用する必要があり、過去のタイムステップのデータを使用することはできません。ロールバックできる能力があれば、投機的に使うことができるかもしれませんが、少なくともこれまでのところ、私たちの答えは「いいえ」、あるいは「できるとしても、まだ方法がわかっていません」です。」

「私たちの計算機は大きく2つのクラスに分けられます。一つは、より古典的な汎用フレキシブルサブシステムです。そこでは、32ビットの固定小数点演算とそのベクトル形式に大きく依存しています。しかし、特化した演算パイプラインは、あらゆる場所で使用されています。5ビットの指数で14ビットの算術演算を行うところもあれば、対数領域の演算を行うところもあります。私たちはアンロール・パイプラインを採用しているので、非常に慎重な数値解析に基づいて、ステージごとに慎重に精度を調整していますが、これが私たちの価値提案の多くを占めています。」

また、オークリッジ国立研究所のJeffrey Vetterからは、「512ノード以上の大規模なMDシミュレーションにも対応できるのか」「材料科学を中心とした他のタイプの分子動力学アプリケーションも検討しているのか」という興味深い質問があった。

Batsonは、「リンク層のネットワークに関しては、物理的に512ノード以上の拡張が可能なハードウェアであることは言えます。しかし、このマシンは、最大でも512ノードで1つの分子動力学シミュレーションを実行するように、正常に動作するように設計されています。それ以上の規模になるのは、複数のシミュレーションを実行し、それらの間でデータを交換したい場合です。そのための機能もあります。」と回答した。

「他のアプリケーション、特に材料科学については、Antonの恩恵を受けられるアプリケーションがあるかもしれません。DESRESでも多少検討しましたが、あまり多くは検討していません。前にも言いましたが、私たちはコンピュータの会社ではありません。私たちの焦点は、病気を治し、人間の不幸を和らげることであり、その範囲外のことに多くの時間を費やすことはありません。」

DESRES社のAtonシステムについては、ライフサイエンスの研究者の間では、常に賞賛と少しの羨望が入り混じっていた。このようなシステムは強力であるが、希少価値がある。

Anton 1とAnton 2

 

生命科学分野の計算機研究のコンサルティングを行うBioTeam社のCEOであるAri Bermanは、「Antonシステムは、この分野の競争力のバランスを常に崩しており、複雑な生体分子システムを解く研究者にとっては常に聖杯のような存在です。Anton 3は、その優位性をまったく新しいレベルに引き上げ、研究者が利用できる最高の汎用スーパーコンピュータと比較して、システムへのアクセスを100倍有利にするものです。この種のシミュレーションは、医薬品の開発期間を短縮したり、病気の原因となる突然変異などの生物学的システムを理解するための変数の絞り込みに役立つなど、非常に重要な役割を果たします。」と述べている。

「最も最近の強力な例は、構造生物学とMDシミュレーションが、SARS-CoV-2ウイルスの表面にあるスパイクタンパク質のメカニズムを素早く解き明かし、解決することで、COVID-19のワクチンを前代未聞の速さで作ることができるようになったことです。ただし、Anton 3は1つしかなく、それを使いたい場合は、(2017年にPSCでAnton 2を使ったように)システムをより広く利用できるようにするまで、D.E.Shawにアピールする必要があります。おそらく今回は、パブリッククラウドがこの技術に関心を持ち、D.E.Shawと協力して、少なくとも適度にこの技術を利用できるようにしてくれるでしょう。」

おそらく、SC21の論文(と、あればプレゼン)では、最新のAntonシステムの事例について詳しく語られるだろう。

[i] D. E. Shaw et al., “Anton 3: Twenty Microseconds of Molecular Dynamics Simulation Before Lunch,” to appear in SC’21: Proceedings of the International Conference for High Performance Computing, Networking, Storage, and Analysis, 2021.